whiteroom-bedrock 0.1.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -0,0 +1,100 @@
1
+ Metadata-Version: 2.4
2
+ Name: whiteroom-bedrock
3
+ Version: 0.1.0
4
+ Summary: Attaches WhiteRoom session attestation to a boto3 bedrock-runtime client
5
+ License: BSL-1.1
6
+ Requires-Python: >=3.9
7
+ Description-Content-Type: text/markdown
8
+ Requires-Dist: botocore>=1.34
9
+ Provides-Extra: test
10
+ Requires-Dist: pytest>=7; extra == "test"
11
+ Requires-Dist: boto3>=1.34; extra == "test"
12
+
13
+ # whiteroom-bedrock
14
+
15
+ Client-side half of WhiteRoom's Bedrock adapter (server side lives in
16
+ `apps/engine/src/bedrock/`).
17
+
18
+ An unmodified AWS SDK has no way to prove its identity to WhiteRoom the way
19
+ `ANTHROPIC_BASE_URL`/`OPENAI_BASE_URL` callers do with a bearer key — every
20
+ Bedrock call is SigV4-signed by the caller's own credentials, and WhiteRoom
21
+ discards that signature (it re-signs outbound with its own assumed-role
22
+ credentials). Proving who called instead happens once, out of band, by
23
+ replaying a signed `sts:GetCallerIdentity` probe to WhiteRoom's `/attest`
24
+ endpoint. This package does that handshake and attaches the resulting
25
+ session token to every subsequent call, so the rest of your agent code never
26
+ has to know WhiteRoom is in the loop.
27
+
28
+ ## Install
29
+
30
+ ```bash
31
+ pip install -e sdks/bedrock-python
32
+ ```
33
+
34
+ ## Test
35
+
36
+ ```bash
37
+ pip install -e "sdks/bedrock-python[test]"
38
+ pytest sdks/bedrock-python/tests
39
+ ```
40
+
41
+ No network access or AWS credentials needed — everything's mocked at the
42
+ botocore/urllib layer. `tests/test_attest.py` is the one worth reading
43
+ closely: it directly checks the signed probe has `host`,
44
+ `x-whiteroom-server-id` and `x-amz-date` in `SignedHeaders`, which is
45
+ exactly what the server (`attest.ts`) requires — a change here that broke
46
+ that would otherwise only surface as every real attestation failing
47
+ against a live proxy.
48
+
49
+ ## Use
50
+
51
+ ```python
52
+ import boto3
53
+ import whiteroom_bedrock
54
+
55
+ client = boto3.client(
56
+ "bedrock-runtime",
57
+ region_name="us-east-1",
58
+ endpoint_url="https://proxy.whiteroom.tech",
59
+ )
60
+ whiteroom_bedrock.install(
61
+ client,
62
+ proxy_base_url="https://proxy.whiteroom.tech",
63
+ server_id="proxy.whiteroom.tech", # must match the proxy's configured serverId
64
+ region="us-east-1",
65
+ )
66
+
67
+ # Everything from here down is a normal bedrock-runtime client.
68
+ response = client.converse(
69
+ modelId="anthropic.claude-3-5-sonnet-20241022-v2:0",
70
+ messages=[{"role": "user", "content": [{"text": "hi"}]}],
71
+ )
72
+ ```
73
+
74
+ ## What `install()` actually does
75
+
76
+ - On first call (and again whenever the cached session goes stale), signs a
77
+ throwaway `GetCallerIdentity` request with whatever credentials this
78
+ process already resolves for AWS (env vars, shared config, an instance or
79
+ task role — the standard botocore chain) and POSTs it to `/attest`.
80
+ - Attaches the session token WhiteRoom hands back as `x-whiteroom-session`
81
+ on every `invoke_model`/`converse` call, via a `before-send` hook — after
82
+ signing, so it never interferes with the client's own (ultimately
83
+ discarded) signature.
84
+ - If the proxy ever answers 403 — the session expired, or WhiteRoom
85
+ restarted and lost its in-memory session table — invalidates the cached
86
+ token and forces exactly one retry, so callers don't see a transient
87
+ session death as a failure.
88
+
89
+ ## Not handled here
90
+
91
+ - **Streaming.** `invoke_model_with_response_stream` and
92
+ `converse_stream` reach the proxy the same as any other call, but the
93
+ server actively refuses to proxy streaming responses today (see
94
+ `routes.ts`) — you'll get a `ValidationException` back, not a stream.
95
+ - **Credential source mismatch.** `region` and credentials are resolved
96
+ from the default botocore chain unless you pass `region=`/
97
+ `botocore_session=` explicitly. If the client you're hooking was built
98
+ with credentials that *don't* come from that chain (e.g. explicit
99
+ `aws_access_key_id=` on `boto3.client(...)`), attest with the same
100
+ credentials by passing a `botocore.session.Session` that resolves them.
@@ -0,0 +1,88 @@
1
+ # whiteroom-bedrock
2
+
3
+ Client-side half of WhiteRoom's Bedrock adapter (server side lives in
4
+ `apps/engine/src/bedrock/`).
5
+
6
+ An unmodified AWS SDK has no way to prove its identity to WhiteRoom the way
7
+ `ANTHROPIC_BASE_URL`/`OPENAI_BASE_URL` callers do with a bearer key — every
8
+ Bedrock call is SigV4-signed by the caller's own credentials, and WhiteRoom
9
+ discards that signature (it re-signs outbound with its own assumed-role
10
+ credentials). Proving who called instead happens once, out of band, by
11
+ replaying a signed `sts:GetCallerIdentity` probe to WhiteRoom's `/attest`
12
+ endpoint. This package does that handshake and attaches the resulting
13
+ session token to every subsequent call, so the rest of your agent code never
14
+ has to know WhiteRoom is in the loop.
15
+
16
+ ## Install
17
+
18
+ ```bash
19
+ pip install -e sdks/bedrock-python
20
+ ```
21
+
22
+ ## Test
23
+
24
+ ```bash
25
+ pip install -e "sdks/bedrock-python[test]"
26
+ pytest sdks/bedrock-python/tests
27
+ ```
28
+
29
+ No network access or AWS credentials needed — everything's mocked at the
30
+ botocore/urllib layer. `tests/test_attest.py` is the one worth reading
31
+ closely: it directly checks the signed probe has `host`,
32
+ `x-whiteroom-server-id` and `x-amz-date` in `SignedHeaders`, which is
33
+ exactly what the server (`attest.ts`) requires — a change here that broke
34
+ that would otherwise only surface as every real attestation failing
35
+ against a live proxy.
36
+
37
+ ## Use
38
+
39
+ ```python
40
+ import boto3
41
+ import whiteroom_bedrock
42
+
43
+ client = boto3.client(
44
+ "bedrock-runtime",
45
+ region_name="us-east-1",
46
+ endpoint_url="https://proxy.whiteroom.tech",
47
+ )
48
+ whiteroom_bedrock.install(
49
+ client,
50
+ proxy_base_url="https://proxy.whiteroom.tech",
51
+ server_id="proxy.whiteroom.tech", # must match the proxy's configured serverId
52
+ region="us-east-1",
53
+ )
54
+
55
+ # Everything from here down is a normal bedrock-runtime client.
56
+ response = client.converse(
57
+ modelId="anthropic.claude-3-5-sonnet-20241022-v2:0",
58
+ messages=[{"role": "user", "content": [{"text": "hi"}]}],
59
+ )
60
+ ```
61
+
62
+ ## What `install()` actually does
63
+
64
+ - On first call (and again whenever the cached session goes stale), signs a
65
+ throwaway `GetCallerIdentity` request with whatever credentials this
66
+ process already resolves for AWS (env vars, shared config, an instance or
67
+ task role — the standard botocore chain) and POSTs it to `/attest`.
68
+ - Attaches the session token WhiteRoom hands back as `x-whiteroom-session`
69
+ on every `invoke_model`/`converse` call, via a `before-send` hook — after
70
+ signing, so it never interferes with the client's own (ultimately
71
+ discarded) signature.
72
+ - If the proxy ever answers 403 — the session expired, or WhiteRoom
73
+ restarted and lost its in-memory session table — invalidates the cached
74
+ token and forces exactly one retry, so callers don't see a transient
75
+ session death as a failure.
76
+
77
+ ## Not handled here
78
+
79
+ - **Streaming.** `invoke_model_with_response_stream` and
80
+ `converse_stream` reach the proxy the same as any other call, but the
81
+ server actively refuses to proxy streaming responses today (see
82
+ `routes.ts`) — you'll get a `ValidationException` back, not a stream.
83
+ - **Credential source mismatch.** `region` and credentials are resolved
84
+ from the default botocore chain unless you pass `region=`/
85
+ `botocore_session=` explicitly. If the client you're hooking was built
86
+ with credentials that *don't* come from that chain (e.g. explicit
87
+ `aws_access_key_id=` on `boto3.client(...)`), attest with the same
88
+ credentials by passing a `botocore.session.Session` that resolves them.
@@ -0,0 +1,23 @@
1
+ [build-system]
2
+ requires = ["setuptools>=68"]
3
+ build-backend = "setuptools.build_meta"
4
+
5
+ [project]
6
+ name = "whiteroom-bedrock"
7
+ version = "0.1.0"
8
+ description = "Attaches WhiteRoom session attestation to a boto3 bedrock-runtime client"
9
+ readme = "README.md"
10
+ requires-python = ">=3.9"
11
+ license = { text = "BSL-1.1" }
12
+ dependencies = [
13
+ "botocore>=1.34",
14
+ ]
15
+
16
+ [project.optional-dependencies]
17
+ test = [
18
+ "pytest>=7",
19
+ "boto3>=1.34",
20
+ ]
21
+
22
+ [tool.setuptools.packages.find]
23
+ where = ["src"]
@@ -0,0 +1,4 @@
1
+ [egg_info]
2
+ tag_build =
3
+ tag_date = 0
4
+
@@ -0,0 +1,10 @@
1
+ from .attest import AttestationError, AttestedIdentity, SessionManager
2
+ from .hook import SESSION_HEADER, install
3
+
4
+ __all__ = [
5
+ "install",
6
+ "SessionManager",
7
+ "AttestationError",
8
+ "AttestedIdentity",
9
+ "SESSION_HEADER",
10
+ ]
@@ -0,0 +1,146 @@
1
+ """
2
+ Client side of the attestation handshake described in
3
+ apps/engine/src/bedrock/attest.ts.
4
+
5
+ WhiteRoom cannot verify a request's SigV4 signature itself (that would need
6
+ the caller's secret key). Instead we sign a throwaway sts:GetCallerIdentity
7
+ request with whatever AWS credentials this process already has, and send it
8
+ to WhiteRoom's /attest endpoint. WhiteRoom replays it to STS, which can
9
+ verify it, and answers with our ARN plus a short-lived session token. That
10
+ token — not the probe — is what rides on every subsequent Bedrock call.
11
+ """
12
+
13
+ from __future__ import annotations
14
+
15
+ import json
16
+ import threading
17
+ import time
18
+ import urllib.error
19
+ import urllib.request
20
+ from dataclasses import dataclass
21
+
22
+ from botocore.auth import SigV4Auth
23
+ from botocore.awsrequest import AWSRequest
24
+ from botocore.session import Session as BotocoreSession
25
+
26
+ # The exact, minimal GetCallerIdentity body attest.ts's assertGetCallerIdentity
27
+ # expects: Action=GetCallerIdentity and nothing else but Version.
28
+ _GET_CALLER_IDENTITY_BODY = "Action=GetCallerIdentity&Version=2011-06-15"
29
+
30
+
31
+ class AttestationError(Exception):
32
+ """Raised when the probe can't be built, or WhiteRoom/STS rejects it."""
33
+
34
+
35
+ @dataclass(frozen=True)
36
+ class AttestedIdentity:
37
+ session_token: str
38
+ arn: str
39
+ issued_at: float
40
+
41
+
42
+ class SessionManager:
43
+ """
44
+ Holds one WhiteRoom session token per (proxy, server_id, region) and
45
+ refreshes it as needed.
46
+
47
+ /attest's response carries no expiry (see routes.ts — deliberately: the
48
+ server owns TTL policy, not the client), so freshness here is a mix of:
49
+ - proactive: refresh after SESSION_TTL_HINT locally, so a long-idle
50
+ process doesn't eat the round trip lazily on its very next call
51
+ - reactive: hook.py invalidates and forces one retry on a 403 from
52
+ the proxy, which covers a token dying earlier than the hint (server
53
+ restart, explicit revocation) without the caller doing anything
54
+
55
+ Thread-safe: a single lock guards refresh, so concurrent callers on a
56
+ cold cache queue behind one /attest call instead of each firing their
57
+ own (the same thundering-herd concern credentials.ts guards against
58
+ server-side for AssumeRole).
59
+ """
60
+
61
+ SESSION_TTL_HINT_SECONDS = 10 * 60
62
+
63
+ def __init__(
64
+ self,
65
+ proxy_base_url: str,
66
+ server_id: str,
67
+ region: str,
68
+ botocore_session: BotocoreSession | None = None,
69
+ timeout_seconds: float = 10.0,
70
+ ) -> None:
71
+ self._proxy_base_url = proxy_base_url.rstrip("/")
72
+ self._server_id = server_id
73
+ self._region = region
74
+ self._botocore_session = botocore_session or BotocoreSession()
75
+ self._timeout_seconds = timeout_seconds
76
+ self._lock = threading.Lock()
77
+ self._identity: AttestedIdentity | None = None
78
+
79
+ def get_token(self, force: bool = False) -> str:
80
+ with self._lock:
81
+ if not force and self._identity is not None:
82
+ age = time.time() - self._identity.issued_at
83
+ if age < self.SESSION_TTL_HINT_SECONDS:
84
+ return self._identity.session_token
85
+ self._identity = self._attest()
86
+ return self._identity.session_token
87
+
88
+ def invalidate(self) -> None:
89
+ with self._lock:
90
+ self._identity = None
91
+
92
+ def _build_probe(self) -> dict:
93
+ creds = self._botocore_session.get_credentials()
94
+ if creds is None:
95
+ raise AttestationError(
96
+ "no AWS credentials available to attest with — "
97
+ "whiteroom_bedrock needs the same credential chain your "
98
+ "bedrock-runtime client uses"
99
+ )
100
+ frozen = creds.get_frozen_credentials()
101
+
102
+ host = f"sts.{self._region}.amazonaws.com"
103
+ request = AWSRequest(
104
+ method="POST",
105
+ url=f"https://{host}/",
106
+ data=_GET_CALLER_IDENTITY_BODY,
107
+ headers={
108
+ "Host": host,
109
+ "Content-Type": "application/x-www-form-urlencoded",
110
+ # Must be present BEFORE signing — SigV4Auth signs whatever
111
+ # headers exist on the request at add_auth() time, and
112
+ # attest.ts refuses any probe where this (and host) aren't
113
+ # in SignedHeaders.
114
+ "x-whiteroom-server-id": self._server_id,
115
+ },
116
+ )
117
+ SigV4Auth(frozen, "sts", self._region).add_auth(request)
118
+
119
+ return {
120
+ "method": "POST",
121
+ "headers": dict(request.headers.items()),
122
+ "body": _GET_CALLER_IDENTITY_BODY,
123
+ }
124
+
125
+ def _attest(self) -> AttestedIdentity:
126
+ probe = self._build_probe()
127
+ req = urllib.request.Request(
128
+ f"{self._proxy_base_url}/attest",
129
+ data=json.dumps(probe).encode("utf-8"),
130
+ headers={"Content-Type": "application/json"},
131
+ method="POST",
132
+ )
133
+ try:
134
+ with urllib.request.urlopen(req, timeout=self._timeout_seconds) as resp:
135
+ payload = json.loads(resp.read())
136
+ except urllib.error.HTTPError as e:
137
+ detail = e.read().decode("utf-8", "replace")
138
+ raise AttestationError(f"WhiteRoom rejected attestation ({e.code}): {detail}") from e
139
+ except urllib.error.URLError as e:
140
+ raise AttestationError(f"could not reach WhiteRoom attest endpoint: {e}") from e
141
+
142
+ token = payload.get("sessionToken")
143
+ arn = payload.get("arn", "")
144
+ if not token:
145
+ raise AttestationError("WhiteRoom attest response missing sessionToken")
146
+ return AttestedIdentity(session_token=token, arn=arn, issued_at=time.time())
@@ -0,0 +1,92 @@
1
+ """
2
+ Wires session attachment into a boto3/botocore bedrock-runtime client.
3
+
4
+ An unmodified AWS SDK has no concept of WhiteRoom's session token — every
5
+ Bedrock call it makes is SigV4-signed with the caller's own credentials and
6
+ sent as-is. WhiteRoom's proxy discards that signature outright (it re-signs
7
+ outbound with its own assumed-role credentials — see sigv4.ts's STRIP list)
8
+ and instead identifies the caller from an `x-whiteroom-session` header. This
9
+ module is what attaches that header, transparently, to every call a hooked
10
+ client makes.
11
+ """
12
+
13
+ from __future__ import annotations
14
+
15
+ from botocore.client import BaseClient
16
+
17
+ from .attest import SessionManager
18
+
19
+ SESSION_HEADER = "x-whiteroom-session"
20
+
21
+ # Botocore's event-name slug for the Bedrock Runtime service. Fixed — this
22
+ # module only ever hooks that one service.
23
+ _SERVICE_EVENT_ID = "bedrock-runtime"
24
+
25
+ # Force at most one re-attest-and-retry per call. Anything past that is a
26
+ # real, persistent rejection (bad server_id, revoked tenant, ...) and should
27
+ # surface to the caller as a normal ClientError rather than loop.
28
+ _MAX_FORCED_RETRIES = 1
29
+
30
+
31
+ def install(
32
+ client: BaseClient,
33
+ *,
34
+ proxy_base_url: str,
35
+ server_id: str,
36
+ region: str | None = None,
37
+ botocore_session=None,
38
+ ) -> SessionManager:
39
+ """
40
+ Call once, right after constructing a bedrock-runtime client pointed at
41
+ a WhiteRoom proxy:
42
+
43
+ client = boto3.client(
44
+ "bedrock-runtime",
45
+ region_name="us-east-1",
46
+ endpoint_url="https://proxy.whiteroom.tech",
47
+ )
48
+ whiteroom_bedrock.install(
49
+ client,
50
+ proxy_base_url="https://proxy.whiteroom.tech",
51
+ server_id="proxy.whiteroom.tech",
52
+ region="us-east-1",
53
+ )
54
+
55
+ Everything past this point — invoke_model, converse, retries, streaming
56
+ refusals — flows through the client exactly as if it were talking to
57
+ real Bedrock. `region` defaults to the client's own configured region if
58
+ omitted; pass it explicitly if the client's credentials weren't resolved
59
+ from the default chain (e.g. an assumed role built by hand).
60
+
61
+ Returns the SessionManager backing this client, in case the caller wants
62
+ manual control (e.g. `manager.invalidate()` after rotating credentials).
63
+ """
64
+ resolved_region = region or client.meta.region_name
65
+ if not resolved_region:
66
+ raise ValueError("region could not be determined from the client; pass region= explicitly")
67
+
68
+ manager = SessionManager(
69
+ proxy_base_url=proxy_base_url,
70
+ server_id=server_id,
71
+ region=resolved_region,
72
+ botocore_session=botocore_session,
73
+ )
74
+
75
+ events = client.meta.events
76
+
77
+ def _attach_session(request, **kwargs):
78
+ request.headers[SESSION_HEADER] = manager.get_token()
79
+
80
+ def _retry_on_session_rejection(response, attempts, **kwargs):
81
+ if response is None or attempts > _MAX_FORCED_RETRIES:
82
+ return None
83
+ http_response, _parsed = response
84
+ if http_response is not None and http_response.status_code == 403:
85
+ manager.invalidate()
86
+ return 0 # retry immediately; _attach_session will fetch a fresh token
87
+ return None
88
+
89
+ events.register(f"before-send.{_SERVICE_EVENT_ID}", _attach_session)
90
+ events.register(f"needs-retry.{_SERVICE_EVENT_ID}", _retry_on_session_rejection)
91
+
92
+ return manager
@@ -0,0 +1,100 @@
1
+ Metadata-Version: 2.4
2
+ Name: whiteroom-bedrock
3
+ Version: 0.1.0
4
+ Summary: Attaches WhiteRoom session attestation to a boto3 bedrock-runtime client
5
+ License: BSL-1.1
6
+ Requires-Python: >=3.9
7
+ Description-Content-Type: text/markdown
8
+ Requires-Dist: botocore>=1.34
9
+ Provides-Extra: test
10
+ Requires-Dist: pytest>=7; extra == "test"
11
+ Requires-Dist: boto3>=1.34; extra == "test"
12
+
13
+ # whiteroom-bedrock
14
+
15
+ Client-side half of WhiteRoom's Bedrock adapter (server side lives in
16
+ `apps/engine/src/bedrock/`).
17
+
18
+ An unmodified AWS SDK has no way to prove its identity to WhiteRoom the way
19
+ `ANTHROPIC_BASE_URL`/`OPENAI_BASE_URL` callers do with a bearer key — every
20
+ Bedrock call is SigV4-signed by the caller's own credentials, and WhiteRoom
21
+ discards that signature (it re-signs outbound with its own assumed-role
22
+ credentials). Proving who called instead happens once, out of band, by
23
+ replaying a signed `sts:GetCallerIdentity` probe to WhiteRoom's `/attest`
24
+ endpoint. This package does that handshake and attaches the resulting
25
+ session token to every subsequent call, so the rest of your agent code never
26
+ has to know WhiteRoom is in the loop.
27
+
28
+ ## Install
29
+
30
+ ```bash
31
+ pip install -e sdks/bedrock-python
32
+ ```
33
+
34
+ ## Test
35
+
36
+ ```bash
37
+ pip install -e "sdks/bedrock-python[test]"
38
+ pytest sdks/bedrock-python/tests
39
+ ```
40
+
41
+ No network access or AWS credentials needed — everything's mocked at the
42
+ botocore/urllib layer. `tests/test_attest.py` is the one worth reading
43
+ closely: it directly checks the signed probe has `host`,
44
+ `x-whiteroom-server-id` and `x-amz-date` in `SignedHeaders`, which is
45
+ exactly what the server (`attest.ts`) requires — a change here that broke
46
+ that would otherwise only surface as every real attestation failing
47
+ against a live proxy.
48
+
49
+ ## Use
50
+
51
+ ```python
52
+ import boto3
53
+ import whiteroom_bedrock
54
+
55
+ client = boto3.client(
56
+ "bedrock-runtime",
57
+ region_name="us-east-1",
58
+ endpoint_url="https://proxy.whiteroom.tech",
59
+ )
60
+ whiteroom_bedrock.install(
61
+ client,
62
+ proxy_base_url="https://proxy.whiteroom.tech",
63
+ server_id="proxy.whiteroom.tech", # must match the proxy's configured serverId
64
+ region="us-east-1",
65
+ )
66
+
67
+ # Everything from here down is a normal bedrock-runtime client.
68
+ response = client.converse(
69
+ modelId="anthropic.claude-3-5-sonnet-20241022-v2:0",
70
+ messages=[{"role": "user", "content": [{"text": "hi"}]}],
71
+ )
72
+ ```
73
+
74
+ ## What `install()` actually does
75
+
76
+ - On first call (and again whenever the cached session goes stale), signs a
77
+ throwaway `GetCallerIdentity` request with whatever credentials this
78
+ process already resolves for AWS (env vars, shared config, an instance or
79
+ task role — the standard botocore chain) and POSTs it to `/attest`.
80
+ - Attaches the session token WhiteRoom hands back as `x-whiteroom-session`
81
+ on every `invoke_model`/`converse` call, via a `before-send` hook — after
82
+ signing, so it never interferes with the client's own (ultimately
83
+ discarded) signature.
84
+ - If the proxy ever answers 403 — the session expired, or WhiteRoom
85
+ restarted and lost its in-memory session table — invalidates the cached
86
+ token and forces exactly one retry, so callers don't see a transient
87
+ session death as a failure.
88
+
89
+ ## Not handled here
90
+
91
+ - **Streaming.** `invoke_model_with_response_stream` and
92
+ `converse_stream` reach the proxy the same as any other call, but the
93
+ server actively refuses to proxy streaming responses today (see
94
+ `routes.ts`) — you'll get a `ValidationException` back, not a stream.
95
+ - **Credential source mismatch.** `region` and credentials are resolved
96
+ from the default botocore chain unless you pass `region=`/
97
+ `botocore_session=` explicitly. If the client you're hooking was built
98
+ with credentials that *don't* come from that chain (e.g. explicit
99
+ `aws_access_key_id=` on `boto3.client(...)`), attest with the same
100
+ credentials by passing a `botocore.session.Session` that resolves them.
@@ -0,0 +1,12 @@
1
+ README.md
2
+ pyproject.toml
3
+ src/whiteroom_bedrock/__init__.py
4
+ src/whiteroom_bedrock/attest.py
5
+ src/whiteroom_bedrock/hook.py
6
+ src/whiteroom_bedrock.egg-info/PKG-INFO
7
+ src/whiteroom_bedrock.egg-info/SOURCES.txt
8
+ src/whiteroom_bedrock.egg-info/dependency_links.txt
9
+ src/whiteroom_bedrock.egg-info/requires.txt
10
+ src/whiteroom_bedrock.egg-info/top_level.txt
11
+ tests/test_attest.py
12
+ tests/test_hook.py
@@ -0,0 +1,5 @@
1
+ botocore>=1.34
2
+
3
+ [test]
4
+ pytest>=7
5
+ boto3>=1.34
@@ -0,0 +1 @@
1
+ whiteroom_bedrock
@@ -0,0 +1,190 @@
1
+ """
2
+ Tests the probe construction and session caching in attest.py.
3
+
4
+ The probe-shape tests are the important ones: attest.ts on the server side
5
+ specifically requires host, x-whiteroom-server-id and x-amz-date to all be
6
+ in SignedHeaders (see apps/engine/src/bedrock/attest.ts and its own test
7
+ suite, bedrock-attest.test.ts). A change here that stops signing one of
8
+ those would only ever surface as every real attestation failing against a
9
+ live proxy -- this is the one place that's caught locally instead.
10
+ """
11
+
12
+ import json
13
+ import time
14
+ from unittest.mock import patch
15
+ from urllib.error import HTTPError, URLError
16
+
17
+ import pytest
18
+ from botocore.credentials import Credentials
19
+
20
+ from whiteroom_bedrock.attest import AttestationError, SessionManager
21
+
22
+
23
+ class FakeCredentialSession:
24
+ """Stands in for botocore.session.Session -- just get_credentials()."""
25
+
26
+ def __init__(self, creds=None):
27
+ self._creds = creds if creds is not None else Credentials(
28
+ "AKIAFAKEEXAMPLE", "fakesecret1234567890", token="fake-session-token"
29
+ )
30
+
31
+ def get_credentials(self):
32
+ return self._creds
33
+
34
+
35
+ def make_manager(**overrides):
36
+ kwargs = dict(
37
+ proxy_base_url="https://proxy.whiteroom.tech",
38
+ server_id="proxy.whiteroom.tech",
39
+ region="us-east-1",
40
+ botocore_session=FakeCredentialSession(),
41
+ )
42
+ kwargs.update(overrides)
43
+ return SessionManager(**kwargs)
44
+
45
+
46
+ class TestBuildProbe:
47
+ def test_method_and_body_are_exactly_get_caller_identity(self):
48
+ probe = make_manager()._build_probe()
49
+ assert probe["method"] == "POST"
50
+ assert probe["body"] == "Action=GetCallerIdentity&Version=2011-06-15"
51
+
52
+ def test_signed_headers_include_everything_the_server_requires(self):
53
+ # Direct regression test for the server-side fix: attest.ts rejects
54
+ # a probe unless host, x-whiteroom-server-id AND x-amz-date are all
55
+ # in SignedHeaders (not just present as headers).
56
+ probe = make_manager()._build_probe()
57
+ auth = probe["headers"]["Authorization"]
58
+ signed_headers = auth.split("SignedHeaders=")[1].split(",")[0].split(";")
59
+ for required in ("host", "x-whiteroom-server-id", "x-amz-date"):
60
+ assert required in signed_headers, f"{required} missing from SignedHeaders: {signed_headers}"
61
+
62
+ def test_server_id_header_matches_what_was_configured(self):
63
+ probe = make_manager(server_id="my-proxy.example.com")._build_probe()
64
+ assert probe["headers"]["x-whiteroom-server-id"] == "my-proxy.example.com"
65
+
66
+ def test_host_targets_sts_in_the_configured_region(self):
67
+ probe = make_manager(region="eu-west-1")._build_probe()
68
+ assert probe["headers"]["Host"] == "sts.eu-west-1.amazonaws.com"
69
+
70
+ def test_includes_security_token_for_temporary_credentials(self):
71
+ probe = make_manager()._build_probe()
72
+ assert probe["headers"]["X-Amz-Security-Token"] == "fake-session-token"
73
+
74
+ def test_omits_security_token_for_static_credentials(self):
75
+ static_creds = Credentials("AKIASTATIC", "staticsecret")
76
+ probe = make_manager(botocore_session=FakeCredentialSession(static_creds))._build_probe()
77
+ assert "X-Amz-Security-Token" not in probe["headers"]
78
+
79
+ def test_raises_when_no_credentials_are_available(self):
80
+ class NoCreds:
81
+ def get_credentials(self):
82
+ return None
83
+
84
+ with pytest.raises(AttestationError, match="no AWS credentials"):
85
+ make_manager(botocore_session=NoCreds())._build_probe()
86
+
87
+
88
+ class TestGetToken:
89
+ def test_caches_within_the_ttl_window(self):
90
+ manager = make_manager()
91
+ calls = []
92
+
93
+ def fake_attest():
94
+ calls.append(1)
95
+ from whiteroom_bedrock.attest import AttestedIdentity
96
+ return AttestedIdentity(session_token=f"tok-{len(calls)}", arn="arn:aws:iam::123:user/a", issued_at=time.time())
97
+
98
+ with patch.object(manager, "_attest", side_effect=fake_attest):
99
+ first = manager.get_token()
100
+ second = manager.get_token()
101
+
102
+ assert first == second == "tok-1"
103
+ assert len(calls) == 1
104
+
105
+ def test_force_refreshes_even_within_the_ttl_window(self):
106
+ manager = make_manager()
107
+ calls = []
108
+
109
+ def fake_attest():
110
+ calls.append(1)
111
+ from whiteroom_bedrock.attest import AttestedIdentity
112
+ return AttestedIdentity(session_token=f"tok-{len(calls)}", arn="arn:aws:iam::123:user/a", issued_at=time.time())
113
+
114
+ with patch.object(manager, "_attest", side_effect=fake_attest):
115
+ manager.get_token()
116
+ second = manager.get_token(force=True)
117
+
118
+ assert second == "tok-2"
119
+ assert len(calls) == 2
120
+
121
+ def test_refreshes_once_the_ttl_hint_has_passed(self):
122
+ manager = make_manager()
123
+ calls = []
124
+
125
+ def fake_attest():
126
+ calls.append(1)
127
+ from whiteroom_bedrock.attest import AttestedIdentity
128
+ # issued_at backdated past the TTL hint, as if time had passed.
129
+ return AttestedIdentity(
130
+ session_token=f"tok-{len(calls)}",
131
+ arn="arn:aws:iam::123:user/a",
132
+ issued_at=time.time() - SessionManager.SESSION_TTL_HINT_SECONDS - 1,
133
+ )
134
+
135
+ with patch.object(manager, "_attest", side_effect=fake_attest):
136
+ manager.get_token()
137
+ second = manager.get_token()
138
+
139
+ assert second == "tok-2"
140
+ assert len(calls) == 2
141
+
142
+ def test_invalidate_forces_a_fresh_attest_on_the_next_call(self):
143
+ manager = make_manager()
144
+ calls = []
145
+
146
+ def fake_attest():
147
+ calls.append(1)
148
+ from whiteroom_bedrock.attest import AttestedIdentity
149
+ return AttestedIdentity(session_token=f"tok-{len(calls)}", arn="arn:aws:iam::123:user/a", issued_at=time.time())
150
+
151
+ with patch.object(manager, "_attest", side_effect=fake_attest):
152
+ manager.get_token()
153
+ manager.invalidate()
154
+ second = manager.get_token()
155
+
156
+ assert second == "tok-2"
157
+ assert len(calls) == 2
158
+
159
+
160
+ class TestAttestErrorHandling:
161
+ def test_wraps_an_http_error_from_the_proxy(self):
162
+ manager = make_manager()
163
+ err = HTTPError("https://proxy.whiteroom.tech/attest", 401, "Unauthorized", {}, None)
164
+ err.read = lambda: b'{"message":"server-ID mismatch"}'
165
+ with patch("urllib.request.urlopen", side_effect=err):
166
+ with pytest.raises(AttestationError, match="401"):
167
+ manager.get_token()
168
+
169
+ def test_wraps_a_network_error_reaching_the_proxy(self):
170
+ manager = make_manager()
171
+ with patch("urllib.request.urlopen", side_effect=URLError("connection refused")):
172
+ with pytest.raises(AttestationError, match="could not reach"):
173
+ manager.get_token()
174
+
175
+ def test_raises_when_the_proxy_response_has_no_session_token(self):
176
+ manager = make_manager()
177
+
178
+ class FakeResponse:
179
+ def read(self):
180
+ return json.dumps({"arn": "arn:aws:iam::123:user/a"}).encode()
181
+
182
+ def __enter__(self):
183
+ return self
184
+
185
+ def __exit__(self, *a):
186
+ return False
187
+
188
+ with patch("urllib.request.urlopen", return_value=FakeResponse()):
189
+ with pytest.raises(AttestationError, match="missing sessionToken"):
190
+ manager.get_token()
@@ -0,0 +1,139 @@
1
+ """
2
+ Tests install()'s session-attachment and 403-retry logic in hook.py.
3
+
4
+ Uses a fake botocore client (just enough of .meta.events.register to capture
5
+ the handlers install() wires up) rather than a real boto3 client, so these
6
+ run without network access or a real AWS SDK event bus -- the actual
7
+ end-to-end path (a real boto3 client, a real HTTP server) was validated
8
+ manually when this adapter was built; these pin the handler logic itself so
9
+ a future change can't silently break it.
10
+ """
11
+
12
+ from unittest.mock import MagicMock, patch
13
+
14
+ import pytest
15
+
16
+ from whiteroom_bedrock.hook import SESSION_HEADER, install
17
+
18
+
19
+ class FakeEvents:
20
+ """Records registrations instead of actually dispatching anything."""
21
+
22
+ def __init__(self):
23
+ self.registered = {}
24
+
25
+ def register(self, event_name, handler):
26
+ self.registered[event_name] = handler
27
+
28
+
29
+ class FakeMeta:
30
+ def __init__(self, region_name="us-east-1"):
31
+ self.region_name = region_name
32
+ self.events = FakeEvents()
33
+
34
+
35
+ class FakeClient:
36
+ def __init__(self, region_name="us-east-1"):
37
+ self.meta = FakeMeta(region_name)
38
+
39
+
40
+ class FakeRequest:
41
+ def __init__(self):
42
+ self.headers = {}
43
+
44
+
45
+ def install_with_fake_manager(client=None, **overrides):
46
+ """install() always builds its own real SessionManager; tests patch its
47
+ get_token/invalidate afterward via the manager it returns."""
48
+ client = client or FakeClient()
49
+ kwargs = dict(proxy_base_url="https://proxy.whiteroom.tech", server_id="proxy.whiteroom.tech")
50
+ kwargs.update(overrides)
51
+ manager = install(client, **kwargs)
52
+ return client, manager
53
+
54
+
55
+ class TestInstall:
56
+ def test_registers_before_send_and_needs_retry_for_bedrock_runtime(self):
57
+ client, _ = install_with_fake_manager()
58
+ assert "before-send.bedrock-runtime" in client.meta.events.registered
59
+ assert "needs-retry.bedrock-runtime" in client.meta.events.registered
60
+
61
+ def test_resolves_region_from_the_client_when_not_given_explicitly(self):
62
+ client = FakeClient(region_name="eu-west-1")
63
+ _, manager = install_with_fake_manager(client)
64
+ assert manager._region == "eu-west-1"
65
+
66
+ def test_explicit_region_overrides_the_client_default(self):
67
+ client = FakeClient(region_name="eu-west-1")
68
+ _, manager = install_with_fake_manager(client, region="ap-southeast-2")
69
+ assert manager._region == "ap-southeast-2"
70
+
71
+ def test_raises_when_no_region_can_be_determined(self):
72
+ client = FakeClient(region_name=None)
73
+ with pytest.raises(ValueError, match="region"):
74
+ install(client, proxy_base_url="https://proxy.whiteroom.tech", server_id="proxy.whiteroom.tech")
75
+
76
+ def test_returns_the_session_manager_for_manual_control(self):
77
+ _, manager = install_with_fake_manager()
78
+ assert hasattr(manager, "get_token")
79
+ assert hasattr(manager, "invalidate")
80
+
81
+
82
+ class TestAttachSession:
83
+ def test_sets_the_session_header_from_get_token(self):
84
+ client, manager = install_with_fake_manager()
85
+ with patch.object(manager, "get_token", return_value="wrbs_testtoken123"):
86
+ handler = client.meta.events.registered["before-send.bedrock-runtime"]
87
+ request = FakeRequest()
88
+ handler(request)
89
+ assert request.headers[SESSION_HEADER] == "wrbs_testtoken123"
90
+
91
+ def test_does_not_force_a_refresh_on_every_call(self):
92
+ # get_token() is called with no force argument -- the hook relies on
93
+ # SessionManager's own TTL caching, it doesn't re-attest every call.
94
+ client, manager = install_with_fake_manager()
95
+ with patch.object(manager, "get_token", return_value="tok") as mock_get_token:
96
+ handler = client.meta.events.registered["before-send.bedrock-runtime"]
97
+ handler(FakeRequest())
98
+ mock_get_token.assert_called_once_with()
99
+
100
+
101
+ class TestRetryOnSessionRejection:
102
+ def _invoke(self, client, response, attempts):
103
+ handler = client.meta.events.registered["needs-retry.bedrock-runtime"]
104
+ return handler(response=response, attempts=attempts)
105
+
106
+ def test_forces_an_immediate_retry_on_the_first_403(self):
107
+ client, manager = install_with_fake_manager()
108
+ http_response = MagicMock(status_code=403)
109
+ with patch.object(manager, "invalidate") as mock_invalidate:
110
+ result = self._invoke(client, (http_response, None), attempts=1)
111
+ assert result == 0
112
+ mock_invalidate.assert_called_once()
113
+
114
+ def test_does_not_force_a_second_retry(self):
115
+ # Two failures in a row is treated as a real, persistent rejection
116
+ # (bad server_id, revoked tenant) -- not something re-attesting a
117
+ # second time will fix.
118
+ client, manager = install_with_fake_manager()
119
+ http_response = MagicMock(status_code=403)
120
+ with patch.object(manager, "invalidate") as mock_invalidate:
121
+ result = self._invoke(client, (http_response, None), attempts=2)
122
+ assert result is None
123
+ mock_invalidate.assert_not_called()
124
+
125
+ def test_ignores_non_403_statuses(self):
126
+ client, manager = install_with_fake_manager()
127
+ http_response = MagicMock(status_code=500)
128
+ with patch.object(manager, "invalidate") as mock_invalidate:
129
+ result = self._invoke(client, (http_response, None), attempts=1)
130
+ assert result is None
131
+ mock_invalidate.assert_not_called()
132
+
133
+ def test_ignores_a_missing_response(self):
134
+ client, _ = install_with_fake_manager()
135
+ assert self._invoke(client, None, attempts=1) is None
136
+
137
+ def test_ignores_a_missing_http_response(self):
138
+ client, _ = install_with_fake_manager()
139
+ assert self._invoke(client, (None, None), attempts=1) is None