Make sure Nexus operation input is also encrypted

Author: Quinn-With-Two-Ns

temporal-sdk/src/test/java/io/temporal/workflow/nexus/EncryptionKeyContextPropagator.java

@@ -0,0 +1,65 @@
+package io.temporal.workflow.nexus;
+
+import io.temporal.api.common.v1.WorkflowExecution;
+import io.temporal.client.ActivityCompletionClient;
+import io.temporal.client.WorkflowOptions;
+import io.temporal.client.WorkflowStub;
+import io.temporal.common.converter.DefaultDataConverter;
+import io.temporal.common.interceptors.WorkflowClientCallsInterceptor;
+import io.temporal.common.interceptors.WorkflowClientCallsInterceptorBase;
+import io.temporal.common.interceptors.WorkflowClientInterceptor;
+import io.temporal.nexus.Nexus;
+import java.util.Optional;
+
+/**
+ * Client interceptor for per-endpoint Nexus encryption.
+ *
+ * <p>When a workflow is started from a Nexus operation handler, injects the endpoint name into the
+ * workflow's header. The {@link PerEndpointEncryptionWorkerInterceptor} reads it on the workflow
+ * thread and sets the codec's thread-local key, ensuring the async workflow result is encrypted
+ * with the correct per-endpoint key.
+ */
+public class PerEndpointEncryptionClientInterceptor implements WorkflowClientInterceptor {
+
+  static final String ENDPOINT_HEADER_KEY = "x-encryption-endpoint";
+
+  @Override
+  public WorkflowClientCallsInterceptor workflowClientCallsInterceptor(
+      WorkflowClientCallsInterceptor next) {
+    return new WorkflowClientCallsInterceptorBase(next) {
+      @Override
+      public WorkflowStartOutput start(WorkflowStartInput input) {
+        // If we're on a Nexus handler thread, inject the endpoint into the workflow header.
+        if (Nexus.isInOperationHandler()) {
+          String endpoint = Nexus.getOperationContext().getInfo().getEndpoint();
+          input
+              .getHeader()
+              .getValues()
+              .put(
+                  ENDPOINT_HEADER_KEY,
+                  DefaultDataConverter.newDefaultInstance().toPayload(endpoint).get());
+        }
+        return super.start(input);
+      }
+    };
+  }
+
+  @Override
+  @Deprecated
+  public WorkflowStub newUntypedWorkflowStub(
+      String workflowType, WorkflowOptions options, WorkflowStub next) {
+    return next;
+  }
+
+  @Override
+  @Deprecated
+  public WorkflowStub newUntypedWorkflowStub(
+      WorkflowExecution execution, Optional<String> workflowType, WorkflowStub next) {
+    return next;
+  }
+
+  @Override
+  public ActivityCompletionClient newActivityCompletionClient(ActivityCompletionClient next) {
+    return next;
+  }
+}

temporal-sdk/src/test/java/io/temporal/workflow/nexus/PerEndpointEncryptionClientInterceptor.java

@@ -20,6 +20,7 @@
 import io.temporal.nexus.Nexus;
 import io.temporal.nexus.WorkflowRunOperation;
 import io.temporal.testing.internal.SDKTestWorkflowRule;
+import io.temporal.worker.WorkerFactoryOptions;
 import io.temporal.workflow.*;
 import java.time.Duration;
 import java.util.Collections;
@@ -37,8 +38,9 @@
  *
  * <ul>
  *   <li>The {@link PerEndpointEncryptionCodec} auto-detects the endpoint from Nexus context on
- *       handler threads, or from the thread-local set by the {@link EncryptionKeyContextPropagator}
- *       on workflow threads. Falls back to a default key for other contexts.
+ *       handler threads, or from the thread-local set by the {@link
+ *       PerEndpointEncryptionInterceptor} on workflow threads. Falls back to a default key for
+ *       other contexts.
  *   <li>{@code decode()} reads the key ID from payload metadata (self-describing), so handler-side
  *       input deserialization works before handler code runs.
  *   <li>Nexus handler code is pure business logic — no encryption concerns.
@@ -60,8 +62,11 @@ public class PerEndpointEncryptionTest extends BaseNexusTest {
                           DefaultDataConverter.newDefaultInstance(),
                           Collections.singletonList(new PerEndpointEncryptionCodec(DEFAULT_KEY_ID)),
                           true))
