tls test ete

Author: balogh.adam@icloud.com

src/opengradient/client/tee_connection.py

@@ -142,12 +142,10 @@ def _connect(self) -> ActiveTEE:
         """Resolve TEE from registry and create a secure HTTP client."""
         tee = self._resolve_tee()
 
-        ssl_ctx = build_ssl_context_from_der(tee.tls_cert_der) if tee.tls_cert_der else None
-        tls_verify: Union[ssl.SSLContext, bool] = ssl_ctx if ssl_ctx else True
-
+        ssl_ctx = build_ssl_context_from_der(tee.tls_cert_der) 
         return ActiveTEE(
             endpoint=tee.endpoint,
-            http_client=x402HttpxClient(self._x402_client, verify=tls_verify),
+            http_client=x402HttpxClient(self._x402_client, verify=ssl_ctx),
             tee_id=tee.tee_id,
             payment_address=tee.payment_address,
         )

tests/tee_connection_test.py

@@ -1,15 +1,25 @@
 """Tests for RegistryTEEConnection and ActiveTEE."""
 
 import asyncio
+import datetime
+import os
 import ssl
+import tempfile
 from unittest.mock import AsyncMock, MagicMock, patch
 
+import httpx
 import pytest
+from cryptography import x509
+from cryptography.hazmat.primitives import hashes, serialization
+from cryptography.hazmat.primitives.asymmetric import rsa
+from cryptography.x509.oid import NameOID
+from x402 import x402Client
 
 from src.opengradient.client.tee_connection import (
     ActiveTEE,
     RegistryTEEConnection,
 )
+from src.opengradient.client.tee_registry import build_ssl_context_from_der
 
 
 # ── Helpers ──────────────────────────────────────────────────────────
@@ -327,3 +337,128 @@ async def test_close_without_refresh_task(self):
         conn = _make_registry_connection(registry=mock_reg)
 
         await conn.close()  # should not raise
+
+
+# ── TLS certificate verification (real handshake) ────────────────────
+
+
+def _make_self_signed_cert():
+    """Generate a self-signed cert. Returns (der_bytes, pem_cert_bytes, pem_key_bytes)."""
+    key = rsa.generate_private_key(public_exponent=65537, key_size=2048)
+    subject = issuer = x509.Name([x509.NameAttribute(NameOID.COMMON_NAME, "localhost")])
+    cert = (
+        x509.CertificateBuilder()
+        .subject_name(subject)
+        .issuer_name(issuer)
+        .public_key(key.public_key())
+        .serial_number(x509.random_serial_number())
+        .not_valid_before(datetime.datetime.now(datetime.UTC))
+        .not_valid_after(datetime.datetime.now(datetime.UTC) + datetime.timedelta(days=1))
+        .sign(key, hashes.SHA256())
+    )
+    return (
+        cert.public_bytes(serialization.Encoding.DER),
+        cert.public_bytes(serialization.Encoding.PEM),
+        key.private_bytes(serialization.Encoding.PEM, serialization.PrivateFormat.TraditionalOpenSSL, serialization.NoEncryption()),
+    )
+
+
+@pytest.fixture
+async def tls_server():
+    """Spin up a local TLS server with a self-signed cert."""
+    der, pem_cert, pem_key = _make_self_signed_cert()
+
+    cert_file = tempfile.NamedTemporaryFile(suffix=".pem", delete=False)
+    key_file = tempfile.NamedTemporaryFile(suffix=".pem", delete=False)
+    try:
+        cert_file.write(pem_cert)
+        cert_file.close()
+        key_file.write(pem_key)
+        key_file.close()
+
+        server_ctx = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER)
+        server_ctx.load_cert_chain(cert_file.name, key_file.name)
+
+        async def handler(reader, writer):
+            await reader.read(4096)
+            writer.write(b"HTTP/1.1 200 OK\r\nContent-Length: 2\r\nConnection: close\r\n\r\nok")
+            await writer.drain()
+            writer.close()
+
+        server = await asyncio.start_server(handler, "127.0.0.1", 0, ssl=server_ctx)
+        port = server.sockets[0].getsockname()[1]
+
+        yield {"port": port, "der": der}
+
+        server.close()
+        await server.wait_closed()
+    finally:
+        os.unlink(cert_file.name)
+        os.unlink(key_file.name)
+
+
+def _registry_with_real_cert(tls_server):
+    """Return a mock registry that serves the local TLS server's real DER cert."""
+    return _mock_registry_with_tee(
+        endpoint=f"https://127.0.0.1:{tls_server['port']}",
+        tls_cert_der=tls_server["der"],
+        tee_id="tee-real",
+        payment_address="0xRealPay",
+    )
+
+
+@pytest.mark.asyncio
+class TestTlsCertVerification:
+    """End-to-end TLS handshake tests through RegistryTEEConnection.
+
+    A real local TLS server is started with a self-signed cert.  The registry
+    mock returns that cert's DER bytes.  RegistryTEEConnection._connect() runs
+    its real code (build_ssl_context_from_der → x402HttpxClient(verify=ctx))
+    so the full cert-pinning path is exercised with an actual TLS handshake.
+    """
+
+    async def test_connect_succeeds_with_matching_cert(self, tls_server):
+        mock_reg = _registry_with_real_cert(tls_server)
+        conn = RegistryTEEConnection(x402_client=x402Client(), registry=mock_reg)
+
+        resp = await conn.get().http_client.get(f"https://127.0.0.1:{tls_server['port']}/")
+        assert resp.status_code == 200
+        assert conn.get().tee_id == "tee-real"
+        assert conn.get().payment_address == "0xRealPay"
+        await conn.close()
+
+    async def test_connect_fails_with_wrong_cert(self, tls_server):
+        wrong_der, _, _ = _make_self_signed_cert()  # different key pair
+        mock_reg = _mock_registry_with_tee(
+            endpoint=f"https://127.0.0.1:{tls_server['port']}",
+            tls_cert_der=wrong_der,
+        )
+        conn = RegistryTEEConnection(x402_client=x402Client(), registry=mock_reg)
+
+        with pytest.raises(httpx.ConnectError):
+            await conn.get().http_client.get(f"https://127.0.0.1:{tls_server['port']}/")
+        await conn.close()
+
+    async def test_connect_fails_with_no_cert_pinning(self, tls_server):
+        """Without a pinned cert (tls_cert_der=None), system CAs are used
+        which won't trust our self-signed server cert."""
+        mock_reg = _mock_registry_with_tee(
+            endpoint=f"https://127.0.0.1:{tls_server['port']}",
+            tls_cert_der=None,
+        )
+        conn = RegistryTEEConnection(x402_client=x402Client(), registry=mock_reg)
+
+        with pytest.raises(httpx.ConnectError):
+            await conn.get().http_client.get(f"https://127.0.0.1:{tls_server['port']}/")
+        await conn.close()
+
+    async def test_reconnect_picks_up_new_cert(self, tls_server):
+        """After reconnect, the connection uses the freshly-resolved cert."""
+        mock_reg = _registry_with_real_cert(tls_server)
+        conn = RegistryTEEConnection(x402_client=x402Client(), registry=mock_reg)
+
+        await conn.reconnect()
+
+        resp = await conn.get().http_client.get(f"https://127.0.0.1:{tls_server['port']}/")
+        assert resp.status_code == 200
+        await conn.close()