diff --git a/lib/restate-constructs/register-service-handler/index.mts b/lib/restate-constructs/register-service-handler/index.mts index 8362cef..df1d632 100644 --- a/lib/restate-constructs/register-service-handler/index.mts +++ b/lib/restate-constructs/register-service-handler/index.mts @@ -40,6 +40,20 @@ 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. 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. */ configurationVersion?: string; @@ -197,7 +211,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 +248,7 @@ export const handler = async function (event: CloudFormationCustomResourceEvent, return; } - const authHeader = await createAuthHeader(props); + const authHeader = await buildBaseHeaders(props); let attempt; @@ -406,9 +420,11 @@ 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> { + const additionalHeaders = props.additionalHeaders ?? {}; + if (!props.authTokenSecretArn) { - return {}; + return { ...additionalHeaders }; } console.log(`Using bearer authentication token from secret ${props.authTokenSecretArn}`); @@ -420,8 +436,29 @@ 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; + } + + // Spread additionalHeaders last so a caller can override the Authorization header if they need to. return { - Authorization: `Bearer ${response.SecretString}`, + Authorization: `Bearer ${token}`, + ...additionalHeaders, }; } diff --git a/lib/restate-constructs/service-deployer.ts b/lib/restate-constructs/service-deployer.ts index 6a8452c..58dc204 100644 --- a/lib/restate-constructs/service-deployer.ts +++ b/lib/restate-constructs/service-deployer.ts @@ -30,6 +30,24 @@ 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. They take precedence over headers the + * handler sets itself (`Authorization`, `Content-Type`, `Accept`); use standard header casing to override one. + */ + 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 +329,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..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: 595464e51d4a001ceaa6194ce91bd75f1475b9359b2b4da2271fe26a9d65260e.zip + S3Key: 1e12bdb6490da25e5bc5f2890d041456e776b60036f3c80482d625bd089988d1.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", {