feat: JWT annotations (#1141)

Author: franciscolopezsancho

sdk/spring-sdk/src/main/java/kalix/springsdk/annotations/JWT.java

@@ -0,0 +1,50 @@
+/*
+ * Copyright 2021 Lightbend Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package kalix.springsdk.annotations;
+
+import java.lang.annotation.*;
+
+@Target(ElementType.METHOD)
+@Retention(RetentionPolicy.RUNTIME)
+@Documented
+public @interface JWT {
+
+  enum JwtMethodMode {
+    // No validation.
+    UNSPECIFIED,
+    // Validate the bearer token.
+    BEARER_TOKEN,
+    // Validate/sign a token field in the message against the message fields.
+    //
+    // If present, the message must have a token annotated field or the message itself must have
+    // validate_bearer_token
+    // set to true.
+    MESSAGE
+  }
+
+  JwtMethodMode[] validate();
+
+  JwtMethodMode[] sign();
+  // If set, then the token extracted from the bearer token must have this issuer.
+  //
+  // This can be used in combination with the issuer field of configuration for JWT secrets, if
+  // there is at least one
+  // secret that has this issuer set, then only those secrets with that issuer set will be used
+  // for validating or
+  // signing this token, so you can be sure that the token did come from a particular issuer.
+  String[] bearerTokenIssuer();
+}

sdk/spring-sdk/src/main/scala/kalix/springsdk/impl/ActionDescriptorFactory.scala

@@ -21,9 +21,11 @@ import kalix.springsdk.impl.ComponentDescriptorFactory.eventingInForTopic
 import kalix.springsdk.impl.ComponentDescriptorFactory.eventingInForValueEntity
 import kalix.springsdk.impl.ComponentDescriptorFactory.eventingOutForTopic
 import kalix.springsdk.impl.ComponentDescriptorFactory.hasEventSourcedEntitySubscription
+import kalix.springsdk.impl.ComponentDescriptorFactory.hasJwtMethodOptions
 import kalix.springsdk.impl.ComponentDescriptorFactory.hasTopicPublication
 import kalix.springsdk.impl.ComponentDescriptorFactory.hasTopicSubscription
 import kalix.springsdk.impl.ComponentDescriptorFactory.hasValueEntitySubscription
+import kalix.springsdk.impl.ComponentDescriptorFactory.jwtMethodOptions
 import kalix.springsdk.impl.ComponentDescriptorFactory.validateRestMethod
 import kalix.springsdk.impl.reflection.CombinedSubscriptionServiceMethod
 import kalix.springsdk.impl.reflection.KalixMethod
@@ -39,7 +41,7 @@ private[impl] object ActionDescriptorFactory extends ComponentDescriptorFactory
     val springAnnotatedMethods =
       RestServiceIntrospector.inspectService(component).methods.map { serviceMethod =>
         validateRestMethod(serviceMethod.javaMethod)
-        KalixMethod(serviceMethod)
+        KalixMethod(serviceMethod).withKalixOptions(buildJWTOptions(serviceMethod.javaMethod))
       }
 
     //TODO make sure no subscription should be exposed via REST.
@@ -121,19 +123,19 @@ private[impl] object ActionDescriptorFactory extends ComponentDescriptorFactory
     val serviceName = nameGenerator.getName(component.getSimpleName)
 
     def filterAndAddKalixOptions(to: Seq[KalixMethod], from: Seq[KalixMethod]): Seq[KalixMethod] = {
-      val common = to.flatMap(toAdd =>
+      val inCommon = to.flatMap(toAdd =>
         from
           .filter { addingFrom =>
             addingFrom.serviceMethod.methodName.equals(toAdd.serviceMethod.methodName)
           }
           .map(addingFrom => toAdd.withKalixOptions(addingFrom.methodOptions)))
-      val notInCommon = to
+      val unique = to
         .filter { toAdd =>
           !from.exists { addingFrom =>
             addingFrom.serviceMethod.methodName.equals(toAdd.serviceMethod.methodName)
           }
         }
-      common ++ notInCommon
+      inCommon ++ unique
     }
 
     def removeDuplicates(springMethods: Seq[KalixMethod], pubSubMethods: Seq[KalixMethod]): Seq[KalixMethod] = {

sdk/spring-sdk/src/main/scala/kalix/springsdk/impl/ComponentDescriptor.scala

@@ -103,10 +103,8 @@ private[impl] object ComponentDescriptor {
           kalixMethod.serviceMethod.streamIn,
           kalixMethod.serviceMethod.streamOut)
 
-      val methodOptions =
-        kalixMethod.methodOptions.foldLeft(MethodOptions.newBuilder()) { (optionsBuilder, options) =>
-          optionsBuilder.setExtension(kalix.Annotations.method, options)
-        }
+      val methodOptions = MethodOptions.newBuilder()
+      kalixMethod.methodOptions.foreach(option => methodOptions.setExtension(kalix.Annotations.method, option))
 
       methodOptions.setExtension(AnnotationsProto.http, httpRuleBuilder.build())
 

sdk/spring-sdk/src/main/scala/kalix/springsdk/impl/ComponentDescriptorFactory.scala

@@ -16,10 +16,14 @@
 
 package kalix.springsdk.impl
 
+import kalix.MethodOptions
+import kalix.springsdk.annotations.JWT
 import kalix.springsdk.annotations.Table
 import kalix.springsdk.annotations.{ Entity, Publish, Subscribe }
+import kalix.springsdk.impl.ComponentDescriptorFactory.hasJwtMethodOptions
+import kalix.springsdk.impl.ComponentDescriptorFactory.jwtMethodOptions
 import kalix.springsdk.impl.reflection._
-import kalix.{ EventDestination, EventSource, Eventing }
+import kalix.{ EventDestination, EventSource, Eventing, JwtMethodOptions }
 
 import java.lang.reflect.{ Method, Modifier }
 
@@ -41,6 +45,10 @@ private[impl] object ComponentDescriptorFactory {
     Modifier.isPublic(javaMethod.getModifiers) &&
     javaMethod.getAnnotation(classOf[Publish.Topic]) != null
 
+  def hasJwtMethodOptions(javaMehod: Method): Boolean =
+    Modifier.isPublic(javaMehod.getModifiers) &&
+    javaMehod.getAnnotation(classOf[JWT]) != null
+
   def findEventSourcedEntityType(javaMethod: Method): String = {
     val ann = javaMethod.getAnnotation(classOf[Subscribe.EventSourcedEntity])
     val entityClass = ann.value()
@@ -74,6 +82,19 @@ private[impl] object ComponentDescriptorFactory {
     ann.value()
   }
 
+  def jwtMethodOptions(javaMethod: Method): JwtMethodOptions = {
+    val ann = javaMethod.getAnnotation(classOf[JWT])
+    val jwt = JwtMethodOptions.newBuilder()
+    ann
+      .validate()
+      .map(springValidate => jwt.addValidate(JwtMethodOptions.JwtMethodMode.forNumber(springValidate.ordinal())))
+    ann
+      .sign()
+      .map(springSign => jwt.addSign(JwtMethodOptions.JwtMethodMode.forNumber(springSign.ordinal())))
+    ann.bearerTokenIssuer().map(jwt.addBearerTokenIssuer)
+    jwt.build()
+  }
+
   def eventingInForValueEntity(javaMethod: Method): Eventing = {
     val entityType = findValueEntityType(javaMethod)
     val eventSource = EventSource.newBuilder().setValueEntity(entityType).build()
@@ -140,6 +161,7 @@ private[impl] trait ComponentDescriptorFactory {
       //Assuming there is only one eventing.in annotation per method, therefore head is as good as any other
       withEventSourcedIn.groupBy(m => m.methodOptions.head.getEventing.getIn.getEventSourcedEntity)
     }
+
     groupByES(subscriptions).collect {
       case (eventSourcedEntity, kMethods) if kMethods.size > 1 =>
         val typeUrl2Method: Seq[TypeUrl2Method] = kMethods.map { k =>
@@ -161,6 +183,12 @@ private[impl] trait ComponentDescriptorFactory {
         kMethod
     }.toSeq
   }
+
+  private[impl] def buildJWTOptions(method: Method): Option[MethodOptions] = {
+    Option.when(hasJwtMethodOptions(method)) {
+      kalix.MethodOptions.newBuilder().setJwt(jwtMethodOptions(method)).build()
+    }
+  }
 }
 
 /**

sdk/spring-sdk/src/main/scala/kalix/springsdk/impl/EntityDescriptorFactory.scala

@@ -28,7 +28,7 @@ private[impl] object EntityDescriptorFactory extends ComponentDescriptorFactory
 
     val kalixMethods =
       RestServiceIntrospector.inspectService(component).methods.map { restMethod =>
-        KalixMethod(restMethod, entityKeys = entityKeys)
+        KalixMethod(restMethod, entityKeys = entityKeys).withKalixOptions(buildJWTOptions(restMethod.javaMethod))
       }
 
     val serviceName = nameGenerator.getName(component.getSimpleName)

sdk/spring-sdk/src/main/scala/kalix/springsdk/impl/ViewDescriptorFactory.scala

@@ -17,7 +17,7 @@
 package kalix.springsdk.impl
 
 import java.lang.reflect.Method
-
+import java.lang.reflect.ParameterizedType
 import kalix.MethodOptions
 import kalix.springsdk.annotations.Query
 import kalix.springsdk.annotations.Subscribe
@@ -31,8 +31,6 @@ import kalix.springsdk.impl.reflection.RestServiceIntrospector
 import kalix.springsdk.impl.reflection.SpringRestServiceMethod
 import kalix.springsdk.impl.reflection.ReflectionUtils
 import kalix.springsdk.impl.reflection.RestServiceIntrospector.BodyParameter
-import java.lang.reflect.ParameterizedType
-
 import kalix.springsdk.impl.ComponentDescriptorFactory.eventingInForEventSourcedEntity
 import kalix.springsdk.impl.ComponentDescriptorFactory.hasEventSourcedEntitySubscription
 import kalix.springsdk.impl.reflection.SubscriptionServiceMethod
@@ -147,7 +145,8 @@ private[impl] object ViewDescriptorFactory extends ComponentDescriptorFactory {
       // since it is a query, we don't actually ever want to handle any request in the SDK
       // the proxy does the work for us, mark the method as non-callable
       (
-        KalixMethod(annotatedMethod.copy(callable = false), methodOptions = Seq(methodOptions)),
+        KalixMethod(annotatedMethod.copy(callable = false), methodOptions = Some(methodOptions))
+          .withKalixOptions(buildJWTOptions(annotatedMethod.javaMethod)),
         queryInputSchemaDescriptor,
         queryOutputSchemaDescriptor)
     }
@@ -279,7 +278,7 @@ private[impl] object ViewDescriptorFactory extends ComponentDescriptorFactory {
 
         KalixMethod(SubscriptionServiceMethod(method, method.getParameterCount - 1))
           .withKalixOptions(methodOptionsBuilder.build())
-
+          .withKalixOptions(buildJWTOptions(method))
       }
       .toSeq
   }

sdk/spring-sdk/src/main/scala/kalix/springsdk/impl/reflection/KalixMethod.scala

@@ -208,14 +208,38 @@ case class SpringRestServiceMethod(
 
 case class KalixMethod(
     serviceMethod: ServiceMethod,
-    methodOptions: Seq[kalix.MethodOptions] = Seq.empty,
+    methodOptions: Option[kalix.MethodOptions] = None,
     entityKeys: Seq[String] = Seq.empty) {
 
+  /**
+   * This method merges the new method options with the existing ones. In case of collision the 'opts' are kept
+   * @param opts
+   * @return
+   */
   def withKalixOptions(opts: kalix.MethodOptions): KalixMethod =
-    copy(methodOptions = methodOptions :+ opts)
+    copy(methodOptions = Some(mergeKalixOptions(methodOptions, opts)))
+
+  /**
+   * This method merges the new method options with the existing ones. In case of collision the 'opts' are kept
+   * @param opts
+   * @return
+   */
+  def withKalixOptions(opts: Option[kalix.MethodOptions]): KalixMethod =
+    opts match {
+      case Some(methodOptions) => withKalixOptions(methodOptions)
+      case None                => this
+    }
 
-  def withKalixOptions(opts: Seq[kalix.MethodOptions]): KalixMethod =
-    copy(methodOptions = methodOptions ++ opts)
+  private[impl] def mergeKalixOptions(
+      source: Option[kalix.MethodOptions],
+      addOn: kalix.MethodOptions): kalix.MethodOptions = {
+    val builder = source match {
+      case Some(src) => src.toBuilder
+      case None      => kalix.MethodOptions.newBuilder()
+    }
+    builder.mergeFrom(addOn)
+    builder.build()
+  }
 }
 
 trait ExtractorCreator {

sdk/spring-sdk/src/test/java/kalix/springsdk/testmodels/action/ActionsTestModels.java

@@ -16,9 +16,9 @@
 
 package kalix.springsdk.testmodels.action;
 
-import akka.NotUsed;
-import akka.stream.javadsl.Source;
+import kalix.JwtMethodOptions.JwtMethodMode;
 import kalix.javasdk.action.Action;
+import kalix.springsdk.annotations.JWT;
 import kalix.springsdk.testmodels.Message;
 import org.springframework.web.bind.annotation.*;
 import reactor.core.publisher.Flux;
@@ -54,6 +54,17 @@ public Action.Effect<Message> message(@RequestBody Message msg) {
     }
   }
 
