Show how to do per endpoint encryption

Author: Quinn-With-Two-Ns

temporal-sdk/src/main/java/io/temporal/internal/nexus/NexusInternal.java

@@ -8,4 +8,8 @@ private NexusInternal() {}
   public static NexusOperationContext getOperationContext() {
     return CurrentNexusOperationContext.get().getUserFacingContext();
   }
+
+  public static boolean isInOperationHandler() {
+    return CurrentNexusOperationContext.isNexusContext();
+  }
 }

temporal-sdk/src/main/java/io/temporal/nexus/Nexus.java

@@ -13,6 +13,15 @@ public static NexusOperationContext getOperationContext() {
     return NexusInternal.getOperationContext();
   }
 
+  /**
+   * Returns true if the current thread is executing inside a Nexus operation handler. Useful for
+   * context-aware components (such as codecs or interceptors) that need to behave differently
+   * inside vs outside a Nexus handler.
+   */
+  public static boolean isInOperationHandler() {
+    return NexusInternal.isInOperationHandler();
+  }
+
   /**
    * Use this to rethrow a checked exception from a Nexus Operation instead of adding the exception
    * to a method signature.

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

@@ -0,0 +1,105 @@
+package io.temporal.workflow.nexus;
+
+import io.temporal.api.common.v1.Payload;
+import io.temporal.common.context.ContextPropagator;
+import io.temporal.common.converter.DefaultDataConverter;
+import io.temporal.nexus.Nexus;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.Map;
+
+/**
+ * A {@link ContextPropagator} that carries the encryption key identifier in workflow headers.
+ *
+ * <p>This propagator bridges the async Nexus operation path: when a Nexus handler starts a
+ * workflow, this propagator auto-detects the endpoint from the Nexus operation context via {@link
+ * #getCurrentContext()} (called on the handler thread), serializes it into the workflow header, and
+ * restores it via {@link #setCurrentContext(Object)} (called on the workflow execution thread). The
+ * codec then reads the thread-local to select the correct encryption key.
+ *
+ * <p>Neither the handler nor the workflow code needs any awareness of encryption. The propagator
+ * automatically discovers the key ID from the Nexus context or the existing thread-local.
+ */
+public class EncryptionKeyContextPropagator implements ContextPropagator {
+
+  static final String HEADER_KEY = "x-encryption-key-id";
+
+  @Override
+  public String getName() {
+    return "encryption-key-propagator";
+  }
+
+  /**
+   * Called on the caller/handler thread to capture context for propagation. Auto-detects the key ID
+   * from: (1) Nexus operation context (on handler threads), (2) existing thread-local (on workflow
+   * threads).
+   */
+  @Override
+  public Object getCurrentContext() {
+    String keyId = null;
+
+    // Try Nexus context first (available on handler threads)
+    try {
+      keyId = Nexus.getOperationContext().getInfo().getEndpoint();
+    } catch (Exception e) {
+      // Not on a Nexus handler thread
+    }
+
+    // Fall back to thread-local (set by setCurrentContext on workflow threads)
+    if (keyId == null) {
+      keyId = PerEndpointEncryptionCodec.getCurrentKeyId();
+    }
+
+    if (keyId == null) {
+      return Collections.emptyMap();
+    }
+    Map<String, String> context = new HashMap<>();
+    context.put(HEADER_KEY, keyId);
+    return context;
+  }
+
+  /**
+   * Called on the workflow execution thread before workflow code runs. Sets the encryption key ID
+   * on the thread-local so the codec can read it during {@code encode()}.
+   */
+  @Override
+  @SuppressWarnings("unchecked")
+  public void setCurrentContext(Object context) {
+    if (context instanceof Map) {
+      Map<String, String> contextMap = (Map<String, String>) context;
+      String keyId = contextMap.get(HEADER_KEY);
+      if (keyId != null) {
+        PerEndpointEncryptionCodec.setCurrentKeyId(keyId);
+      }
+    }
+  }
+
+  @Override
+  public Map<String, Payload> serializeContext(Object context) {
+    if (context instanceof Map) {
+      @SuppressWarnings("unchecked")
+      Map<String, String> contextMap = (Map<String, String>) context;
+      String keyId = contextMap.get(HEADER_KEY);
+      if (keyId != null) {
+        Map<String, Payload> serialized = new HashMap<>();
+        serialized.put(
+            HEADER_KEY, DefaultDataConverter.newDefaultInstance().toPayload(keyId).get());
+        return serialized;
+      }
+    }
+    return Collections.emptyMap();
+  }
+
+  @Override
+  public Object deserializeContext(Map<String, Payload> header) {
+    Map<String, String> context = new HashMap<>();
+    Payload keyIdPayload = header.get(HEADER_KEY);
+    if (keyIdPayload != null) {
+      String keyId =
+          DefaultDataConverter.newDefaultInstance()
+              .fromPayload(keyIdPayload, String.class, String.class);
+      context.put(HEADER_KEY, keyId);
+    }
+    return context;
+  }
+}

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

