Implement request identity (#297)

* Add RequestIdentityVerifier interface, to be used to implement request identity.

* Add RestateRequestIdentityVerifier implementation.

* Fix case with unsigned signature scheme.

* Fix description of module

* Better name for the factory method

* Add test and remove prefixing with ASN1, this seems not needed.

Author: slinkydeveloper

sdk-common/src/main/java/dev/restate/sdk/auth/RequestIdentityVerifier.java

@@ -0,0 +1,26 @@
+// Copyright (c) 2023 - Restate Software, Inc., Restate GmbH
+//
+// This file is part of the Restate Java SDK,
+// which is released under the MIT license.
+//
+// You can find a copy of the license in file LICENSE in the root
+// directory of this repository or package, or at
+// https://github.com/restatedev/sdk-java/blob/main/LICENSE
+package dev.restate.sdk.auth;
+
+import org.jspecify.annotations.Nullable;
+
+/** Interface to verify requests. */
+public interface RequestIdentityVerifier {
+
+  /** Abstraction for headers map. */
+  @FunctionalInterface
+  interface Headers {
+    @Nullable String get(String key);
+  }
+
+  /**
+   * @throws Exception if the request cannot be verified
+   */
+  void verifyRequest(Headers headers) throws Exception;
+}

sdk-core/src/main/java/dev/restate/sdk/core/Entries.java

@@ -217,8 +217,8 @@ public Result<Collection<String>> parseCompletionResult(CompletionMessage actual
         } catch (InvalidProtocolBufferException e) {
           throw new ProtocolException(
               "Cannot parse get state keys completion",
-              e,
-              ProtocolException.PROTOCOL_VIOLATION_CODE);
+              ProtocolException.PROTOCOL_VIOLATION_CODE,
+              e);
         }
         return Result.success(
             stateKeys.getKeysList().stream()

sdk-core/src/main/java/dev/restate/sdk/core/InvocationStateMachine.java

@@ -209,8 +209,8 @@ void onStartMessage(MessageLite msg) {
       this.fail(
           new ProtocolException(
               "Expected at least one entry with Input, got " + this.entriesToReplay + " entries",
-              null,
-              TerminalException.INTERNAL_SERVER_ERROR_CODE));
+              TerminalException.INTERNAL_SERVER_ERROR_CODE,
+              null));
       return;
     }
 

sdk-core/src/main/java/dev/restate/sdk/core/ProtocolException.java

@@ -14,6 +14,7 @@
 
 public class ProtocolException extends RuntimeException {
 
+  static final int UNAUTHORIZED_CODE = 401;
   static final int NOT_FOUND_CODE = 404;
   static final int JOURNAL_MISMATCH_CODE = 570;
   static final int PROTOCOL_VIOLATION_CODE = 571;
@@ -28,10 +29,10 @@ private ProtocolException(String message) {
   }
 
   private ProtocolException(String message, int code) {
-    this(message, null, code);
+    this(message, code, null);
   }
 
-  public ProtocolException(String message, Throwable cause, int code) {
+  public ProtocolException(String message, int code, Throwable cause) {
     super(message, cause);
     this.code = code;
   }
@@ -77,7 +78,11 @@ static ProtocolException methodNotFound(String serviceName, String handlerName)
   static ProtocolException invalidSideEffectCall() {
     return new ProtocolException(
         "A syscall was invoked from within a side effect closure.",
-        null,
-        TerminalException.INTERNAL_SERVER_ERROR_CODE);
+        TerminalException.INTERNAL_SERVER_ERROR_CODE,
+        null);
+  }
+
+  static ProtocolException unauthorized(Throwable e) {
+    return new ProtocolException("Unauthorized", UNAUTHORIZED_CODE, e);
   }
 }

sdk-core/src/main/java/dev/restate/sdk/core/RestateEndpoint.java

@@ -8,6 +8,7 @@
 // https://github.com/restatedev/sdk-java/blob/main/LICENSE
 package dev.restate.sdk.core;
 
+import dev.restate.sdk.auth.RequestIdentityVerifier;
 import dev.restate.sdk.common.BindableServiceFactory;
 import dev.restate.sdk.common.syscalls.HandlerDefinition;
 import dev.restate.sdk.common.syscalls.ServiceDefinition;