+  public static class PostWithoutParamWithJWT extends Action {
+    @PostMapping("/message")
+    @JWT(
+        validate = JWT.JwtMethodMode.BEARER_TOKEN,
+        sign = JWT.JwtMethodMode.MESSAGE,
+        bearerTokenIssuer = {"a", "b"})
+    public Action.Effect<Message> message(@RequestBody Message msg) {
+      return effects().reply(msg);
+    }
+  }
+
   public static class PostWithOneParam extends Action {
     @PostMapping("/message/{one}")
     public Action.Effect<Message> message(@PathVariable String one, @RequestBody Message msg) {

sdk/spring-sdk/src/test/java/kalix/springsdk/testmodels/eventsourcedentity/EventSourcedEntitiesTestModels.java

@@ -16,9 +16,13 @@
 
 package kalix.springsdk.testmodels.eventsourcedentity;
 
+import kalix.JwtMethodOptions;
 import kalix.javasdk.eventsourcedentity.EventSourcedEntity;
 import kalix.springsdk.annotations.Entity;
 import kalix.springsdk.annotations.EventHandler;
+import kalix.springsdk.annotations.JWT;
+import kalix.springsdk.annotations.Subscribe;
+import kalix.springsdk.testmodels.valueentity.UserEntity;
 import org.springframework.web.bind.annotation.*;
 
 public class EventSourcedEntitiesTestModels {
@@ -74,6 +78,29 @@ private Integer privateMethodSimilarSignature(Integer event) {
     }
   }
 
+  @Entity(entityKey = "id", entityType = "counter")
+  @RequestMapping("/eventsourced/{id}")
+  public static class WellAnnotatedESEntityWithJWT extends EventSourcedEntity<Integer> {
+
+    @GetMapping("/int/{number}")
+    @JWT(
+        validate = JWT.JwtMethodMode.BEARER_TOKEN,
+        sign = JWT.JwtMethodMode.MESSAGE,
+        bearerTokenIssuer = {"a", "b"})
+    public Integer getInteger(@PathVariable Integer number) {
+      return number;
+    }
+
+    @PostMapping("/changeInt/{number}")
+    @JWT(
+        validate = JWT.JwtMethodMode.BEARER_TOKEN,
+        sign = JWT.JwtMethodMode.MESSAGE,
+        bearerTokenIssuer = {"a", "b"})
+    public Integer changeInteger(@PathVariable Integer number) {
+      return number;
+    }
+  }
+
   @Entity(entityKey = "id", entityType = "counter")
   public static class ErrorDuplicatedEventsEntity extends EventSourcedEntity<Integer> {
 

sdk/spring-sdk/src/test/java/kalix/springsdk/testmodels/view/ViewTestModels.java

@@ -17,6 +17,7 @@
 package kalix.springsdk.testmodels.view;
 
 import kalix.javasdk.view.View;
+import kalix.springsdk.annotations.JWT;
 import kalix.springsdk.annotations.Query;
 import kalix.springsdk.annotations.Subscribe;
 import kalix.springsdk.annotations.Table;
@@ -92,6 +93,27 @@ public TransformedUser getUser(@RequestBody ByEmail byEmail) {
     }
   }
 
+  @Table("users_view")
+  public static class TransformedUserViewWithJWT extends View<TransformedUser> {
+
+    // when methods are annotated, it's implicitly a transform = true
+    @Subscribe.ValueEntity(UserEntity.class)
+    public UpdateEffect<TransformedUser> onChange(User user) {
+      return effects()
+          .updateState(new TransformedUser(user.lastName + ", " + user.firstName, user.email));
+    }
+
+    @Query("SELECT * FROM users_view WHERE email = :email")
+    @PostMapping("/users/by-email")
+    @JWT(
+        validate = JWT.JwtMethodMode.BEARER_TOKEN,
+        sign = JWT.JwtMethodMode.MESSAGE,
+        bearerTokenIssuer = {"a", "b"})
+    public TransformedUser getUser(@RequestBody ByEmail byEmail) {
+      return null;
+    }
+  }
+
   @Table("users_view")
   public static class TransformedUserViewUsingState extends View<TransformedUser> {