diff --git a/kernelguard.py b/kernelguard.py index d59e926..1fd8f09 100644 --- a/kernelguard.py +++ b/kernelguard.py @@ -111,13 +111,42 @@ # "return cache[...]" but exclude compiled-kernel / module caches (legitimate) RE_RETURN_CACHE_INDEX = re.compile( r'return\s+(?!_?(?:compiled|kernel|module|func|op)_?\w*cache)' - r'\w*(?:cache|reuse)\w*\s*\[', + r'(?P\w*(?:cache|reuse)\w*)\s*\[', re.IGNORECASE, ) RE_CACHE_STORE_OUTPUT = re.compile( r'\w*(?:cache|reuse)\w*\s*\[[^\]]+\]\s*=\s*(?:result|output|out|c_ref)\b', re.IGNORECASE, ) +# Cache populated by a `.compile(...)` method call — stores a compiled +# callable (not an output tensor), so a subsequent `return cache[key]` +# returns the callable, not a result. Recognizes the standard tile-DSL +# compile-cache idiom used by TileLang / Triton / CuTeDSL / Pallas: +# +# _KERNEL_CACHE[(M, N, K)] = my_kernel.compile(M=M, N=N, K=K, ...) +# _KCACHE[key] = _gemm.compile(...) +# +# Variable-name allowlists (e.g. RE_RETURN_CACHE_INDEX's +# `compiled|kernel|module|func|op` prefix list) miss abbreviated names +# (e.g. `_KCACHE`, `_CCACHE`, `_TILECACHE`); checking the RHS of the +# store is more robust because the *value* in the cache is what +# distinguishes legitimate compile-cache from output replay. +# +# Deliberately matches only the `.compile(` method-call form, not a +# function-call form like `_compile_kernel(out)` — the latter is too +# easy to forge with an identity passthrough (e.g. attacker defines +# `_compile_passthrough(x): return x` and writes +# `_CACHE[key] = _compile_passthrough(output)`). The method-call form +# requires the cached value to actually have a `.compile()` attribute, +# which PyTorch tensors do not. +# Whitespace tolerance: callers apply this against strip_python_strings_and_comments(...) +# which emits a space between tokens, so source `_gemm.compile(` becomes +# `_gemm . compile (`. The `\s*` around the `.` covers both forms. +RE_CACHE_STORES_COMPILED_CALLABLE = re.compile( + r'(?P\w*(?:cache|reuse)\w*)\s*\[[^\]]+\]\s*=\s*' + r'[^=\n]*?\s*\.\s*compile\s*\(', + re.IGNORECASE, +) RE_CPP_CACHE_FAST_RETURN = re.compile( r'if\s*\([^)]*(?:cache\.last|cache\.prev|lastA|lastB|prevA|prevB)[^)]*\)\s*(?:\{[^{}]{0,200}?\breturn\b|return\b)', re.DOTALL, @@ -219,6 +248,119 @@ def strip_python_comments(code: str) -> str: return tokenize.untokenize(tokens) +def strip_python_strings_and_comments(code: str) -> str: + """Replace Python COMMENT and STRING tokens with same-length whitespace. + + Used to defeat string-literal / comment decoys: an attacker can plant + text matching a legitimate-looking pattern inside a docstring or a + `NOTE = "..."` constant to flip a detector signal. Most detectors + operate on `python_active` (which strips comments only) — when a + detector needs to be robust against in-string decoys, run it against + this stripped form instead. + + Preserves line/column offsets so error messages stay roughly intact. + Falls back to the original code on tokenize / syntax errors.""" + try: + out = [] + for tok in tokenize.generate_tokens(io.StringIO(code).readline): + if tok.type in (tokenize.COMMENT, tokenize.STRING): + out.append(" " * len(tok.string)) + elif tok.type in ( + tokenize.NL, tokenize.NEWLINE, + tokenize.INDENT, tokenize.DEDENT, + tokenize.ENCODING, tokenize.ENDMARKER, + ): + if tok.type in (tokenize.NL, tokenize.NEWLINE): + out.append("\n") + else: + out.append(tok.string) + out.append(" ") + return "".join(out) + except (tokenize.TokenError, IndentationError, SyntaxError, ValueError): + return code + + +def _compiled_callable_cache_names(scope: str) -> set[str]: + outputish_names = {"result", "output", "out", "c_ref"} + + try: + tree = ast.parse(scope) + except SyntaxError: + stripped = strip_python_strings_and_comments(scope) + names: set[str] = set() + for match in RE_CACHE_STORES_COMPILED_CALLABLE.finditer(stripped): + rhs = match.group(0).split("=", 1)[-1] + if not re.search(r'\b(?:result|output|out|c_ref)\b', rhs, re.IGNORECASE): + names.add(match.group("cache").lower()) + return names + + local_compile_classes = { + node.name + for node in ast.walk(tree) + if isinstance(node, ast.ClassDef) + and any( + isinstance(item, (ast.FunctionDef, ast.AsyncFunctionDef)) + and item.name == "compile" + for item in node.body + ) + } + local_compile_instances: set[str] = set() + outputish_instances: set[str] = set() + compiled_caches: set[str] = set() + + for node in ast.walk(tree): + if isinstance(node, ast.Assign): + for target in node.targets: + if ( + isinstance(target, ast.Name) + and isinstance(node.value, ast.Call) + and isinstance(node.value.func, ast.Name) + and node.value.func.id in local_compile_classes + ): + local_compile_instances.add(target.id) + if ( + isinstance(target, ast.Name) + and isinstance(node.value, ast.Call) + and (_expr_names(node.value) & outputish_names) + ): + outputish_instances.add(target.id) + + def is_safe_compile_call(expr: ast.AST) -> bool: + if not ( + isinstance(expr, ast.Call) + and isinstance(expr.func, ast.Attribute) + and expr.func.attr == "compile" + ): + return False + if _expr_names(expr) & outputish_names: + return False + receiver = expr.func.value + if ( + isinstance(receiver, ast.Call) + and isinstance(receiver.func, ast.Name) + and receiver.func.id in local_compile_classes + ): + return False + if isinstance(receiver, ast.Name) and receiver.id in local_compile_instances: + return False + if isinstance(receiver, ast.Name) and receiver.id in outputish_instances: + return False + return True + + for node in ast.walk(tree): + if not isinstance(node, ast.Assign): + continue + if not is_safe_compile_call(node.value): + continue + for target in node.targets: + if isinstance(target, ast.Subscript): + root = _ast_root_name(target.value) + if root and re.search(r'(?:cache|reuse)', root, re.IGNORECASE): + compiled_caches.add(root.lower()) + + return compiled_caches + + def extract_function_block(code: str, func_name: str) -> str: """Best-effort extraction of a Python function block from source text.""" lines = code.splitlines() @@ -1722,8 +1864,28 @@ def detect_result_caching(code: str | SubmissionFacts) -> list[dict]: workspace_signals.append("id(data) cache key") if RE_DATA_PTR_CACHE_KEY.search(cache_scope): workspace_signals.append("data_ptr() cache key") - if RE_RETURN_CACHE_INDEX.search(cache_scope): - output_replay_signals.append("direct return from cache[...]") + returned_caches = [ + match.group("cache").lower() + for match in RE_RETURN_CACHE_INDEX.finditer(cache_scope) + ] + if returned_caches: + # Suppress when the cache provably stores a compiled callable rather + # than an output tensor — `cache[key] = factory.compile(...)` is the + # standard tile-DSL compile-cache idiom (TileLang, Triton, CuTeDSL, + # Pallas). The subsequent `return cache[key]` returns the callable, + # not a cached result, so it isn't output replay. See + # RE_CACHE_STORES_COMPILED_CALLABLE for the rationale on why the + # method-call form is the only one allowlisted (forgery resistance). + # + # Strip strings AND comments from the scope before matching: an + # attacker can plant a decoy `_CACHE[key] = foo.compile(...)` + # inside a module-level docstring / constant string to enable + # suppression while keeping a real output-replay cache in the + # active code. cache_scope is already comment-stripped via + # strip_python_comments; this adds string-literal stripping. + compiled_caches = _compiled_callable_cache_names(cache_scope) + if any(cache_name not in compiled_caches for cache_name in returned_caches): + output_replay_signals.append("direct return from cache[...]") if stores_output and output_replay_signals: output_replay_signals.append("cache[...] stores output/result tensor") elif stores_output: