From 7be9bff6457bbe738fcc4174a2a6971ef3192e9b Mon Sep 17 00:00:00 2001 From: Dmitry Kropachev Date: Thu, 16 Jul 2026 15:55:31 -0400 Subject: [PATCH] Fix discovery drain timeout and session tokens --- README.md | 6 ++++-- src/client-base.ts | 22 ++++++++++++++++++-- src/discovery.ts | 46 ++++++++++++++++++++++++++++++++++++----- test/discovery.test.ts | 44 +++++++++++++++++++++++++++++++++++++++ test/middleware.test.ts | 29 +++++++++++++++++++++++++- 5 files changed, 137 insertions(+), 10 deletions(-) diff --git a/README.md b/README.md index 1131517..60b62f0 100644 --- a/README.md +++ b/README.md @@ -221,8 +221,10 @@ Header optimization is disabled by default. When enabled, headers are whitelisted, not removed by a strip list. The default whitelist is `Host`, `X-Amz-Target`, `Content-Length`, `Accept-Encoding`, and `Content-Encoding`; when credentials are configured, `Authorization` and -`X-Amz-Date` are also kept. The Alternator `User-Agent` is applied after this -filter, so it is kept unless `userAgent: false` is configured. +`X-Amz-Date` are also kept. Alternator does not use AWS session tokens, so +`sessionToken` is not sent even when provided in credentials. The Alternator +`User-Agent` is applied after this filter, so it is kept unless +`userAgent: false` is configured. By default, the client replaces the AWS SDK `User-Agent` with the ScyllaDB Alternator client identity: diff --git a/src/client-base.ts b/src/client-base.ts index c958fe3..aa66cdf 100644 --- a/src/client-base.ts +++ b/src/client-base.ts @@ -6,7 +6,7 @@ import { type ServiceOutputTypes, } from "@aws-sdk/client-dynamodb"; import type { HttpHandler, HttpHandlerUserInput } from "@smithy/protocol-http"; -import type { HttpHandlerOptions } from "@smithy/types"; +import type { AwsCredentialIdentity, HttpHandlerOptions } from "@smithy/types"; import { DEFAULT_REGION, firstEndpointUrl, NO_AUTH_CREDENTIALS, normalizeConfig } from "./config.js"; import { AlternatorDiscovery } from "./discovery.js"; import { KeyRouteAffinityPlanner } from "./affinity.js"; @@ -137,7 +137,7 @@ function buildDynamoConfig( ...awsConfig, endpoint: firstEndpointUrl(alternatorConfig), region: region ?? DEFAULT_REGION, - credentials: credentials ?? NO_AUTH_CREDENTIALS, + credentials: dropSessionToken(credentials) ?? NO_AUTH_CREDENTIALS, requestHandler, }; if (alternatorConfig.headerOptimization.enabled) { @@ -146,4 +146,22 @@ function buildDynamoConfig( return dynamoConfig; } +function dropSessionToken( + credentials: DynamoDBClientConfig["credentials"] | undefined, +): DynamoDBClientConfig["credentials"] | undefined { + if (credentials === undefined) { + return undefined; + } + if (typeof credentials === "function") { + return async (identityProperties?: Record) => + removeSessionToken(await credentials(identityProperties)); + } + return removeSessionToken(credentials); +} + +function removeSessionToken(credentials: AwsCredentialIdentity): AwsCredentialIdentity { + const { sessionToken: _sessionToken, ...withoutSessionToken } = credentials; + return withoutSessionToken; +} + export type AlternatorRequestHandler = HttpHandler; diff --git a/src/discovery.ts b/src/discovery.ts index 6ec112c..2eb05f3 100644 --- a/src/discovery.ts +++ b/src/discovery.ts @@ -250,7 +250,7 @@ export class AlternatorDiscovery { }); if (response.response.statusCode < 200 || response.response.statusCode >= 300) { - await drainResponseBody(response.response.body); + await drainResponseBody(response.response.body, this.config.discovery.timeoutMs); throw new Error(`/localnodes returned HTTP ${response.response.statusCode}`); } @@ -272,12 +272,48 @@ export class AlternatorDiscovery { } } -async function drainResponseBody(body: unknown): Promise { - try { - await bodyToString(body); - } catch (_error) { +async function drainResponseBody(body: unknown, timeoutMs: number): Promise { + const drain = bodyToString(body).then( + () => undefined, + () => undefined, + ); + let timeout: ReturnType | undefined; + const timeoutPromise = new Promise((resolve) => { + timeout = setTimeout(() => { + destroyResponseBody(body); + resolve(); + }, Math.max(1, timeoutMs)); + timeout.unref?.(); + }); + + await Promise.race([drain, timeoutPromise]); + if (timeout) { + clearTimeout(timeout); + } +} + +function destroyResponseBody(body: unknown): void { + if (typeof body !== "object" || body === null) { + return; + } + + if ("destroy" in body && typeof (body as { destroy?: unknown }).destroy === "function") { + try { + (body as { destroy(): void }).destroy(); + } catch (_error) { + return; + } return; } + + if ("cancel" in body && typeof (body as { cancel?: unknown }).cancel === "function") { + try { + const cancellation = (body as { cancel(): unknown }).cancel(); + void Promise.resolve(cancellation).catch(() => undefined); + } catch (_error) { + return; + } + } } function queryToRequestQuery(query: LocalNodesQuery): Record { diff --git a/test/discovery.test.ts b/test/discovery.test.ts index 4affc5c..1850f6d 100644 --- a/test/discovery.test.ts +++ b/test/discovery.test.ts @@ -325,6 +325,50 @@ describe("Alternator discovery", () => { } }); + it("bounds draining non-terminating non-2xx discovery bodies", async () => { + let requests = 0; + const server = createServer((request, response) => { + expect(request.url).toBe("/localnodes"); + requests += 1; + response.setHeader("content-type", "application/json"); + if (requests === 1) { + response.statusCode = 500; + response.write(JSON.stringify({ error: "temporary failure" })); + return; + } + response.end(JSON.stringify(["node-a.internal"])); + }); + const address = await listen(server); + const client = new AlternatorDynamoDBClient({ + seeds: [address.address, address.address], + port: address.port, + discovery: { + background: false, + timeoutMs: 20, + }, + connection: { + keepAlive: true, + maxSockets: 1, + }, + }); + + try { + await expect(client.alternator.refreshNodes()).resolves.toEqual([ + { + host: "node-a.internal", + scheme: "http", + port: address.port, + url: `http://node-a.internal:${address.port}`, + }, + ]); + expect(requests).toBe(2); + } finally { + client.destroy(); + server.closeAllConnections?.(); + await close(server); + } + }); + it("keeps the DynamoDB socket reusable after repeated non-2xx responses", async () => { let requests = 0; let connections = 0; diff --git a/test/middleware.test.ts b/test/middleware.test.ts index d391759..be13d5a 100644 --- a/test/middleware.test.ts +++ b/test/middleware.test.ts @@ -156,7 +156,7 @@ describe("Alternator middleware", () => { ]); }); - it("uses the default optimized header whitelist with session credentials", async () => { + it("drops session tokens before signing Alternator requests", async () => { const handler = new RecordingHandler(() => ({ TableNames: [] })); const client = new AlternatorDynamoDBClient({ seeds: ["seed"], @@ -177,6 +177,33 @@ describe("Alternator middleware", () => { expect(headers["x-amz-date"]).toBeDefined(); expect(headers["x-amz-target"]).toBe("DynamoDB_20120810.ListTables"); expect(headers["x-amz-security-token"]).toBeUndefined(); + expect(signedHeaderNames(headers.authorization)).toEqual([ + "content-length", + "host", + "x-amz-date", + "x-amz-target", + ]); + }); + + it("drops session tokens from credential providers before signing Alternator requests", async () => { + const handler = new RecordingHandler(() => ({ TableNames: [] })); + const client = new AlternatorDynamoDBClient({ + seeds: ["seed"], + requestHandler: handler, + discovery: { background: false }, + credentials: () => Promise.resolve({ + accessKeyId: "key", + secretAccessKey: "secret", + sessionToken: "session-token", + }), + }); + + await client.send(new ListTablesCommand({})); + + const headers = commandRequests(handler)[0]?.headers ?? {}; + expect(headers.authorization).toContain("AWS4-HMAC-SHA256"); + expect(headers["x-amz-security-token"]).toBeUndefined(); + expect(signedHeaderNames(headers.authorization)).not.toContain("x-amz-security-token"); }); it("compresses JSON request bodies when enabled", async () => {