-                  .setContextPropagators(
-                      Collections.singletonList(new EncryptionKeyContextPropagator()))
+                  .setInterceptors(new PerEndpointEncryptionClientInterceptor())
+                  .build())
+          .setWorkerFactoryOptions(
+              WorkerFactoryOptions.newBuilder()
+                  .setWorkerInterceptors(new PerEndpointEncryptionWorkerInterceptor())
                   .build())
           .build();
 
@@ -202,6 +207,15 @@ public void testPayloadsEncryptedWithCorrectKeys() {
       if (event.hasNexusOperationScheduledEventAttributes()) {
         Payload nexusInput = event.getNexusOperationScheduledEventAttributes().getInput();
         assertPayloadIsEncrypted(nexusInput, "NexusOperationScheduled input");
+        // Verify the Nexus operation input uses the endpoint key, not the default key.
+        String nexusInputKeyId =
+            nexusInput
+                .getMetadataOrThrow(PerEndpointEncryptionCodec.METADATA_KEY_ID)
+                .toStringUtf8();
+        assertEquals(
+            "Nexus operation input must use endpoint key, not default",
+            expectedEndpointKey,
+            nexusInputKeyId);
         foundEncryptedPayload = true;
       }
       if (event.hasWorkflowExecutionCompletedEventAttributes()) {
@@ -236,7 +250,7 @@ public void testPayloadsEncryptedWithCorrectKeys() {
         assertPayloadIsEncrypted(resultPayload, "AsyncWorkflow result");
 
         // This is the critical assertion: the async workflow's result must be encrypted
-        // with the endpoint-specific key, proving the ContextPropagator carried it through.
+        // with the endpoint-specific key, proving the interceptor carried it through.
         String keyId =
             resultPayload
                 .getMetadataOrThrow(PerEndpointEncryptionCodec.METADATA_KEY_ID)

temporal-sdk/src/test/java/io/temporal/workflow/nexus/PerEndpointEncryptionTest.java

@@ -0,0 +1,60 @@
+package io.temporal.workflow.nexus;
+
+import io.temporal.api.common.v1.Payload;
+import io.temporal.common.converter.DefaultDataConverter;
+import io.temporal.common.interceptors.*;
+
+/**
+ * Worker interceptor for per-endpoint Nexus encryption.
+ *
+ * <p>Two interception points:
+ *
+ * <ul>
+ *   <li><b>Caller workflow (outbound)</b>: Before a Nexus operation executes, sets the codec's
+ *       thread-local key from the endpoint name. This ensures the codec encrypts the Nexus
+ *       operation input with the correct per-endpoint key.
+ *   <li><b>Async workflow (inbound)</b>: When a workflow started by a Nexus handler begins, reads
+ *       the endpoint from a header (injected by {@link PerEndpointEncryptionClientInterceptor}) and
+ *       sets the thread-local. This ensures the codec encrypts the workflow result with the
+ *       endpoint key.
+ * </ul>
+ */
+public class PerEndpointEncryptionWorkerInterceptor extends WorkerInterceptorBase {
+
+  @Override
+  public WorkflowInboundCallsInterceptor interceptWorkflow(WorkflowInboundCallsInterceptor next) {
+    return new WorkflowInboundCallsInterceptorBase(next) {
+      @Override
+      public void init(WorkflowOutboundCallsInterceptor outboundCalls) {
+        next.init(
+            new WorkflowOutboundCallsInterceptorBase(outboundCalls) {
+              @Override
+              public <R> ExecuteNexusOperationOutput<R> executeNexusOperation(
+                  ExecuteNexusOperationInput<R> input) {
+                // Set the key from the endpoint BEFORE the input is serialized.
+                PerEndpointEncryptionCodec.setCurrentKeyId(input.getEndpoint());
+                return super.executeNexusOperation(input);
+              }
+            });
+      }
+
+      @Override
+      public WorkflowOutput execute(WorkflowInput input) {
+        // If this workflow was started by a Nexus handler, the client interceptor injected
+        // the endpoint into the header. Read it and set the thread-local.
+        Payload endpointPayload =
+            input
+                .getHeader()
+                .getValues()
+                .get(PerEndpointEncryptionClientInterceptor.ENDPOINT_HEADER_KEY);
+        if (endpointPayload != null) {
+          String endpoint =
+              DefaultDataConverter.newDefaultInstance()
+                  .fromPayload(endpointPayload, String.class, String.class);
+          PerEndpointEncryptionCodec.setCurrentKeyId(endpoint);
+        }
+        return super.execute(input);
+      }
+    };
+  }
+}