Metadata/Context Propagation

Server Context

First define the context as a server constant. In chat-service/src/main/java/com/example/chat/grpc/Constant.java:

// TODO Add a JWT Context Key
public static final Context.Key<DecodedJWT> JWT_CTX_KEY = Context.key("jwt");

public static final Context.Key<String> USER_ID_CTX_KEY = 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:

// TODO Get token from Metadata
// TODO If token is nul, or is invalid,
// close the call with Status.UNAUTHENTICATED
public <ReqT, RespT> ServerCall.Listener<ReqT> interceptCall(
  ServerCall<ReqT, RespT> serverCall, Metadata metadata,
  ServerCallHandler<ReqT, RespT> serverCallHandler) {

  String token = metadata.get(Constant.JWT_METADATA_KEY);
  if (token == null) {
    serverCall.close(Status.UNAUTHENTICATED
      .withDescription("JWT Token is missing from Metadata"), metadata);
    return NOOP_LISTENER;
  }

  try {
    DecodedJWT jwt = verifier.verify(token);
    Context ctx = Context.current()
        .withValue(Constant.USER_ID_CTX_KEY, jwt.getSubject())
        .withValue(Constant.JWT_CTX_KEY, jwt);
    return Contexts.interceptCall(ctx, serverCall, metadata, serverCallHandler);
  } catch (Exception e) {
    System.out.println("Verification failed - Unauthenticated!");
    serverCall.close(Status.UNAUTHENTICATED
      .withDescription(e.getMessage()).withCause(e), metadata);
    return NOOP_LISTENER;
  }
}

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 chat-service/src/main/java/com/example/chat/grpc/JwtClientInterceptor.java, and implement the SimpleForwardingCall.start(…) method.

public class JwtClientInterceptor implements ClientInterceptor {
  @Override
  public <ReqT, RespT> ClientCall<ReqT, RespT> interceptCall(...) {

    return new ForwardingClientCall
      .SimpleForwardingClientCall<ReqT, RespT>(...) {
      @Override
      public void start(...) {
        // TODO Convert JWT Context to Metadata header
        DecodedJWT jwt = 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:

final ManagedChannel authChannel = ManagedChannelBuilder
  .forTarget("localhost:9091")
  .intercept(new JwtClientInterceptor())
  .usePlaintext(true)
  .build();