Bound the cost of identify() on very large files

sflock/abstracts.py

@@ -14,7 +14,7 @@
 import subprocess
 import tempfile
 
-from sflock.config import iter_passwords
+from sflock.config import MAX_IDENT_SCAN_SIZE, iter_passwords
 from sflock.exception import UnpackException
 from sflock.misc import data_file, make_list
 from sflock.pick import package, platform
@@ -264,6 +264,7 @@ def __init__(
         self._ole = None
         self._ole_tried = False
         self._header = None
+        self._scan_buffer = None
 
         # Filepaths of all child entries if this is an archive.
         self.filepaths = []
@@ -314,6 +315,21 @@ def header(self):
             self._header = self.stream.read(1024 * 1024)
         return self._header or b""
 
+    @property
+    def scan_buffer(self):
+        """Head of the file, for content-sniffing identifiers. Bounded so
+        that identification stays cheap on very large files."""
+        if self._scan_buffer is None:
+            if self._contents is not None:
+                self._scan_buffer = self._contents[:MAX_IDENT_SCAN_SIZE]
+            elif self._stream is not None:
+                self._stream.seek(0)
+                self._scan_buffer = self._stream.read(MAX_IDENT_SCAN_SIZE)
+            elif self.filepath:
+                with open(self.filepath, "rb") as fh:
+                    self._scan_buffer = fh.read(MAX_IDENT_SCAN_SIZE)
+        return self._scan_buffer or b""
+
     @property
     def magic(self):
         if not self._magic and self.filesize:

sflock/config.py

@@ -7,6 +7,11 @@
 # May be tweaked in the future including modifying this at runtime.
 MAX_TOTAL_SIZE = 1024 * 1024 * 1024
 
+# Content-scanning identifiers only inspect the head of a file. File-type
+# identification does not need the whole buffer, and unbounded scans over very
+# large files took minutes.
+MAX_IDENT_SCAN_SIZE = 16 * 1024 * 1024
+
 
 def iter_passwords():
     from importlib.resources import as_file, files

sflock/ident.py

@@ -209,7 +209,7 @@
 
 
 def is_executable(f):
-    return f.contents.startswith((b"MZ", b"\x7fELF"))
+    return f.header.startswith((b"MZ", b"\x7fELF"))
 
 
 def have_enough_memory_for_unicorn():
@@ -455,73 +455,84 @@ def powershell(f):
     if found > 1:
         return "ps1"
 
+# Each pattern is paired with the literal substring(s) that must be present
+# for it to have any chance of matching. Most of these patterns are \b-anchored,
+# which defeats re's literal-prefix optimization and makes them scan at ~60
+# MB/s instead of the ~3500 MB/s of a plain substring search; gating on the
+# (much cheaper) literal first avoids running the regex at all on the
+# overwhelming majority of non-matching content.
 nodejs_patterns = {
     "Explicit Directives (Highest Confidence)": [
         # Catches #!/usr/bin/env node
-        rb"^#!.*\bnode\b",
+        ((b"node",), rb"^#!.*\bnode\b"),
         # Catches import ... from 'node:fs'
-        rb"['\"]node:[a-zA-Z\/]+['\"]"
+        ((b"node:",), rb"['\"]node:[a-zA-Z\/]+['\"]"),
     ],
 
     "Core Globals": [
         # Robust process detection
-        rb"\bprocess\.(env|argv|cwd|exit|platform|versions|nextTick)\b",
-        rb"\bglobal\.(?!\.)",
+        ((b"process.",), rb"\bprocess\.(env|argv|cwd|exit|platform|versions|nextTick)\b"),
+        ((b"global.",), rb"\bglobal\.(?!\.)"),
         # Legacy Buffer usage
-        rb"\bBuffer\.(from|alloc|allocUnsafe|concat)\b",
-        rb"\b__dirname\b",
-        rb"\b__filename\b"
+        ((b"Buffer.",), rb"\bBuffer\.(from|alloc|allocUnsafe|concat)\b"),
+        ((b"__dirname",), rb"\b__dirname\b"),
+        ((b"__filename",), rb"\b__filename\b"),
     ],
 
     "System Execution (Critical)": [
         # Catches require('child_process') OR from 'child_process'
-        rb"(?:require\s*\(|from\s+)['\"]child_process['\"]",
-        rb"\bspawn\(",
-        rb"\bexec\(",
-        rb"\bexecSync\(",
-        rb"\bfork\("
+        ((b"child_process",), rb"(?:require\s*\(|from\s+)['\"]child_process['\"]"),
+        ((b"spawn(",), rb"\bspawn\("),
+        ((b"exec(",), rb"\bexec\("),
+        ((b"execSync(",), rb"\bexecSync\("),
+        ((b"fork(",), rb"\bfork\("),
     ],
 
     "File System Access": [
         # Catches require('fs'), require('fs/promises'), from 'fs', etc.
-        rb"(?:require\s*\(|from\s+)['\"](fs|fs\/promises|path)['\"]",
-        rb"\bfs\.readFile",
-        rb"\bfs\.writeFile",
-        rb"\bfs\.promises\."
+        ((b"'fs", b'"fs', b"'path", b'"path'), rb"(?:require\s*\(|from\s+)['\"](fs|fs\/promises|path)['\"]"),
+        ((b"fs.readFile",), rb"\bfs\.readFile"),
+        ((b"fs.writeFile",), rb"\bfs\.writeFile"),
+        ((b"fs.promises.",), rb"\bfs\.promises\."),
     ],
 
     "Networking & OS": [
         # Catches require('net'), require('os'), require('dgram'), etc.
-        rb"(?:require\s*\(|from\s+)['\"](net|os|dgram|dns|tls|http|https)['\"]",
-        rb"\bnet\.createServer",
-        rb"\bnet\.connect",
-        rb"\bos\.cpus",
-        rb"\bos\.userInfo",
-        rb"\bos\.networkInterfaces"
+        (
+            (b"'net", b'"net', b"'os", b'"os', b"'dgram", b'"dgram', b"'dns", b'"dns', b"'tls", b'"tls', b"'http", b'"http'),
+            rb"(?:require\s*\(|from\s+)['\"](net|os|dgram|dns|tls|http|https)['\"]",
+        ),
+        ((b"net.createServer",), rb"\bnet\.createServer"),
+        ((b"net.connect",), rb"\bnet\.connect"),
+        ((b"os.cpus",), rb"\bos\.cpus"),
+        ((b"os.userInfo",), rb"\bos\.userInfo"),
+        ((b"os.networkInterfaces",), rb"\bos\.networkInterfaces"),
     ],
 
     "Module System": [
         # CommonJS exports (Node specific vs Browser ES modules)
-        rb"\bmodule\.exports\b",
-        rb"\bexports\.\w+\s*="
-    ]
+        ((b"module.exports",), rb"\bmodule\.exports\b"),
+        ((b"exports.",), rb"\bexports\.\w+\s*="),
+    ],
 }
 nodejs_compiled_patterns = {}
 for category, patterns in nodejs_patterns.items():
