feat: implement Layer 2 per-file transit, RAM session, mock buffer, and tests (v0.2.0)

tests/test_upload_refs.py

@@ -0,0 +1,120 @@
+import threading
+from http.server import HTTPServer, BaseHTTPRequestHandler
+import pytest
+from flask import Flask
+from upload.backend.upload_refs.blueprint import create_upload_refs_blueprint
+from upload.backend.session import get_files, cleanup
+
+
+class MockBufferHandler(BaseHTTPRequestHandler):
+    storage = {}
+
+    def do_PUT(self):
+        length = int(self.headers.get("Content-Length", 0))
+        data = self.rfile.read(length)
+        MockBufferHandler.storage[self.path] = data
+        self.send_response(201)
+        self.end_headers()
+
+    def do_GET(self):
+        data = MockBufferHandler.storage.get(self.path)
+        if data is None:
+            self.send_response(404)
+            self.end_headers()
+            return
+        self.send_response(200)
+        self.send_header("Content-Length", str(len(data)))
+        self.end_headers()
+        self.wfile.write(data)
+
+    def do_DELETE(self):
+        MockBufferHandler.storage.pop(self.path, None)
+        self.send_response(204)
+        self.end_headers()
+
+    def log_message(self, format, *args):
+        pass  # suppress console logs in tests
+
+
+@pytest.fixture(scope="module")
+def mock_server():
+    server = HTTPServer(("127.0.0.1", 0), MockBufferHandler)
+    port = server.server_port
+    thread = threading.Thread(target=server.serve_forever, daemon=True)
+    thread.start()
+    yield f"http://127.0.0.1:{port}/buffer/"
+    server.shutdown()
+
+
+def test_upload_refs_pull_and_delete(mock_server):
+    # Put a file into mock buffer
+    file_path = "/buffer/token123_0"
+    content = b"Content of test document for layer 2 transit"
+    MockBufferHandler.storage[file_path] = content
+
+    received_events = []
+    def on_received(sid, name, data):
+        received_events.append((sid, name, data))
+
+    app = Flask(__name__)
+    bp = create_upload_refs_blueprint({
+        "vmUploadPrefix": mock_server,
+        "pullRetries": 1,
+        "pullRetryDelay": 0.1,
+        "onFileReceived": on_received,
+    })
+    app.register_blueprint(bp)
+    client = app.test_client()
+
+    file_url = mock_server + "token123_0"
+    resp = client.post("/api/upload_refs", json={
+        "files": [
+            {"name": "test_doc.pdf", "size": len(content), "url": file_url}
+        ]
+    })
+
+    assert resp.status_code == 200
+    data = resp.get_json()
+    assert data["ok"] is True
+    sid = data["session"]
+    assert data["count"] == 1
+    assert data["added"] == 1
+
+    # Verify file is in session memory
+    files = get_files(sid)
+    assert len(files) == 1
+    assert files[0] == ("test_doc.pdf", content)
+
+    # Verify callback for Layer 3 was triggered
+    assert len(received_events) == 1
+    assert received_events[0] == (sid, "test_doc.pdf", content)
+
+    # Verify file was DELETED from mock buffer (RAM clean!)
+    assert file_path not in MockBufferHandler.storage
+
+    cleanup(sid)
+
+
+def test_upload_refs_ssrf_protection(mock_server):
+    app = Flask(__name__)
+    bp = create_upload_refs_blueprint({
+        "vmUploadPrefix": mock_server,
+    })
+    app.register_blueprint(bp)
+    client = app.test_client()
+
+    # Try to pass an evil URL outside vmUploadPrefix
+    evil_url = "http://169.254.169.254/latest/meta-data/"
+    resp = client.post("/api/upload_refs", json={
+        "files": [
+            {"name": "evil.txt", "size": 100, "url": evil_url}
+        ]
+    })
+    assert resp.status_code == 200
+    data = resp.get_json()
+    assert data["ok"] is True
+    sid = data["session"]
+    assert data["added"] == 0
+    assert data["count"] == 0
+
+    cleanup(sid)