Java Code Examples for io.grpc.Metadata#Key

@Test
public void statsHeaderMalformed() {
  // Construct a malformed header and make sure parsing it will throw
  byte[] statsHeaderValue = new byte[]{1};
  Metadata.Key<byte[]> arbitraryStatsHeader =
      Metadata.Key.of("grpc-tags-bin", Metadata.BINARY_BYTE_MARSHALLER);
  try {
    tagCtxSerializer.fromByteArray(statsHeaderValue);
    fail("Should have thrown");
  } catch (Exception e) {
    // Expected
  }

  // But the header key will return a default context for it
  Metadata headers = new Metadata();
  assertNull(headers.get(censusStats.statsHeader));
  headers.put(arbitraryStatsHeader, statsHeaderValue);
  assertSame(tagger.empty(), headers.get(censusStats.statsHeader));
}
 

/**
 * Echo the request headers from a client into response headers and trailers. Useful for
 * testing end-to-end metadata propagation.
 */
private static ServerInterceptor echoRequestHeadersInterceptor(final Metadata.Key<?>... keys) {
  final Set<Metadata.Key<?>> keySet = new HashSet<>(Arrays.asList(keys));
  return new ServerInterceptor() {
    @Override
    public <ReqT, RespT> ServerCall.Listener<ReqT> interceptCall(
        ServerCall<ReqT, RespT> call,
        final Metadata requestHeaders,
        ServerCallHandler<ReqT, RespT> next) {
      return next.startCall(new SimpleForwardingServerCall<ReqT, RespT>(call) {
            @Override
            public void sendHeaders(Metadata responseHeaders) {
              responseHeaders.merge(requestHeaders, keySet);
              super.sendHeaders(responseHeaders);
            }

            @Override
            public void close(Status status, Metadata trailers) {
              trailers.merge(requestHeaders, keySet);
              super.close(status, trailers);
            }
          }, requestHeaders);
    }
  };
}
 

/**
 * Echoes request headers with the specified key(s) from a client into response headers only.
 */
private static ServerInterceptor echoRequestMetadataInHeaders(final Metadata.Key<?>... keys) {
  final Set<Metadata.Key<?>> keySet = new HashSet<>(Arrays.asList(keys));
  return new ServerInterceptor() {
    @Override
    public <ReqT, RespT> ServerCall.Listener<ReqT> interceptCall(
        ServerCall<ReqT, RespT> call,
        final Metadata requestHeaders,
        ServerCallHandler<ReqT, RespT> next) {
      return next.startCall(new SimpleForwardingServerCall<ReqT, RespT>(call) {
        @Override
        public void sendHeaders(Metadata responseHeaders) {
          responseHeaders.merge(requestHeaders, keySet);
          super.sendHeaders(responseHeaders);
        }

        @Override
        public void close(Status status, Metadata trailers) {
          super.close(status, trailers);
        }
      }, requestHeaders);
    }
  };
}
 

@Override
public <ReqT, RespT> ClientCall<ReqT, RespT> interceptCall(
    MethodDescriptor<ReqT, RespT> methodDescriptor, CallOptions callOptions, Channel channel) {
  return new ForwardingClientCall.SimpleForwardingClientCall<ReqT, RespT>(
      channel.newCall(methodDescriptor, callOptions)) {
    @Override
    public void start(Listener<RespT> responseListener, Metadata headers) {
      // TODO: Here set request metadata
      Metadata.Key<String> email_key =
          Metadata.Key.of("email", Metadata.ASCII_STRING_MARSHALLER);
      Metadata.Key<String> dev_key =
          Metadata.Key.of("developer_key", Metadata.ASCII_STRING_MARSHALLER);
      Metadata.Key<String> source_key =
          Metadata.Key.of("source", Metadata.ASCII_STRING_MARSHALLER);

      headers.put(email_key, client2Email);
      headers.put(dev_key, client2DevKey);
      headers.put(source_key, "PythonClient");
      super.start(responseListener, headers);
    }
  };
}
 

