Metadata/Context Propagation

Server Context

First define the context as a server constant. In chatservice/src/main/scala/chatroom/grpc/Constant.scala:

// Add a JWT Context Key
val JWT_CTX_KEY: Context.Key[DecodedJWT] = Context.key("jwt")

// Add a JWT Context Key
val USER_ID_CTX_KEY: Context.Key[String] = Context.key("userId")

Server Interceptor - Metadata to Context

Since the server interceptor can capture the Metadata, we can also use it to propagate the information into a Context variable.

Let’s implement the full on JWT Interceptor so it will:

1. Capture the JWT token from Metadata
2. Verify that the token is valid
3. Converting the token into a DecodedJWT object, and store both the DecodedJWT and the User ID values into respective contexts:

override def interceptCall[ReqT, RespT](serverCall: ServerCall[ReqT, RespT], metadata: Metadata,
                                        serverCallHandler: ServerCallHandler[ReqT, RespT]) = {
  //  Get token from Metadata
  val token = metadata.get(Constant.JWT_METADATA_KEY)
  logger.info(s"interceptCall token: $token")

  if (token == null) {
    serverCall.close(Status.UNAUTHENTICATED.withDescription("JWT Token is missing from Metadata"), metadata)
    JwtServerInterceptor.NOOP_LISTENER[ReqT]
  }
  else {
    try {
      val jwt = verifier.verify(token)
      val ctx = Context.current.withValue(Constant.USER_ID_CTX_KEY, jwt.getSubject).withValue(Constant.JWT_CTX_KEY, jwt)
      logger.info(s"jwt.getPayload ${jwt.getPayload}")
      Contexts.interceptCall(ctx, serverCall, metadata, serverCallHandler)
    }
    catch  {
      case e:Exception =>
        logger.info("Verification failed - Unauthenticated!")
        serverCall.close(Status.UNAUTHENTICATED.withDescription(e.getMessage).withCause(e), metadata)
        JwtServerInterceptor.NOOP_LISTENER[ReqT]
    }
  }
}

Note: The magic here is Context ctx = Context.current().withValue(…) to capture the context value, and subsequently, using Contexts.interceptCall(…) to propagate the context to the service implementation.

Client Interceptor - Context to Metadata

Similarly, you can propagate context value to another service over network boundary by converting the context value into a Metadata. You can do this in a Client Interceptor.

Open chatservice/src/main/scala/chatroom/grpc/JwtClientInterceptor.scala, and implement the SimpleForwardingCall.start(…) method.

override def interceptCall[ReqT, RespT](methodDescriptor: MethodDescriptor[ReqT, RespT], callOptions: CallOptions, channel: Channel) = {
  new ForwardingClientCall.SimpleForwardingClientCall[ReqT, RespT](channel.newCall(methodDescriptor, callOptions)) {
    override def start(responseListener:ClientCall.Listener[RespT], headers:Metadata):Unit = {
      // TODO Convert JWT Context to Metadata header
      val jwt: DecodedJWT = Constant.JWT_CTX_KEY.get
      if (jwt != null) {
        headers.put(Constant.JWT_METADATA_KEY, jwt.getToken)
      }
      super.start(responseListener, headers)
    }
  }
}

Register client interceptor

You can attach a client interceptor to a channel similarly to how we attached the metadata interceptor. For example, in ChatServer.main(…), we can attach this client interceptor to authChannel:

// TODO Add trace interceptor
val authChannel = ManagedChannelBuilder.forTarget("localhost:9091")
  .intercept(new JwtClientInterceptor)
  .usePlaintext(true)
  .asInstanceOf[ManagedChannelBuilder[_]]
  .build()