-    nodejs_compiled_patterns[category] = [re.compile(p) for p in patterns]
+    nodejs_compiled_patterns[category] = [(literals, re.compile(p)) for literals, p in patterns]
 
 def nodejs(f):
-    count = 0
-    if not f.contents:
+    buf = f.scan_buffer
+    if not buf:
         return
 
-    for category, pattern_list in nodejs_compiled_patterns.items():
-        for pattern in pattern_list:
-            if pattern.search(f.contents):
+    count = 0
+    for pattern_list in nodejs_compiled_patterns.values():
+        for literals, pattern in pattern_list:
+            if not any(literal in buf for literal in literals):
+                continue
+            if pattern.search(buf):
                 count += 1
-
-    if count >= 3:
-        return "nodejs"
+                if count >= 3:
+                    return "nodejs"
 
 def javascript(f):
     JS_STRS = [
@@ -672,7 +683,7 @@ def identify(f, check_shellcode: bool = False):
 
     if f.filename:
         for package, extensions in file_extensions.items():
-            if f.filename.endswith(extensions) and not f.contents.startswith(b"MZ"):
+            if f.filename.endswith(extensions) and not f.header.startswith(b"MZ"):
                 return package
 
     for identifier in identifiers_special:

tests/test_ident.py

@@ -3,10 +3,13 @@
 # See the file 'docs/LICENSE.txt' for copying permission.
 
 import os
+import re
 import tempfile
 
+import sflock.abstracts
+
 from sflock.abstracts import File
-from sflock.ident import identify
+from sflock.ident import identify, nodejs, nodejs_patterns
 from sflock.main import unpack
 
 
@@ -92,3 +95,82 @@ def test_iqy():
     f = unpack(b"tests/files/1.iqy")
     assert f.package == "xls"
     assert f.platform == "windows"
+
+
+def test_nodejs_literal_prefilter():
+    """Each nodejs() pattern is gated behind mandatory literal substring(s);
+    a minimal sample matching the pattern must always pass the literal gate,
+    or that pattern is silently disabled."""
+    samples = {
+        rb"^#!.*\bnode\b": b"#!/usr/bin/env node\n",
+        rb"['\"]node:[a-zA-Z\/]+['\"]": b"import fs from 'node:fs'",
+        rb"\bprocess\.(env|argv|cwd|exit|platform|versions|nextTick)\b": b"process.env",
+        rb"\bglobal\.(?!\.)": b"global.foo",
+        rb"\bBuffer\.(from|alloc|allocUnsafe|concat)\b": b"Buffer.from",
+        rb"\b__dirname\b": b"__dirname",
+        rb"\b__filename\b": b"__filename",
+        rb"(?:require\s*\(|from\s+)['\"]child_process['\"]": b"require('child_process')",
+        rb"\bspawn\(": b"spawn(",
+        rb"\bexec\(": b"exec(",
+        rb"\bexecSync\(": b"execSync(",
+        rb"\bfork\(": b"fork(",
+        rb"(?:require\s*\(|from\s+)['\"](fs|fs\/promises|path)['\"]": b"require('fs')",
+        rb"\bfs\.readFile": b"fs.readFile",
+        rb"\bfs\.writeFile": b"fs.writeFile",
+        rb"\bfs\.promises\.": b"fs.promises.",
+        rb"(?:require\s*\(|from\s+)['\"](net|os|dgram|dns|tls|http|https)['\"]": b"require('https')",
+        rb"\bnet\.createServer": b"net.createServer",
+        rb"\bnet\.connect": b"net.connect",
+        rb"\bos\.cpus": b"os.cpus",
+        rb"\bos\.userInfo": b"os.userInfo",
+        rb"\bos\.networkInterfaces": b"os.networkInterfaces",
+        rb"\bmodule\.exports\b": b"module.exports",
+        rb"\bexports\.\w+\s*=": b"exports.foo =",
+    }
+
+    seen = set()
+    for pattern_list in nodejs_patterns.values():
+        for literals, pattern in pattern_list:
+            seen.add(pattern)
+            sample = samples[pattern]
+            assert any(literal in sample for literal in literals), (literals, pattern)
+            assert re.search(pattern, sample), pattern
+
+    # Guards against a pattern being added/removed without updating this test.
+    assert seen == set(samples)
+
+
+def test_nodejs_scan_buffer_cap_contents(monkeypatch):
+    monkeypatch.setattr(sflock.abstracts, "MAX_IDENT_SCAN_SIZE", 4096)
+
+    markers = b" process.env module.exports require('child_process') "
+    filler = b"A" * 4096
+
+    within_cap = File(contents=markers + filler)
+    assert nodejs(within_cap) == "nodejs"
+
+    past_cap = File(contents=filler + markers)
+    assert nodejs(past_cap) is None
+
+
+def test_nodejs_scan_buffer_cap_stream(monkeypatch):
+    monkeypatch.setattr(sflock.abstracts, "MAX_IDENT_SCAN_SIZE", 4096)
+
+    markers = b" process.env module.exports require('child_process') "
+    filler = b"A" * 4096
+
+    fd, filepath = tempfile.mkstemp()
+    with os.fdopen(fd, "wb") as fh:
+        fh.write(filler + markers)
+    try:
+        assert nodejs(File.from_path(filepath.encode())) is None
+    finally:
+        os.unlink(filepath)
+
+    fd, filepath = tempfile.mkstemp()
+    with os.fdopen(fd, "wb") as fh:
+        fh.write(markers + filler)
+    try:
+        assert nodejs(File.from_path(filepath.encode())) == "nodejs"
+    finally:
+        os.unlink(filepath)