Skip to content

Commit 99fd49e

Browse files
sunnywuclaude
sunnywu
and
claude
authored
UID2-6717: validate operator key against core service before enclave startup (#2477)
* UID2-6717: validate operator key against core service before enclave startup Adds a pre-flight POST to `{core_base_url}/attest` with the operator key in the Authorization header. The core service returns 401 before inspecting the attestation payload for any unrecognised key, so a 401 here unambiguously means the key is wrong. Any other response (400 for missing payload, 5xx, timeout) is logged and does not block startup, keeping the change safe to roll out without new failure modes. This catches subtle transcription errors (e.g. I vs l) that pass the existing regex and environment-alignment checks but are rejected at attestation time, saving operators from a confusing failure deep inside a running enclave. Adds unit tests covering: 401 rejection, 400 pass-through, 200 pass-through, connection error, timeout, unexpected exception, and endpoint URL/header correctness. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> * Rename validate_operator_key_with_service to validate_operator_key_with_core_service Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> --------- Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com>
1 parent 7d6e67c commit 99fd49e

2 files changed

Lines changed: 168 additions & 0 deletions

File tree

‎scripts/confidential_compute.py‎

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -51,6 +51,10 @@ class UID2ServicesUnreachableError(ConfidentialComputeStartupError):
5151
def __init__(self, cls, ip=None):
5252
super().__init__(error_name=f"E06: {self.__class__.__name__}", provider=cls, extra_message=ip)
5353

54+
class OperatorKeyRejectedError(ConfidentialComputeStartupError):
55+
def __init__(self, cls):
56+
super().__init__(error_name=f"E07: {self.__class__.__name__}", provider=cls)
57+
5458
class OperatorKeyPermissionError(ConfidentialComputeStartupError):
5559
def __init__(self, cls, message = None):
5660
super().__init__(error_name=f"E08: {self.__class__.__name__}", provider=cls, extra_message=message)
@@ -97,6 +101,30 @@ def validate_connectivity() -> None:
97101
raise UID2ServicesUnreachableError(self.__class__.__name__, core_ip)
98102
except Exception as e:
99103
raise UID2ServicesUnreachableError(self.__class__.__name__)
104+
105+
def validate_operator_key_with_core_service() -> None:
106+
"""Pre-flight check: verifies the operator key is accepted by the core service.
107+
POSTs to /attest with only the Authorization header; core returns 401 for an
108+
invalid key before it even inspects the attestation payload."""
109+
core_url = self.configs["core_base_url"]
110+
operator_key = self.configs.get("operator_key")
111+
try:
112+
response = requests.post(
113+
f"{core_url}/attest",
114+
headers={"Authorization": f"Bearer {operator_key}"},
115+
json={},
116+
timeout=5
117+
)
118+
if response.status_code == 401:
119+
logging.error(f"Operator key rejected by core service. Response: {response.text}")
120+
raise OperatorKeyRejectedError(self.__class__.__name__)
121+
logging.info(f"Operator key verified with core service (HTTP {response.status_code})")
122+
except OperatorKeyRejectedError:
123+
raise
124+
except (requests.ConnectionError, requests.Timeout) as e:
125+
logging.warning(f"Could not reach core service for key pre-verification: {e}")
126+
except Exception as e:
127+
logging.warning(f"Unexpected error during operator key pre-verification: {e}")
100128

101129
type_hints = get_type_hints(ConfidentialComputeConfig, include_extras=True)
102130
required_keys = [field for field, hint in type_hints.items() if "NotRequired" not in str(hint)]
@@ -115,6 +143,7 @@ def validate_connectivity() -> None:
115143
validate_url("optout_base_url", environment)
116144
validate_operator_key()
117145
validate_connectivity()
146+
validate_operator_key_with_core_service()
118147
logging.info("Completed static validation of confidential compute config values")
119148

120149
@abstractmethod

‎scripts/tests/test_confidential_compute.py‎

Lines changed: 139 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,139 @@
1+
import pytest
2+
import requests
3+
import sys
4+
import os
5+
from unittest.mock import patch, MagicMock
6+
7+
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
8+
from confidential_compute import (
9+
ConfidentialCompute,
10+
ConfidentialComputeConfig,
11+
OperatorKeyRejectedError,
12+
OperatorKeyValidationError,
13+
UID2ServicesUnreachableError,
14+
ConfigurationMissingError,
15+
ConfigurationValueError,
16+
)
17+
18+
19+
class ConcreteConfidentialCompute(ConfidentialCompute):
20+
"""Minimal concrete implementation for testing the base class."""
21+
22+
def _set_confidential_config(self, secret_identifier):
23+
pass
24+
25+
def _setup_auxiliaries(self):
26+
pass
27+
28+
def _validate_auxiliaries(self):
29+
pass
30+
31+
def run_compute(self):
32+
pass
33+
34+
35+
VALID_CONFIG = {
36+
"operator_key": "UID2-O-I-1-abcdefghijklmnop",
37+
"core_base_url": "https://core-integ.uidapi.com",
38+
"optout_base_url": "https://optout-integ.uidapi.com",
39+
"environment": "integ",
40+
"uid_instance_id_prefix": "ec2-abc123-ami-xyz",
41+
}
42+
43+
44+
def make_instance(config_overrides=None):
45+
cc = ConcreteConfidentialCompute()
46+
cc.configs = {**VALID_CONFIG, **(config_overrides or {})}
47+
return cc
48+
49+
50+
class TestValidateOperatorKeyWithService:
51+
"""Tests for the pre-flight operator key verification against the core service."""
52+
53+
def _run_validate(self, cc, mock_response):
54+
with patch("confidential_compute.socket.gethostbyname", return_value="1.2.3.4"), \
55+
patch("confidential_compute.requests.get") as mock_get, \
56+
patch("confidential_compute.requests.post", return_value=mock_response) as mock_post:
57+
mock_get.return_value = MagicMock(status_code=200)
58+
cc.validate_configuration()
59+
return mock_post
60+
61+
def test_invalid_key_raises_operator_key_rejected_error(self):
62+
cc = make_instance()
63+
mock_resp = MagicMock()
64+
mock_resp.status_code = 401
65+
mock_resp.text = '{"status":"Unauthorized"}'
66+
67+
with pytest.raises(OperatorKeyRejectedError):
68+
self._run_validate(cc, mock_resp)
69+
70+
def test_valid_key_with_no_payload_passes(self):
71+
cc = make_instance()
72+
mock_resp = MagicMock()
73+
mock_resp.status_code = 400 # valid key, missing attestation_request
74+
75+
self._run_validate(cc, mock_resp) # should not raise
76+
77+
def test_valid_key_200_response_passes(self):
78+
cc = make_instance()
79+
mock_resp = MagicMock()
80+
mock_resp.status_code = 200
81+
82+
self._run_validate(cc, mock_resp) # should not raise
83+
84+
def test_server_error_is_non_blocking(self):
85+
cc = make_instance()
86+
mock_resp = MagicMock()
87+
mock_resp.status_code = 500
88+
89+
self._run_validate(cc, mock_resp) # should not raise
90+
91+
def test_connection_error_is_non_blocking(self):
92+
cc = make_instance()
93+
94+
with patch("confidential_compute.socket.gethostbyname", return_value="1.2.3.4"), \
95+
patch("confidential_compute.requests.get") as mock_get, \
96+
patch("confidential_compute.requests.post", side_effect=requests.ConnectionError("refused")):
97+
mock_get.return_value = MagicMock(status_code=200)
98+
cc.validate_configuration() # should not raise
99+
100+
def test_timeout_is_non_blocking(self):
101+
cc = make_instance()
102+
103+
with patch("confidential_compute.socket.gethostbyname", return_value="1.2.3.4"), \
104+
patch("confidential_compute.requests.get") as mock_get, \
105+
patch("confidential_compute.requests.post", side_effect=requests.Timeout("timed out")):
106+
mock_get.return_value = MagicMock(status_code=200)
107+
cc.validate_configuration() # should not raise
108+
109+
def test_unexpected_exception_is_non_blocking(self):
110+
cc = make_instance()
111+
112+
with patch("confidential_compute.socket.gethostbyname", return_value="1.2.3.4"), \
113+
patch("confidential_compute.requests.get") as mock_get, \
114+
patch("confidential_compute.requests.post", side_effect=RuntimeError("unexpected")):
115+
mock_get.return_value = MagicMock(status_code=200)
116+
cc.validate_configuration() # should not raise
117+
118+
def test_post_sent_to_correct_endpoint(self):
119+
cc = make_instance()
120+
mock_resp = MagicMock()
121+
mock_resp.status_code = 400
122+
123+
mock_post = self._run_validate(cc, mock_resp)
124+
mock_post.assert_called_once_with(
125+
"https://core-integ.uidapi.com/attest",
126+
headers={"Authorization": f"Bearer {VALID_CONFIG['operator_key']}"},
127+
json={},
128+
timeout=5,
129+
)
130+
131+
def test_skip_validations_bypasses_key_service_check(self):
132+
cc = make_instance({"skip_validations": True})
133+
134+
with patch("confidential_compute.requests.post") as mock_post:
135+
# validate_configuration is not called when skip_validations is True
136+
# (the cloud scripts check this flag before calling validate_configuration)
137+
# But we can confirm the function itself is gated correctly:
138+
# skip_validations=True means validate_configuration() is never called.
139+
mock_post.assert_not_called()

0 commit comments

Comments
 (0)