diff --git a/strands-py/src/strands/vended_plugins/context_offloader/storage.py b/strands-py/src/strands/vended_plugins/context_offloader/storage.py index 652b7a8fb2..a984c5eea1 100644 --- a/strands-py/src/strands/vended_plugins/context_offloader/storage.py +++ b/strands-py/src/strands/vended_plugins/context_offloader/storage.py @@ -514,7 +514,10 @@ async def retrieve(self, reference: str) -> tuple[bytes, str]: """Retrieve content from an S3 object. Accepts both ``s3://`` URIs (as returned by ``store()``) and raw - S3 keys for backward compatibility. + S3 keys for backward compatibility. References are constrained to the + configured ``bucket`` and ``prefix``: a reference that resolves to a key + outside the prefix (or to a different bucket) is rejected, mirroring the + scope that ``store()`` enforces. Args: reference: The S3 URI or object key returned by store(). @@ -523,7 +526,8 @@ async def retrieve(self, reference: str) -> tuple[bytes, str]: A tuple of (content bytes, content type). Raises: - KeyError: If the object does not exist. + KeyError: If the object does not exist or the reference resolves + outside the configured bucket and prefix. """ s3_key = reference if reference.startswith("s3://"): @@ -531,6 +535,8 @@ async def retrieve(self, reference: str) -> tuple[bytes, str]: if not reference.startswith(expected_prefix): raise KeyError(f"Reference not found: {reference}") s3_key = reference[len(expected_prefix) :] + if self._prefix and not s3_key.startswith(self._prefix): + raise KeyError(f"Reference not found: {reference}") try: response = self._client.get_object(Bucket=self._bucket, Key=s3_key) content: bytes = response["Body"].read() diff --git a/strands-py/tests/strands/vended_plugins/context_offloader/test_storage.py b/strands-py/tests/strands/vended_plugins/context_offloader/test_storage.py index 053cd28a09..5a7b1b3c0f 100644 --- a/strands-py/tests/strands/vended_plugins/context_offloader/test_storage.py +++ b/strands-py/tests/strands/vended_plugins/context_offloader/test_storage.py @@ -445,6 +445,30 @@ async def test_retrieve_rejects_wrong_bucket_uri(self, storage): with pytest.raises(KeyError, match="Reference not found"): await storage.retrieve("s3://wrong-bucket/artifacts/some_key") + @pytest.mark.asyncio + async def test_retrieve_rejects_raw_key_outside_prefix(self, storage, mock_s3_client): + # Regression for report 3858663: a raw key pointing outside the configured + # prefix must not be readable, even when the object exists in the same bucket. + mock_s3_client.put_object(Bucket="test-bucket", Key="private/secret.txt", Body=b"SECRET") + with pytest.raises(KeyError, match="Reference not found"): + await storage.retrieve("private/secret.txt") + + @pytest.mark.asyncio + async def test_retrieve_rejects_s3_uri_outside_prefix(self, storage, mock_s3_client): + # Regression for report 3858663: a same-bucket s3:// URI outside the configured + # prefix must be rejected, not silently read. + mock_s3_client.put_object(Bucket="test-bucket", Key="private/secret.txt", Body=b"SECRET") + with pytest.raises(KeyError, match="Reference not found"): + await storage.retrieve("s3://test-bucket/private/secret.txt") + + @pytest.mark.asyncio + async def test_retrieve_rejects_sibling_prefix_key(self, storage, mock_s3_client): + # Regression for report 3858663: the trailing-slash normalization must make + # prefix "artifacts/" reject "artifactsX/...", not just fully-disjoint keys. + mock_s3_client.put_object(Bucket="test-bucket", Key="artifactsX/secret.txt", Body=b"SECRET") + with pytest.raises(KeyError, match="Reference not found"): + await storage.retrieve("artifactsX/secret.txt") + @pytest.mark.asyncio async def test_put_object_called_with_correct_params(self, storage, mock_s3_client): await storage.store("key_1", b"test content", "application/json") @@ -462,7 +486,7 @@ async def test_non_nosuchkey_error_propagates(self, storage, mock_s3_client): mock_s3_client.get_object.side_effect = ClientError(error_response, "GetObject") with pytest.raises(ClientError, match="Forbidden"): - await storage.retrieve("some_key") + await storage.retrieve("artifacts/some_key") class TestFileStorageWithSandbox: diff --git a/strands-ts/src/vended-plugins/context-offloader/__tests__/storage.test.ts b/strands-ts/src/vended-plugins/context-offloader/__tests__/storage.test.ts index b993b12e64..7f727db85b 100644 --- a/strands-ts/src/vended-plugins/context-offloader/__tests__/storage.test.ts +++ b/strands-ts/src/vended-plugins/context-offloader/__tests__/storage.test.ts @@ -228,12 +228,59 @@ describe('S3Storage', () => { expect(command.input.Key).toBe('some/raw/key') }) + it('round-trips a stored reference under a configured prefix', async () => { + mockSend.mockResolvedValueOnce({}).mockResolvedValueOnce({ + Body: { transformToByteArray: () => Promise.resolve(new TextEncoder().encode('hello')) }, + ContentType: 'text/plain', + }) + + const storage = new S3Storage('b', { prefix: 'artifacts', s3Client: mockS3Client as never }) + const ref = await storage.store('key1', new TextEncoder().encode('hello')) + const result = await storage.retrieve(ref) + + expect(new TextDecoder().decode(result.content)).toBe('hello') + }) + it('throws on bucket mismatch', async () => { const storage = new S3Storage('my-bucket', { s3Client: mockS3Client as never }) await expect(storage.retrieve('s3://wrong-bucket/key')).rejects.toThrow('bucket mismatch') }) + it('rejects a raw key outside the configured prefix', async () => { + const storage = new S3Storage('b', { prefix: 'artifacts', s3Client: mockS3Client as never }) + + await expect(storage.retrieve('private/secret.txt')).rejects.toThrow('Reference not found') + expect(mockSend).not.toHaveBeenCalled() + }) + + it('rejects a same-bucket s3:// URI outside the configured prefix', async () => { + const storage = new S3Storage('b', { prefix: 'artifacts', s3Client: mockS3Client as never }) + + await expect(storage.retrieve('s3://b/private/secret.txt')).rejects.toThrow('Reference not found') + expect(mockSend).not.toHaveBeenCalled() + }) + + it('rejects a sibling-prefix key that shares a string prefix but not a path segment', async () => { + const storage = new S3Storage('b', { prefix: 'artifacts', s3Client: mockS3Client as never }) + + await expect(storage.retrieve('artifactsX/secret.txt')).rejects.toThrow('Reference not found') + expect(mockSend).not.toHaveBeenCalled() + }) + + it('preserves a leading-slash prefix so legacy references remain retrievable', async () => { + mockSend.mockResolvedValue({ + Body: { transformToByteArray: () => Promise.resolve(new TextEncoder().encode('hello')) }, + ContentType: 'text/plain', + }) + const storage = new S3Storage('b', { prefix: '/artifacts', s3Client: mockS3Client as never }) + const result = await storage.retrieve('s3://b//artifacts/foo') + + expect(new TextDecoder().decode(result.content)).toBe('hello') + const command = mockSend.mock.calls[0]![0] + expect(command.input.Key).toBe('/artifacts/foo') + }) + it('throws on NoSuchKey error', async () => { const noSuchKey = new Error('not found') noSuchKey.name = 'NoSuchKey' diff --git a/strands-ts/src/vended-plugins/context-offloader/storage.ts b/strands-ts/src/vended-plugins/context-offloader/storage.ts index 5e250d3da9..ac1fb52bda 100644 --- a/strands-ts/src/vended-plugins/context-offloader/storage.ts +++ b/strands-ts/src/vended-plugins/context-offloader/storage.ts @@ -427,6 +427,11 @@ export class S3Storage implements Storage { s3Key = uriMatch[2] } + // Constrain references to the configured prefix, mirroring the scope store() enforces. + if (this._prefix && !s3Key.startsWith(this._prefix)) { + throw new Error(`Reference not found: ${reference}`) + } + try { const response = await client.send( new GetObjectCommand({