test: guard the legacy pickle scan of spans above INT_MAX

tests/native/pickle_span_test.cpp

@@ -0,0 +1,161 @@
+// Runs upstream's pickle scanner, compiled from the pinned submodule, on the
+// spans read_torch_legacy_file gives it. That reader passes the rest of the
+// whole checkpoint, tensor storage included, so a small pickle can sit in a
+// span longer than INT_MAX.
+//
+// Modes:
+//   bounded    each pickle in a buffer of exactly its size
+//   large      malformed pickles at the start of spans above INT_MAX
+//   large-ok   valid pickles at the start of spans above INT_MAX
+//
+// Large spans are mapped read-only, so they cost no memory; only the first
+// page is writable and holds the pickle.
+
+#include "model_io/pickle_io.h"
+
+#include <sys/mman.h>
+#include <unistd.h>
+
+#include <cstdarg>
+#include <cstdint>
+#include <cstdio>
+#include <cstdlib>
+#include <cstring>
+#include <string>
+#include <vector>
+
+#define CHECK(condition)                                                                       \
+    do {                                                                                       \
+        if (!(condition)) {                                                                    \
+            std::fprintf(stderr, "%s:%d: CHECK(%s) failed\n", __FILE__, __LINE__, #condition); \
+            std::_Exit(1);                                                                     \
+        }                                                                                      \
+    } while (0)
+
+// Stands in for upstream's util.cpp; pickle_io.cpp calls it only to format errors.
+std::string sd_format(const char* format, ...) {
+    char buffer[1024];
+    va_list args;
+    va_start(args, format);
+    std::vsnprintf(buffer, sizeof(buffer), format, args);
+    va_end(args);
+    return buffer;
+}
+
+namespace {
+
+struct Pickle {
+    const char* name;
+    std::string bytes;
+};
+
+// Keeps embedded NULs; drops the literal's terminator.
+template <size_t N>
+std::string bytes(const char (&data)[N]) {
+    return std::string(data, N - 1);
+}
+
+// Every opcode whose operand skip_pickle_object scans to a newline, plus a
+// length-prefixed string, which does not depend on the span.
+std::vector<Pickle> valid_pickles() {
+    return {
+        {"GLOBAL", bytes("\x80\x02" "ccollections\nOrderedDict\n)R.")},
+        {"INST", bytes("\x80\x02(icollections\nOrderedDict\n.")},
+        {"PERSID", bytes("\x80\x02P0\n.")},
+        {"UNICODE", bytes("\x80\x02Vweight\n.")},
+        {"BINUNICODE", bytes("\x80\x02X\x06\x00\x00\x00weight.")},
+    };
+}
+
+std::vector<Pickle> malformed_pickles() {
+    return {
+        {"GLOBAL without a name line", bytes("\x80\x02" "ccollections\nOrderedDict")},
+        {"GLOBAL without a module line", bytes("\x80\x02" "ccollections")},
+        {"PERSID without a newline", bytes("\x80\x02P0")},
+        {"BINUNICODE past the end", bytes("\x80\x02X\xff\x00\x00\x00weight.")},
+        {"unknown opcode", bytes("\x80\x02\x00.")},
+        {"no STOP", bytes("\x80\x02)")},
+    };
+}
+
+bool skip(const uint8_t* data, size_t size, size_t* consumed) {
+    *consumed = 0;
+    return skip_pickle_object(data, size, consumed);
+}
+
+void check_bounded() {
+    for (const Pickle& pickle : valid_pickles()) {
+        // An exact-size heap copy lets a sanitizer catch reads past the end.
+        std::vector<uint8_t> data(pickle.bytes.begin(), pickle.bytes.end());
+        size_t consumed = 0;
+        std::printf("bounded %s\n", pickle.name);
+        CHECK(skip(data.data(), data.size(), &consumed));
+        CHECK(consumed == data.size());
+    }
+    for (const Pickle& pickle : malformed_pickles()) {
+        std::vector<uint8_t> data(pickle.bytes.begin(), pickle.bytes.end());
+        size_t consumed = 0;
+        std::printf("bounded %s\n", pickle.name);
+        CHECK(!skip(data.data(), data.size(), &consumed));
+    }
+}
+
+// Above INT_MAX, and above UINT32_MAX, where narrowing to int keeps only the
+// low 32 bits.
+const size_t kLargeSpans[] = {(size_t(1) << 31) + 4096, (size_t(1) << 32) + 16};
+
+uint8_t* map_span(size_t span, const std::string& pickle) {
+    const size_t page = (size_t)sysconf(_SC_PAGESIZE);
+    void* mapped      = mmap(nullptr, span, PROT_READ, MAP_PRIVATE | MAP_ANONYMOUS, -1, 0);
+    CHECK(mapped != MAP_FAILED);
+    CHECK(mprotect(mapped, page, PROT_READ | PROT_WRITE) == 0);
+    CHECK(pickle.size() <= page);
+    std::memcpy(mapped, pickle.data(), pickle.size());
+    return static_cast<uint8_t*>(mapped);
+}
+
+void check_large(const std::vector<Pickle>& pickles, bool valid) {
+    for (size_t span : kLargeSpans) {
+        for (const Pickle& pickle : pickles) {
+            uint8_t* data   = map_span(span, pickle.bytes);
+            size_t consumed = 0;
+            std::printf("span %zu %s\n", span, pickle.name);
+            std::fflush(stdout);
+            const bool skipped = skip(data, span, &consumed);
+            CHECK(munmap(data, span) == 0);
+            if (valid) {
+                CHECK(skipped);
+                CHECK(consumed == pickle.bytes.size());
+            } else {
+                CHECK(!skipped);
+            }
+        }
+    }
+}
+
+// The rest of a malformed pickle in a large span is zeros, which the scanner
+// rejects as an unknown opcode; a scan for a missing newline would read up to
+// INT_MAX bytes of them, so those cases stay bounded only.
+std::vector<Pickle> malformed_large_pickles() {
+    return {
+        {"no STOP", bytes("\x80\x02)")},
+        {"BINUNICODE past the pickle", bytes("\x80\x02X\x00\x00\x00\x80")},
+    };
+}
+
+}  // namespace
+
+int main(int argc, char** argv) {
+    const std::string mode = argc > 1 ? argv[1] : "";
+    if (mode == "bounded") {
+        check_bounded();
+    } else if (mode == "large") {
+        check_large(malformed_large_pickles(), false);
+    } else if (mode == "large-ok") {
+        check_large(valid_pickles(), true);
+    } else {
+        std::fprintf(stderr, "usage: %s bounded|large|large-ok\n", argv[0]);
+        return 2;
+    }
+    return 0;
+}

tests/test_pickle_span.py

@@ -0,0 +1,86 @@
+"""Builds and runs `tests/native/pickle_span_test.cpp`.
+
+The test compiles upstream's `pickle_io.cpp` from the pinned submodule, so it
+needs the submodules' sources but no built runtime and no model. Spans above
+INT_MAX are reserved read-only and cost no memory.
+"""
+
+import os
+import shutil
+import subprocess
+import sys
+import tempfile
+import unittest
+from pathlib import Path
+
+REPO_ROOT = Path(__file__).resolve().parent.parent
+UPSTREAM = REPO_ROOT / "third_party" / "stable-diffusion.cpp"
+SOURCES = [
+    UPSTREAM / "src" / "model_io" / "pickle_io.cpp",
+    REPO_ROOT / "tests" / "native" / "pickle_span_test.cpp",
+]
+INCLUDE_DIRS = [UPSTREAM / "src", UPSTREAM / "include", UPSTREAM / "ggml" / "include"]
+COMPILER = os.environ.get("CXX") or shutil.which("c++")
+ADDRESS_SANITIZER = "-fsanitize=address"
+
+
+@unittest.skipIf(sys.platform == "win32", "uses a GCC or Clang command line and mmap")
+class PickleSpanTest(unittest.TestCase):
+    def require(self, condition: bool, reason: str) -> None:
+        """Skips locally, but fails in CI, where a skip would hide the test."""
+        if condition:
+            return
+        if os.environ.get("CI"):
+            self.fail(reason)
+        self.skipTest(reason)
+
+    def setUp(self) -> None:
+        self.require(COMPILER is not None, "no C++ compiler; set CXX")
+        self.require((INCLUDE_DIRS[2] / "ggml.h").is_file() and SOURCES[0].is_file(),
+                     "third_party/stable-diffusion.cpp is not checked out; run "
+                     "`git submodule update --init --recursive`")
+
+    def run_mode(self, mode: str, *flags: str) -> subprocess.CompletedProcess:
+        includes = [f"-I{directory}" for directory in INCLUDE_DIRS]
+        with tempfile.TemporaryDirectory() as directory:
+            binary = Path(directory) / "pickle_span_test"
+            subprocess.run([COMPILER, "-std=c++17", "-O1", "-g", *flags, *includes,
+                            *map(str, SOURCES), "-o", str(binary)], check=True)
+            return subprocess.run([str(binary), mode], capture_output=True, text=True,
+                                  timeout=600)
+
+    def assert_passes(self, mode: str, *flags: str) -> None:
+        result = self.run_mode(mode, *flags)
+        self.assertEqual(0, result.returncode, result.stdout + result.stderr)
+
+    def supports(self, flag: str) -> bool:
+        with tempfile.TemporaryDirectory() as directory:
+            source = Path(directory) / "empty.cpp"
+            source.write_text("int main() { return 0; }\n")
+            binary = Path(directory) / "empty"
+            built = subprocess.run([COMPILER, flag, str(source), "-o", str(binary)],
+                                   capture_output=True)
+            return built.returncode == 0 and subprocess.run(
+                [str(binary)], capture_output=True).returncode == 0
+
+    def test_skips_valid_and_rejects_malformed_bounded_pickles(self) -> None:
+        self.assert_passes("bounded")
+
+    def test_bounded_pickles_stay_in_bounds_under_address_sanitizer(self) -> None:
+        self.require(self.supports(ADDRESS_SANITIZER),
+                     f"{COMPILER} cannot build or run {ADDRESS_SANITIZER} binaries")
+        self.assert_passes("bounded", ADDRESS_SANITIZER)
+
+    def test_rejects_malformed_pickles_in_spans_above_int_max(self) -> None:
+        self.assert_passes("large")
+
+    # Upstream narrows the remaining span to int before scanning for a newline
+    # (#23), so the pinned scanner rejects these. A submodule bump that fixes it
+    # makes this an unexpected success, which fails the run: remove the marker.
+    @unittest.expectedFailure
+    def test_skips_valid_pickles_in_spans_above_int_max(self) -> None:
+        self.assert_passes("large-ok")
+
+
+if __name__ == "__main__":
+    unittest.main()