@Test
public void statsHeaderMalformed() {
  // Construct a malformed header and make sure parsing it will throw
  byte[] statsHeaderValue = new byte[]{1};
  Metadata.Key<byte[]> arbitraryStatsHeader =
      Metadata.Key.of("grpc-tags-bin", Metadata.BINARY_BYTE_MARSHALLER);
  try {
    tagCtxSerializer.fromByteArray(statsHeaderValue);
    fail("Should have thrown");
  } catch (Exception e) {
    // Expected
  }

  // But the header key will return a default context for it
  Metadata headers = new Metadata();
  assertNull(headers.get(censusStats.statsHeader));
  headers.put(arbitraryStatsHeader, statsHeaderValue);
  assertSame(tagger.empty(), headers.get(censusStats.statsHeader));
}
 

private static Metadata.Key<String> getHeader(String headerName) {
    Metadata.Key<String> key = headerCache.get(headerName);
    if (key == null) {
        key = Metadata.Key.of(headerName, Metadata.ASCII_STRING_MARSHALLER);
        headerCache.put(headerName, key);
    }
    return key;
}
 

@Test
public void createRequestHeaders_sanitizes() {
  Metadata metaData = new Metadata();

  // Intentionally being explicit here rather than relying on any pre-defined lists of headers,
  // since the goal of this test is to validate the correctness of such lists in the first place.
  metaData.put(GrpcUtil.CONTENT_TYPE_KEY, "to-be-removed");
  metaData.put(GrpcUtil.USER_AGENT_KEY, "to-be-removed");
  metaData.put(GrpcUtil.TE_HEADER, "to-be-removed");


  Metadata.Key<String> userKey = Metadata.Key.of("user-key", Metadata.ASCII_STRING_MARSHALLER);
  String userValue = "user-value";
  metaData.put(userKey, userValue);

  String path = "//testServerice/test";
  String authority = "localhost";
  String userAgent = "useragent";

  List<Header> headers = Headers.createRequestHeaders(
      metaData,
      path,
      authority,
      userAgent,
      false,
      false);

  // 7 reserved headers, 1 user header
  assertEquals(7 + 1, headers.size());
  // Check the 3 reserved headers that are non pseudo
  // Users can not create pseudo headers keys so no need to check for them here
  assertThat(headers).contains(Headers.CONTENT_TYPE_HEADER);
  assertThat(headers).contains(new Header(GrpcUtil.USER_AGENT_KEY.name(), userAgent));
  assertThat(headers).contains(new Header(GrpcUtil.TE_HEADER.name(), GrpcUtil.TE_TRAILERS));
  // Check the user header is in tact
  assertThat(headers).contains(new Header(userKey.name(), userValue));
}
 

@Test
public void contextScopeStackingWorks() {
    Metadata.Key<String> key = Metadata.Key.of("k", Metadata.ASCII_STRING_MARSHALLER);
    AmbientContext.initialize(Context.current()).run(() -> {
        AmbientContext.current().put(key, "outer");
        assertThat(AmbientContext.current().get(key)).isEqualTo("outer");

        AmbientContext.current().fork(Context.current()).run(() -> {
            AmbientContext.current().put(key, "inner");
            assertThat(AmbientContext.current().get(key)).isEqualTo("inner");
        });

        assertThat(AmbientContext.current().get(key)).isEqualTo("outer");
    });
}
 

@Test
public void alwaysLoggedMetadata_grpcTraceBin() throws Exception {
  Metadata.Key<byte[]> key
      = Metadata.Key.of("grpc-trace-bin", Metadata.BINARY_BYTE_MARSHALLER);
  Metadata metadata = new Metadata();
  metadata.put(key, new byte[1]);
  int zeroHeaderBytes = 0;
  MaybeTruncated<io.grpc.binarylog.v1.Metadata.Builder> pair =
      createMetadataProto(metadata, zeroHeaderBytes);
  assertEquals(
      key.name(),
      Iterables.getOnlyElement(pair.proto.getEntryBuilderList()).getKey());
  assertFalse(pair.truncated);
}
 

