From 9675f63bdabd3b3b6771ba839e33144a37f58476 Mon Sep 17 00:00:00 2001 From: Nolan Murphey Date: Fri, 17 Jul 2026 11:02:48 -0400 Subject: [PATCH 1/2] Add authTokenJsonField and additionalHeaders to ServiceDeployer registration Support extracting a bearer token from a JSON field of the auth-token secret, and forwarding static headers on every admin API request. Both options are opt-in and only forwarded to the custom resource when set, avoiding CFN property diffs for existing users. JSON extraction happens inside the handler so the plaintext token is never exposed as a custom-resource property. Co-Authored-By: Claude Code --- .../register-service-handler/index.mts | 52 +++++++++++++++++-- lib/restate-constructs/service-deployer.ts | 22 ++++++++ .../restate-constructs.test.ts.snap | 2 +- test/restate-constructs.test.ts | 52 +++++++++++++++++++ 4 files changed, 122 insertions(+), 6 deletions(-) diff --git a/lib/restate-constructs/register-service-handler/index.mts b/lib/restate-constructs/register-service-handler/index.mts index 8362cef..66d8ab1 100644 --- a/lib/restate-constructs/register-service-handler/index.mts +++ b/lib/restate-constructs/register-service-handler/index.mts @@ -40,6 +40,16 @@ export interface RegistrationProperties { */ authTokenSecretArn?: string; + /** + * When set, treat the secret value as a JSON object and extract the bearer token from this top-level field instead + * of using the raw secret string. Extraction happens here in the handler so the plaintext token is never exposed as + * a custom-resource property. + */ + authTokenJsonField?: string; + + /** Static headers to add to every admin API request. Reserved headers set by the handler take precedence. */ + additionalHeaders?: Record; + /** Not used by the handler, purely used to trick CloudFormation to perform an update when it otherwise would not. */ configurationVersion?: string; @@ -197,7 +207,7 @@ export const handler = async function (event: CloudFormationCustomResourceEvent, let authHeader: Record = {}; try { - authHeader = await createAuthHeader(props); + authHeader = await buildBaseHeaders(props); } catch (e) { console.warn(`Failed to load auth token for deletion: ${(e as Error)?.message}`); console.warn("Proceeding with deletion without auth header."); @@ -234,7 +244,7 @@ export const handler = async function (event: CloudFormationCustomResourceEvent, return; } - const authHeader = await createAuthHeader(props); + const authHeader = await buildBaseHeaders(props); let attempt; @@ -406,9 +416,21 @@ export const handler = async function (event: CloudFormationCustomResourceEvent, throw new Error(failureReason ?? "Restate service registration failed. Please see logs for details."); }; -async function createAuthHeader(props: RegistrationProperties): Promise> { +async function buildBaseHeaders(props: RegistrationProperties): Promise> { + // Static extra headers form the base; reserved headers set by the handler must always win, so we strip any + // caller-provided entries that collide with them (case-insensitive) before layering the auth header on top. + const reserved = new Set(["authorization", "content-type", "accept"]); + const baseHeaders: Record = {}; + for (const [key, value] of Object.entries(props.additionalHeaders ?? {})) { + if (reserved.has(key.toLowerCase())) { + console.warn(`Ignoring reserved header "${key}" from additionalHeaders.`); + continue; + } + baseHeaders[key] = value; + } + if (!props.authTokenSecretArn) { - return {}; + return baseHeaders; } console.log(`Using bearer authentication token from secret ${props.authTokenSecretArn}`); @@ -420,8 +442,28 @@ async function createAuthHeader(props: RegistrationProperties): Promise; + try { + parsed = JSON.parse(response.SecretString ?? ""); + } catch { + // Deliberately avoid echoing the secret material in the error. + throw new Error(`Secret value is not valid JSON; cannot extract field "${props.authTokenJsonField}".`); + } + const field = parsed?.[props.authTokenJsonField]; + if (typeof field !== "string") { + throw new Error( + `Secret JSON field "${props.authTokenJsonField}" is missing or not a string; cannot use it as a bearer token.`, + ); + } + token = field; + } + return { - Authorization: `Bearer ${response.SecretString}`, + ...baseHeaders, + Authorization: `Bearer ${token}`, }; } diff --git a/lib/restate-constructs/service-deployer.ts b/lib/restate-constructs/service-deployer.ts index 6a8452c..85543bc 100644 --- a/lib/restate-constructs/service-deployer.ts +++ b/lib/restate-constructs/service-deployer.ts @@ -30,6 +30,25 @@ export interface ServiceRegistrationProps { */ authToken?: secrets.ISecret; + /** + * When the {@link authToken} secret stores a JSON object rather than a raw string, extract the bearer token from + * this top-level field instead of using the whole secret value. For example, given a secret value of + * `{"token":"rst_xxx","version":3}`, set this to `"token"`. Only flat, top-level keys are supported; nested paths + * are not. + */ + authTokenJsonField?: string; + + /** + * Static headers to add to every admin API request made during registration (health check, deployment + * registration, service visibility patch, and any pruning/deletion queries). Useful for tagging requests or + * satisfying a proxy/gateway in front of the Restate admin endpoint. + * + * These are applied by the shipped handler and do not require bundling. Reserved headers set by the handler + * itself (`Authorization`, `Content-Type`, `Accept`, matched case-insensitively) take precedence over entries + * provided here. + */ + additionalHeaders?: Record; + /** * The external invoker role that Restate can assume to execute service handlers. If left unset, it's assumed that * the Restate deployment has sufficient permissions to invoke the handler directly. Takes precedence over the @@ -311,6 +330,9 @@ export class ServiceDeployer extends Construct { servicePath: serviceName, adminUrl: options?.adminUrl ?? environment.adminUrl, authTokenSecretArn: authToken?.secretArn, + // Forward JSON-field extraction and extra headers only when set, to avoid CFN property diffs for existing users. + ...(options?.authTokenJsonField !== undefined ? { authTokenJsonField: options.authTokenJsonField } : {}), + ...(options?.additionalHeaders !== undefined ? { additionalHeaders: options.additionalHeaders } : {}), serviceLambdaArn: handler.functionArn, invokeRoleArn: invokerRole?.roleArn, removalPolicy: options?.removalPolicy === cdk.RemovalPolicy.DESTROY ? "destroy" : ("retain" as const), diff --git a/test/__snapshots__/restate-constructs.test.ts.snap b/test/__snapshots__/restate-constructs.test.ts.snap index 459e71e..c1e1bd8 100644 --- a/test/__snapshots__/restate-constructs.test.ts.snap +++ b/test/__snapshots__/restate-constructs.test.ts.snap @@ -1851,7 +1851,7 @@ exports[`Restate constructs Service Deployer overrides 1`] = ` - arm64 Code: S3Bucket: cdk-hnb659fds-assets-account-id-region - S3Key: 595464e51d4a001ceaa6194ce91bd75f1475b9359b2b4da2271fe26a9d65260e.zip + S3Key: 0b7d6fa10d1fe685a7edcad5f8d2c1ec288a077408e3e36cd88feca84d7dee5b.zip Description: Restate custom registration handler Handler: entrypoint.handler MemorySize: 128 diff --git a/test/restate-constructs.test.ts b/test/restate-constructs.test.ts index a7b8f42..bd96186 100644 --- a/test/restate-constructs.test.ts +++ b/test/restate-constructs.test.ts @@ -303,6 +303,58 @@ describe("Restate constructs", () => { expect("healthCheckMaxBackoffSeconds" in customResource).toBe(false); }); + test("Service Deployer forwards authTokenJsonField and additionalHeaders to the custom resource", () => { + const app = new cdk.App(); + const stack = new cdk.Stack(app, "ServiceDeployerAuthOptions", { + env: { account: "account-id", region: "region" }, + }); + + const authToken = new secrets.Secret(stack, "RestateApiKey", { + secretStringValue: cdk.SecretValue.unsafePlainText('{"token":"rst_xxx"}'), + }); + + const restateEnvironment = RestateEnvironment.fromAttributes({ + adminUrl: "https://restate.example.com:9070", + authToken, + }); + + const handler = mockHandler(stack); + const serviceDeployer = new ServiceDeployer(stack, "ServiceDeployer", { + code: lambda.Code.fromAsset("dist/register-service-handler"), + }); + serviceDeployer.register(handler.currentVersion, restateEnvironment, { + authTokenJsonField: "token", + additionalHeaders: { "X-Deploy-Source": "cdk", "X-Env": "staging" }, + }); + + const properties = Template.fromStack(stack).findResources("Custom::RestateServiceDeployment"); + const customResource = Object.values(properties)[0]!.Properties as Record; + expect(customResource.authTokenJsonField).toBe("token"); + expect(customResource.additionalHeaders).toEqual({ "X-Deploy-Source": "cdk", "X-Env": "staging" }); + }); + + test("Service Deployer omits auth options when not set", () => { + const app = new cdk.App(); + const stack = new cdk.Stack(app, "ServiceDeployerAuthDefaults", { + env: { account: "account-id", region: "region" }, + }); + + const restateEnvironment = RestateEnvironment.fromAttributes({ + adminUrl: "https://restate.example.com:9070", + }); + + const handler = mockHandler(stack); + const serviceDeployer = new ServiceDeployer(stack, "ServiceDeployer", { + code: lambda.Code.fromAsset("dist/register-service-handler"), + }); + serviceDeployer.register(handler.currentVersion, restateEnvironment); + + const properties = Template.fromStack(stack).findResources("Custom::RestateServiceDeployment"); + const customResource = Object.values(properties)[0]!.Properties as Record; + expect("authTokenJsonField" in customResource).toBe(false); + expect("additionalHeaders" in customResource).toBe(false); + }); + test("[Experimental] Create a self-hosted Restate environment deployed on ECS Fargate", () => { const app = new cdk.App(); const stack = new cdk.Stack(app, "RestateOnFargateStack", { From a5a93ce28b7a70085e37534fc524198da93ffc0e Mon Sep 17 00:00:00 2001 From: Nolan Murphey Date: Fri, 17 Jul 2026 12:40:51 -0400 Subject: [PATCH 2/2] Let additionalHeaders override handler-set headers Simplify header handling: instead of stripping reserved headers (Authorization, Content-Type, Accept) from additionalHeaders, spread additionalHeaders last so a caller can override them when a proxy or gateway in front of the admin endpoint requires it. Co-Authored-By: Claude Code --- .../register-service-handler/index.mts | 23 ++++++++----------- lib/restate-constructs/service-deployer.ts | 5 ++-- .../restate-constructs.test.ts.snap | 2 +- 3 files changed, 12 insertions(+), 18 deletions(-) diff --git a/lib/restate-constructs/register-service-handler/index.mts b/lib/restate-constructs/register-service-handler/index.mts index 66d8ab1..df1d632 100644 --- a/lib/restate-constructs/register-service-handler/index.mts +++ b/lib/restate-constructs/register-service-handler/index.mts @@ -47,7 +47,11 @@ export interface RegistrationProperties { */ authTokenJsonField?: string; - /** Static headers to add to every admin API request. Reserved headers set by the handler take precedence. */ + /** + * Static headers to add to every admin API request. These take precedence over headers the handler sets itself + * (`Authorization`, `Content-Type`, `Accept`), so a caller can override them when a proxy or gateway in front of the + * admin endpoint requires it. Use standard header casing to override a handler-set header. + */ additionalHeaders?: Record; /** Not used by the handler, purely used to trick CloudFormation to perform an update when it otherwise would not. */ @@ -417,20 +421,10 @@ export const handler = async function (event: CloudFormationCustomResourceEvent, }; async function buildBaseHeaders(props: RegistrationProperties): Promise> { - // Static extra headers form the base; reserved headers set by the handler must always win, so we strip any - // caller-provided entries that collide with them (case-insensitive) before layering the auth header on top. - const reserved = new Set(["authorization", "content-type", "accept"]); - const baseHeaders: Record = {}; - for (const [key, value] of Object.entries(props.additionalHeaders ?? {})) { - if (reserved.has(key.toLowerCase())) { - console.warn(`Ignoring reserved header "${key}" from additionalHeaders.`); - continue; - } - baseHeaders[key] = value; - } + const additionalHeaders = props.additionalHeaders ?? {}; if (!props.authTokenSecretArn) { - return baseHeaders; + return { ...additionalHeaders }; } console.log(`Using bearer authentication token from secret ${props.authTokenSecretArn}`); @@ -461,9 +455,10 @@ async function buildBaseHeaders(props: RegistrationProperties): Promise; diff --git a/test/__snapshots__/restate-constructs.test.ts.snap b/test/__snapshots__/restate-constructs.test.ts.snap index c1e1bd8..d2ae982 100644 --- a/test/__snapshots__/restate-constructs.test.ts.snap +++ b/test/__snapshots__/restate-constructs.test.ts.snap @@ -1851,7 +1851,7 @@ exports[`Restate constructs Service Deployer overrides 1`] = ` - arm64 Code: S3Bucket: cdk-hnb659fds-assets-account-id-region - S3Key: 0b7d6fa10d1fe685a7edcad5f8d2c1ec288a077408e3e36cd88feca84d7dee5b.zip + S3Key: 1e12bdb6490da25e5bc5f2890d041456e776b60036f3c80482d625bd089988d1.zip Description: Restate custom registration handler Handler: entrypoint.handler MemorySize: 128