@@ -0,0 +1,144 @@
+package io.temporal.workflow.nexus;
+
+import com.google.protobuf.ByteString;
+import io.temporal.api.common.v1.Payload;
+import io.temporal.nexus.Nexus;
+import io.temporal.payload.codec.PayloadCodec;
+import java.util.List;
+import java.util.stream.Collectors;
+import javax.annotation.Nonnull;
+
+/**
+ * A simulated per-endpoint encryption codec that demonstrates how to select different encryption
+ * keys based on the Nexus endpoint being called.
+ *
+ * <p>This codec uses a simple prefix-based transformation (not real encryption) to keep focus on
+ * the Nexus wiring pattern rather than cryptographic details. The design is asymmetric:
+ *
+ * <ul>
+ *   <li>{@code encode()}: Auto-detects the key ID by checking (1) the Nexus operation context if on
+ *       a handler thread, (2) the thread-local set by {@link EncryptionKeyContextPropagator} if on
+ *       a workflow thread, or (3) a default key. This means handler and workflow code need zero
+ *       awareness of encryption.
+ *   <li>{@code decode()}: Reads the key ID from the payload's own metadata ({@code
+ *       encryption-key-id} field). This is self-describing and does not depend on any external
+ *       context, which is necessary because handler-side input deserialization occurs before
+ *       handler code runs.
+ * </ul>
+ */
+public class PerEndpointEncryptionCodec implements PayloadCodec {
+
+  static final String METADATA_ENCRYPTED = "encrypted";
+  static final ByteString METADATA_ENCRYPTED_VALUE = ByteString.copyFromUtf8("true");
+  static final String METADATA_KEY_ID = "encryption-key-id";
+  static final String ENC_PREFIX = "ENC:";
+
+  private static final ThreadLocal<String> CURRENT_KEY_ID = new ThreadLocal<>();
+
+  private final String defaultKeyId;
+
+  /**
+   * @param defaultKeyId the default key ID used when neither Nexus context nor thread-local is
+   *     available (e.g., on the test thread or internal SDK serialization paths)
+   */
+  public PerEndpointEncryptionCodec(String defaultKeyId) {
+    this.defaultKeyId = defaultKeyId;
+  }
+
+  /** Sets the current encryption key ID on the calling thread (used by ContextPropagator). */
+  public static void setCurrentKeyId(String keyId) {
+    CURRENT_KEY_ID.set(keyId);
+  }
+
+  /** Clears the current encryption key ID on the calling thread. */
+  public static void clearCurrentKeyId() {
+    CURRENT_KEY_ID.remove();
+  }
+
+  /** Returns the current key ID for the calling thread, or null if not set. */
+  static String getCurrentKeyId() {
+    return CURRENT_KEY_ID.get();
+  }
+
+  @Nonnull
+  @Override
+  public List<Payload> encode(@Nonnull List<Payload> payloads) {
+    return payloads.stream().map(this::encodePayload).collect(Collectors.toList());
+  }
+
+  @Nonnull
+  @Override
+  public List<Payload> decode(@Nonnull List<Payload> payloads) {
+    return payloads.stream().map(this::decodePayload).collect(Collectors.toList());
+  }
+
+  /**
+   * Resolves the key ID to use for encoding. Checks sources in priority order:
+   *
+   * <ol>
+   *   <li>Nexus operation context (available on handler threads during serialization)
+   *   <li>Thread-local (set by {@link EncryptionKeyContextPropagator} on workflow threads)
+   *   <li>Default key (fallback for test threads and internal SDK operations)
+   * </ol>
+   */
+  private String resolveKeyId() {
+    // 1. Check if we're on a Nexus handler thread — the endpoint is the key ID
+    if (Nexus.isInOperationHandler()) {
+      return Nexus.getOperationContext().getInfo().getEndpoint();
+    }
+
+    // 2. Check thread-local (set by ContextPropagator for workflow threads)
+    String threadLocalKey = CURRENT_KEY_ID.get();
+    if (threadLocalKey != null) {
+      return threadLocalKey;
+    }
+
+    // 3. Default key
+    return defaultKeyId;
+  }
+
+  private Payload encodePayload(Payload payload) {
+    String keyId = resolveKeyId();
+
+    // Simulated encryption: prefix data with "ENC:<keyId>:" to make it visibly non-plaintext.
+    // Use toBuilder() to preserve existing metadata (e.g., "encoding: json/plain").
+    ByteString prefix = ByteString.copyFromUtf8(ENC_PREFIX + keyId + ":");
+    ByteString encodedData = prefix.concat(payload.getData());
+
+    return payload.toBuilder()
+        .putMetadata(METADATA_ENCRYPTED, METADATA_ENCRYPTED_VALUE)
+        .putMetadata(METADATA_KEY_ID, ByteString.copyFromUtf8(keyId))
+        .setData(encodedData)
+        .build();
+  }
+
+  private Payload decodePayload(Payload payload) {
+    // Self-describing: check our encryption marker, then read key-id from metadata.
+    ByteString encryptedMeta = payload.getMetadataOrDefault(METADATA_ENCRYPTED, null);
+    if (encryptedMeta == null || !encryptedMeta.equals(METADATA_ENCRYPTED_VALUE)) {
+      return payload;
+    }
+
+    ByteString keyIdBytes = payload.getMetadataOrDefault(METADATA_KEY_ID, null);
+    if (keyIdBytes == null) {
+      throw new IllegalStateException("Encrypted payload missing encryption-key-id metadata");
+    }
+    String keyId = keyIdBytes.toStringUtf8();
+
+    // Reverse the simulated encryption: strip the "ENC:<keyId>:" prefix
+    String expectedPrefix = ENC_PREFIX + keyId + ":";
+    ByteString expectedPrefixBytes = ByteString.copyFromUtf8(expectedPrefix);
+    ByteString data = payload.getData();
+    if (!data.startsWith(expectedPrefixBytes)) {
+      throw new IllegalStateException(
+          "Encrypted payload data does not start with expected prefix: " + expectedPrefix);
+    }
+    ByteString decodedData = data.substring(expectedPrefixBytes.size());
+
+    return payload.toBuilder()
+        .removeMetadata(METADATA_ENCRYPTED)
+        .removeMetadata(METADATA_KEY_ID)
+        .setData(decodedData)
+        .build();
+  }
+}