scripts/test_public_tunnel_fixture.py
Rama main · réplica de lectura del repositorio original.
#!/usr/bin/env python3
"""Exercise the disposable Tunnel app and its OCI image without a public service."""
import base64
import contextlib
import gzip
import hashlib
import http.client
import io
import json
from pathlib import Path
import socket
import struct
import subprocess
import tarfile
import tempfile
import time
import unittest
import public_tunnel_fixture as fixture
def fake_elf(machine=62, dynamic=False):
data = bytearray(120)
data[:6] = b"\x7fELF\x02\x01"
struct.pack_into("<HH", data, 16, 2, machine)
struct.pack_into("<Q", data, 32, 64)
struct.pack_into("<HH", data, 54, 56, 1)
struct.pack_into("<I", data, 64, 3 if dynamic else 1)
return bytes(data)
def read_exact(sock, size):
data = b""
while len(data) < size:
received = sock.recv(size - len(data))
if not received:
raise EOFError("fixture closed the socket")
data += received
return data
class FixtureTests(unittest.TestCase):
@classmethod
def setUpClass(cls):
cls.temp = tempfile.TemporaryDirectory()
cls.binary = Path(cls.temp.name) / "fixture"
subprocess.run(["rustc", "--edition=2021", "-O", str(Path(__file__).with_name("public_tunnel_fixture.rs")),
"-o", str(cls.binary)], check=True, capture_output=True)
@classmethod
def tearDownClass(cls):
cls.temp.cleanup()
def launch(self, label):
with socket.socket() as reserve:
reserve.bind(("127.0.0.1", 0))
port = reserve.getsockname()[1]
process = subprocess.Popen([str(self.binary), label, f"127.0.0.1:{port}"],
stdout=subprocess.DEVNULL, stderr=subprocess.PIPE)
self.addCleanup(self.stop, process)
deadline = time.monotonic() + 5
while True:
try:
with socket.create_connection(("127.0.0.1", port), timeout=.2):
return port
except OSError:
if process.poll() is not None or time.monotonic() >= deadline:
self.fail("fixture did not start")
time.sleep(.02)
@staticmethod
def stop(process):
process.terminate()
try:
process.communicate(timeout=5)
except subprocess.TimeoutExpired:
process.kill()
process.communicate(timeout=5)
def test_distinct_http_targets_and_websocket_echo(self):
for label in ("A", "B"):
with self.subTest(label=label):
port = self.launch(label)
conn = http.client.HTTPConnection("127.0.0.1", port, timeout=3)
self.addCleanup(conn.close)
for path, expected in (("/app", f"fixture-{label}".encode()), ("/health", b"healthy")):
conn.request("GET", path)
response = conn.getresponse()
self.assertEqual((response.status, response.read()), (200, expected))
key = base64.b64encode(b"0123456789abcdef").decode()
with socket.create_connection(("127.0.0.1", port), timeout=3) as stream:
stream.sendall(("GET /ws HTTP/1.1\r\nHost: test\r\nConnection: Upgrade\r\nUpgrade: websocket\r\n"
f"Sec-WebSocket-Key: {key}\r\nSec-WebSocket-Version: 13\r\n\r\n").encode())
headers = b""
while not headers.endswith(b"\r\n\r\n"):
headers += read_exact(stream, 1)
accept = base64.b64encode(hashlib.sha1((key + "258EAFA5-E914-47DA-95CA-C5AB0DC85B11").encode()).digest())
self.assertTrue(headers.startswith(b"HTTP/1.1 101 "))
self.assertIn(b"Sec-WebSocket-Accept: " + accept, headers)
message, mask = b"qa-echo-123", b"\0\xff\1\x80"
masked = bytes(byte ^ mask[i % 4] for i, byte in enumerate(message))
stream.sendall(bytes((0x81, 0x80 | len(message))) + mask + masked)
self.assertEqual(read_exact(stream, 2), bytes((0x81, len(message))))
self.assertEqual(read_exact(stream, len(message)), message)
def test_image_contains_only_the_static_fixture_and_known_label(self):
image = fake_elf()
first = fixture.image_fixture(image, "A")
self.assertEqual(first, fixture.image_fixture(image, "A"))
digest, raw_manifest, blobs = first
self.assertEqual(digest, fixture.digest(raw_manifest))
manifest = json.loads(raw_manifest)
config = json.loads(blobs[manifest["config"]["digest"]])
self.assertEqual(config["architecture"], "amd64")
self.assertEqual(config["config"]["Entrypoint"], ["/bin/public-tunnel-fixture", "A"])
raw_layer = gzip.decompress(blobs[manifest["layers"][0]["digest"]])
self.assertEqual(config["rootfs"]["diff_ids"], [fixture.digest(raw_layer)])
with tarfile.open(fileobj=io.BytesIO(raw_layer)) as archive:
self.assertEqual(archive.getnames(), ["bin", "bin/public-tunnel-fixture"])
self.assertEqual(archive.extractfile("bin/public-tunnel-fixture").read(), image)
self.assertNotEqual(digest, fixture.image_fixture(image, "B")[0])
def test_connection_limit_preserves_service_after_slow_clients_disconnect(self):
port = self.launch("A")
clients = []
self.addCleanup(lambda: [client.close() for client in clients])
for _ in range(16):
client = socket.create_connection(("127.0.0.1", port), timeout=3)
client.sendall(b"G")
clients.append(client)
deadline = time.monotonic() + 3
while True:
with socket.create_connection(("127.0.0.1", port), timeout=3) as extra:
extra.sendall(b"GET /health HTTP/1.1\r\nHost: test\r\n\r\n")
try:
refused = extra.recv(1) == b""
except ConnectionResetError:
refused = True
if refused:
break
if time.monotonic() >= deadline:
self.fail("fixture did not enforce its connection limit")
time.sleep(.02)
for client in clients:
client.close()
clients.clear()
deadline = time.monotonic() + 5
while True:
try:
with contextlib.closing(http.client.HTTPConnection("127.0.0.1", port, timeout=1)) as conn:
conn.request("GET", "/health")
response = conn.getresponse()
self.assertEqual((response.status, response.read()), (200, b"healthy"))
break
except (OSError, http.client.HTTPException):
if time.monotonic() >= deadline:
raise
time.sleep(.02)
def test_rejects_unsafe_or_unusable_images(self):
for binary in [b"not elf", fake_elf(dynamic=True), fake_elf(machine=40), fake_elf()[:80]]:
with self.subTest(binary=binary[:20]), self.assertRaises(ValueError):
fixture.image_fixture(binary, "A")
with self.assertRaises(ValueError):
fixture.image_fixture(fake_elf(), "other")
if __name__ == "__main__":
unittest.main()