@Test
public void multipleContextTransfersOneHopSync() throws Exception {
    Metadata.Key<String> ctxKey = Metadata.Key.of("ctx-context-key", Metadata.ASCII_STRING_MARSHALLER);
    Metadata.Key<String> l5dKey = Metadata.Key.of("l5d-context-key", Metadata.ASCII_STRING_MARSHALLER);
    String expectedCtxValue = "context-value";
    AtomicReference<String> ctxValue = new AtomicReference<>();
    AtomicReference<String> l5dValue = new AtomicReference<>();

    // Service
    GreeterGrpc.GreeterImplBase svc = new GreeterGrpc.GreeterImplBase() {
        @Override
        public void sayHello(HelloRequest request, StreamObserver<HelloResponse> responseObserver) {
            ctxValue.set(AmbientContext.current().get(ctxKey));
            l5dValue.set(AmbientContext.current().get(l5dKey));
            responseObserver.onNext(HelloResponse.newBuilder().setMessage("Hello " + request.getName()).build());
            responseObserver.onCompleted();
        }
    };

    // Plumbing
    serverRule1.getServiceRegistry().addService(ServerInterceptors.intercept(svc,
            new AmbientContextServerInterceptor("ctx-"),
            new AmbientContextServerInterceptor("l5d-")));

    GreeterGrpc.GreeterBlockingStub stub = GreeterGrpc
            .newBlockingStub(serverRule1.getChannel())
            .withInterceptors(
                    new AmbientContextClientInterceptor("ctx-"),
                    new AmbientContextClientInterceptor("l5d-"));

    // Test
    AmbientContext.initialize(Context.current()).run(() -> {
        AmbientContext.current().put(ctxKey, expectedCtxValue);
        AmbientContext.current().put(l5dKey, expectedCtxValue);
        stub.sayHello(HelloRequest.newBuilder().setName("world").build());
    });

    assertThat(ctxValue.get()).isEqualTo(expectedCtxValue);
    assertThat(l5dValue.get()).isEqualTo(expectedCtxValue);
}
 

@Test
public void neverLoggedMetadata_grpcStatusDetilsBin() throws Exception {
  Metadata.Key<byte[]> key
      = Metadata.Key.of("grpc-status-details-bin", Metadata.BINARY_BYTE_MARSHALLER);
  Metadata metadata = new Metadata();
  metadata.put(key, new byte[1]);
  int unlimitedHeaderBytes = Integer.MAX_VALUE;
  MaybeTruncated<io.grpc.binarylog.v1.Metadata.Builder> pair
      = createMetadataProto(metadata, unlimitedHeaderBytes);
  assertThat(pair.proto.getEntryBuilderList()).isEmpty();
  assertFalse(pair.truncated);
}
 

/** This assumes the client limits metadata size to GrpcUtil.DEFAULT_MAX_HEADER_LIST_SIZE. */
@Test
public void clientChecksInboundMetadataSize_trailer() throws Exception {
  server.start(serverListener);
  client = newClientTransport(server);
  startTransport(client, mockClientTransportListener);
  MockServerTransportListener serverTransportListener
      = serverListener.takeListenerOrFail(TIMEOUT_MS, TimeUnit.MILLISECONDS);
  serverTransport = serverTransportListener.transport;

  Metadata.Key<String> tellTaleKey
      = Metadata.Key.of("tell-tale", Metadata.ASCII_STRING_MARSHALLER);
  Metadata tooLargeMetadata = new Metadata();
  tooLargeMetadata.put(tellTaleKey, "true");
  tooLargeMetadata.put(
      Metadata.Key.of("foo-bin", Metadata.BINARY_BYTE_MARSHALLER),
      new byte[GrpcUtil.DEFAULT_MAX_HEADER_LIST_SIZE]);

  ClientStream clientStream =
      client.newStream(methodDescriptor, new Metadata(), callOptions);
  ClientStreamListenerBase clientStreamListener = new ClientStreamListenerBase();
  clientStream.start(clientStreamListener);

  clientStream.writeMessage(methodDescriptor.streamRequest("foo"));
  clientStream.halfClose();
  clientStream.request(1);

  StreamCreation serverStreamCreation
      = serverTransportListener.takeStreamOrFail(TIMEOUT_MS, TimeUnit.MILLISECONDS);

  serverStreamCreation.stream.request(1);
  serverStreamCreation.stream.writeHeaders(new Metadata());
  serverStreamCreation.stream.writeMessage(methodDescriptor.streamResponse("response"));
  serverStreamCreation.stream.close(Status.OK, tooLargeMetadata);

  Status status = clientStreamListener.status.get(TIMEOUT_MS, TimeUnit.MILLISECONDS);
  List<Status.Code> codeOptions = Arrays.asList(
      Status.Code.UNKNOWN, Status.Code.RESOURCE_EXHAUSTED, Status.Code.INTERNAL);
  if (!codeOptions.contains(status.getCode())) {
    fail("Status code was not expected: " + status);
  }
  Metadata metadata = clientStreamListener.trailers.get(TIMEOUT_MS, TimeUnit.MILLISECONDS);
  assertNull(metadata.get(tellTaleKey));
}
 

