121 lines
3.5 KiB
Python
121 lines
3.5 KiB
Python
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)
|