@@ -32,14 +33,17 @@ public class RestateEndpoint {
 
   private final Map<String, ServiceAndOptions<?>> services;
   private final Tracer tracer;
+  private final RequestIdentityVerifier requestIdentityVerifier;
   private final DeploymentManifest deploymentManifest;
 
   private RestateEndpoint(
       DeploymentManifestSchema.ProtocolMode protocolMode,
       Map<String, ServiceAndOptions<?>> services,
-      Tracer tracer) {
+      Tracer tracer,
+      RequestIdentityVerifier requestIdentityVerifier) {
     this.services = services;
     this.tracer = tracer;
+    this.requestIdentityVerifier = requestIdentityVerifier;
     this.deploymentManifest =
         new DeploymentManifest(protocolMode, services.values().stream().map(c -> c.service));
 
@@ -49,6 +53,7 @@ private RestateEndpoint(
   public ResolvedEndpointHandler resolve(
       String componentName,
       String handlerName,
+      RequestIdentityVerifier.Headers headers,
       io.opentelemetry.context.Context otelContext,
       LoggingContextSetter loggingContextSetter,
       @Nullable Executor syscallExecutor)
@@ -65,6 +70,15 @@ public ResolvedEndpointHandler resolve(
       throw ProtocolException.methodNotFound(componentName, handlerName);
     }
 
+    // Verify request
+    if (requestIdentityVerifier != null) {
+      try {
+        requestIdentityVerifier.verifyRequest(headers);
+      } catch (Exception e) {
+        throw ProtocolException.unauthorized(e);
+      }
+    }
+
     // Generate the span
     Span span =
         tracer
@@ -108,6 +122,7 @@ public static class Builder {
 
     private final List<ServiceAndOptions<?>> services = new ArrayList<>();
     private final DeploymentManifestSchema.ProtocolMode protocolMode;
+    private RequestIdentityVerifier requestIdentityVerifier;
     private Tracer tracer = OpenTelemetry.noop().getTracer("NOOP");
 
     public Builder(DeploymentManifestSchema.ProtocolMode protocolMode) {
@@ -124,12 +139,18 @@ public Builder withTracer(Tracer tracer) {
       return this;
     }
 
+    public Builder withRequestIdentityVerifier(RequestIdentityVerifier requestIdentityVerifier) {
+      this.requestIdentityVerifier = requestIdentityVerifier;
+      return this;
+    }
+
     public RestateEndpoint build() {
       return new RestateEndpoint(
           this.protocolMode,
           this.services.stream()
               .collect(Collectors.toMap(c -> c.service.getServiceName(), Function.identity())),
-          tracer);
+          tracer,
+          requestIdentityVerifier);
     }
   }
 

sdk-core/src/test/java/dev/restate/sdk/core/MockMultiThreaded.java

@@ -56,6 +56,7 @@ public void executeTest(TestDefinitions.TestDefinition definition) {
         server.resolve(
             serviceDefinition.get(0).getServiceName(),
             definition.getMethod(),
+            k -> null,
             io.opentelemetry.context.Context.current(),
             RestateEndpoint.LoggingContextSetter.THREAD_LOCAL_INSTANCE,
             syscallsExecutor);

sdk-core/src/test/java/dev/restate/sdk/core/MockSingleThread.java

@@ -54,6 +54,7 @@ public void executeTest(TestDefinition definition) {
         server.resolve(
             serviceDefinition.get(0).getServiceName(),
             definition.getMethod(),
+            k -> null,
             io.opentelemetry.context.Context.current(),
             RestateEndpoint.LoggingContextSetter.THREAD_LOCAL_INSTANCE,
             null);

sdk-http-vertx/src/main/java/dev/restate/sdk/http/vertx/RequestHttpServerHandler.java

@@ -116,6 +116,7 @@ public void handle(HttpServerRequest request) {
           restateEndpoint.resolve(
               serviceName,
               handlerName,
+              request::getHeader,
               otelContext,
               ContextualData::put,
               currentContextExecutor(vertxCurrentContext));

sdk-http-vertx/src/main/java/dev/restate/sdk/http/vertx/RestateHttpEndpointBuilder.java

@@ -8,6 +8,7 @@
 // https://github.com/restatedev/sdk-java/blob/main/LICENSE
 package dev.restate.sdk.http.vertx;
 
+import dev.restate.sdk.auth.RequestIdentityVerifier;
 import dev.restate.sdk.common.BindableService;
 import dev.restate.sdk.common.syscalls.ServiceDefinition;
 import dev.restate.sdk.core.RestateEndpoint;
@@ -103,7 +104,7 @@ public <O> RestateHttpEndpointBuilder bind(BindableService<O> service, O options
   }
 
   /**
-   * Add a {@link OpenTelemetry} implementation for tracing and metrics.
+   * Set the {@link OpenTelemetry} implementation for tracing and metrics.
    *
    * @see OpenTelemetry
    */
@@ -112,6 +113,18 @@ public RestateHttpEndpointBuilder withOpenTelemetry(OpenTelemetry openTelemetry)
     return this;
   }
 
+  /**
+   * Set the request identity verifier for this endpoint.
+   *
+   * <p>For the Restate implementation to use with Restate Cloud, check the module {@code
+   * sdk-request-identity}.
+   */
+  public RestateHttpEndpointBuilder withRequestIdentityVerifier(
+      RequestIdentityVerifier requestIdentityVerifier) {
+    this.endpointBuilder.withRequestIdentityVerifier(requestIdentityVerifier);
+    return this;
+  }
+
   /** Build and listen on the specified port. */
   public void buildAndListen(int port) {
     build().listen(port).onComplete(RestateHttpEndpointBuilder::handleStart);

sdk-lambda/src/main/java/dev/restate/sdk/lambda/RestateLambdaEndpoint.java

@@ -122,6 +122,7 @@ private APIGatewayProxyResponseEvent handleInvoke(APIGatewayProxyRequestEvent in
           this.restateEndpoint.resolve(
               serviceName,
               handlerName,
+              input.getHeaders()::get,
               otelContext,
               RestateEndpoint.LoggingContextSetter.THREAD_LOCAL_INSTANCE,
               null);