Add JSON-field auth token extraction and additional header support to ServiceDeployer

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<string, string>;
+
   /** 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<string, string> = {};
     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<Record<string, string>> {
+async function buildBaseHeaders(props: RegistrationProperties): Promise<Record<string, string>> {
+  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<Record<s
   );
 
   console.log(`Successfully retrieved secret "${response.Name}" version ${response.VersionId}`);
+
+  let token = response.SecretString;
+  if (props.authTokenJsonField) {
+    let parsed: Record<string, unknown>;
+    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,
   };
 }
 

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<string, string>;
+
   /**
    * 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),

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

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<string, unknown>;
+    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<string, unknown>;
+    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", {