public Metadata.Key<byte[]> asMetaBytesKey() {
    return bytesMetaKey;
}
 

private String validateId(String id, Metadata.Key key) {
    if (!IdValidateUtils.validateId(id)) {
        throw Status.INVALID_ARGUMENT.withDescription("invalid " + key.name()).asRuntimeException();
    }
    return id;
}
 

/** This assumes the client limits metadata size to GrpcUtil.DEFAULT_MAX_HEADER_LIST_SIZE. */
@Test
public void clientChecksInboundMetadataSize_trailer() throws Exception {
  server.start(serverListener);
  client = newClientTransport(server);
  startTransport(client, mockClientTransportListener);
  MockServerTransportListener serverTransportListener
      = serverListener.takeListenerOrFail(TIMEOUT_MS, TimeUnit.MILLISECONDS);
  serverTransport = serverTransportListener.transport;

  Metadata.Key<String> tellTaleKey
      = Metadata.Key.of("tell-tale", Metadata.ASCII_STRING_MARSHALLER);
  Metadata tooLargeMetadata = new Metadata();
  tooLargeMetadata.put(tellTaleKey, "true");
  tooLargeMetadata.put(
      Metadata.Key.of("foo-bin", Metadata.BINARY_BYTE_MARSHALLER),
      new byte[GrpcUtil.DEFAULT_MAX_HEADER_LIST_SIZE]);

  ClientStream clientStream =
      client.newStream(methodDescriptor, new Metadata(), callOptions);
  ClientStreamListenerBase clientStreamListener = new ClientStreamListenerBase();
  clientStream.start(clientStreamListener);

  clientStream.writeMessage(methodDescriptor.streamRequest("foo"));
  clientStream.halfClose();
  clientStream.request(1);

  StreamCreation serverStreamCreation
      = serverTransportListener.takeStreamOrFail(TIMEOUT_MS, TimeUnit.MILLISECONDS);

  serverStreamCreation.stream.request(1);
  serverStreamCreation.stream.writeHeaders(new Metadata());
  serverStreamCreation.stream.writeMessage(methodDescriptor.streamResponse("response"));
  serverStreamCreation.stream.close(Status.OK, tooLargeMetadata);

  Status status = clientStreamListener.status.get(TIMEOUT_MS, TimeUnit.MILLISECONDS);
  List<Status.Code> codeOptions = Arrays.asList(
      Status.Code.UNKNOWN, Status.Code.RESOURCE_EXHAUSTED, Status.Code.INTERNAL);
  if (!codeOptions.contains(status.getCode())) {
    fail("Status code was not expected: " + status);
  }
  Metadata metadata = clientStreamListener.trailers.get(TIMEOUT_MS, TimeUnit.MILLISECONDS);
  assertNull(metadata.get(tellTaleKey));
}
 

@Override
public Metadata.Key<byte[]> getTrailerKey() {
  return createKey("google.ads.googleads.v3.errors.googleadsfailure-bin");
}
 

protected static Metadata.Key<byte[]> createKey(String trailerKey) {
  return Metadata.Key.of(trailerKey, Metadata.BINARY_BYTE_MARSHALLER);
}
 

/**
 * Remove all values for the given key. If there were no values, {@code null} is returned.
 *
 * @throws IllegalStateException  if the AmbientContext is frozen
 */
public <T> Iterable<T> removeAll(Metadata.Key<T> key) {
    checkFreeze();
    return contextMetadata.removeAll(key);
}
 

/**
 * Adds the {@code key, value} pair. If {@code key} already has values, {@code value} is added to
 * the end. Duplicate values for the same key are permitted.
 *
 * @throws NullPointerException if key or value is null
 * @throws IllegalStateException  if the AmbientContext is frozen
 */
public <T> void put(Metadata.Key<T> key, T value) {
    checkFreeze();
    contextMetadata.put(key, value);
}
 

/**
 * Returns true if a value is defined for the given key.
 *
 * <p>This is done by linear search, so if it is followed by {@link #get} or {@link #getAll},
 * prefer calling them directly and checking the return value against {@code null}.
 */
public boolean containsKey(Metadata.Key<?> key) {
    return contextMetadata.containsKey(key);
}