Java Code Examples for io.grpc.stub.MetadataUtils#attachHeaders()

@Test
public void exchangeMetadataUnaryCall() throws Exception {
  TestServiceGrpc.TestServiceBlockingStub stub = blockingStub;

  // Capture the metadata exchange
  Metadata fixedHeaders = new Metadata();
  // Send a context proto (as it's in the default extension registry)
  Messages.SimpleContext contextValue =
      Messages.SimpleContext.newBuilder().setValue("dog").build();
  fixedHeaders.put(Util.METADATA_KEY, contextValue);
  stub = MetadataUtils.attachHeaders(stub, fixedHeaders);
  // .. and expect it to be echoed back in trailers
  AtomicReference<Metadata> trailersCapture = new AtomicReference<>();
  AtomicReference<Metadata> headersCapture = new AtomicReference<>();
  stub = MetadataUtils.captureMetadata(stub, headersCapture, trailersCapture);

  assertNotNull(stub.emptyCall(EMPTY));

  // Assert that our side channel object is echoed back in both headers and trailers
  Assert.assertEquals(contextValue, headersCapture.get().get(Util.METADATA_KEY));
  Assert.assertEquals(contextValue, trailersCapture.get().get(Util.METADATA_KEY));
}
 

/**
 * getStubWithJwt adds an AttachHeadersInterceptor to the stub, which will eventually attach a
 * header whose key is accessJwt and value is the access JWT stored in the current
 * DgraphAsyncClient object.
 *
 * @param stub the original stub that we should attach JWT to
 * @return the augmented stub with JWT
 */
protected DgraphGrpc.DgraphStub getStubWithJwt(DgraphGrpc.DgraphStub stub) {
  Lock readLock = jwtLock.readLock();
  readLock.lock();
  try {
    if (jwt != null && !jwt.getAccessJwt().isEmpty()) {
      Metadata metadata = new Metadata();
      metadata.put(
          Metadata.Key.of("accessJwt", Metadata.ASCII_STRING_MARSHALLER), jwt.getAccessJwt());
      return MetadataUtils.attachHeaders(stub, metadata);
    }

    return stub;
  } finally {
    readLock.unlock();
  }
}
 

public BrokerBlockingStub addSPNEGOTokenToMetadata(BrokerBlockingStub stub, String principal) {
    Metadata metadata = new Metadata();
    Metadata.Key<String> key = Metadata.Key.of("authorization", Metadata.ASCII_STRING_MARSHALLER);
    metadata.put(key, "Negotiate " + principal);
    stub = MetadataUtils.attachHeaders(stub, metadata);
    return stub;
}
 

public BrokerBlockingStub addSessionTokenToMetadata(BrokerBlockingStub stub, Session session) {
    Metadata metadata = new Metadata();
    Metadata.Key<String> key = Metadata.Key.of("authorization", Metadata.ASCII_STRING_MARSHALLER);
    metadata.put(key, "BrokerSession " + SessionTokenUtils.marshallSessionToken(session));
    stub = MetadataUtils.attachHeaders(stub, metadata);
    return stub;
}
 

/** Appends RemoteExecutionMetadata to the GRPC headers. */
public static <Stub extends AbstractStub<Stub>> Stub wrapStubToSendMetadata(
    Stub grpcStub, RemoteExecutionMetadata metadata) {
  Metadata extraHeaders = new Metadata();
  extraHeaders.put(REMOTE_EXECUTION_METADATA_KEY, metadata.toByteArray());
  return MetadataUtils.attachHeaders(grpcStub, extraHeaders);
}
 

public void setSessionToken(String sessionToken) {
    // Set the session token in the 'authorization' header
    Metadata metadata = new Metadata();
    Metadata.Key<String> key = Metadata.Key.of("authorization", Metadata.ASCII_STRING_MARSHALLER);
    metadata.put(key, REQUEST_AUTH_HEADER + " " + sessionToken);
    stub = MetadataUtils.attachHeaders(stub, metadata);
}
 

@Test
void exchangeHeadersUnaryCall_grpcMetadata() throws Exception {
    final Metadata metadata = new Metadata();
    metadata.put(TestServiceImpl.EXTRA_HEADER_KEY, "dog");

    TestServiceBlockingStub stub = MetadataUtils.attachHeaders(blockingStub, metadata);

    final AtomicReference<Metadata> headers = new AtomicReference<>();
    final AtomicReference<Metadata> trailers = new AtomicReference<>();
    stub = MetadataUtils.captureMetadata(stub, headers, trailers);

    assertThat(stub.emptyCall(EMPTY)).isNotNull();

    final HttpHeaders clientHeaders = CLIENT_HEADERS_CAPTURE.get();
    assertThat(clientHeaders.get(HttpHeaderNames.TE))
            .isEqualTo(HttpHeaderValues.TRAILERS.toString());

    // Assert that our side channel object is echoed back in both headers and trailers
    assertThat(clientHeaders.get(TestServiceImpl.EXTRA_HEADER_NAME)).isEqualTo("dog");
    assertThat(SERVER_TRAILERS_CAPTURE.get().get(TestServiceImpl.EXTRA_HEADER_NAME)).isEqualTo("dog");

    assertThat(headers.get()).isNull();
    assertThat(trailers.get().get(TestServiceImpl.EXTRA_HEADER_KEY)).isEqualTo("dog");
    assertThat(trailers.get().getAll(TestServiceImpl.STRING_VALUE_KEY)).containsExactly(
            StringValue.newBuilder().setValue("hello").build(),
            StringValue.newBuilder().setValue("world").build());

    checkRequestLog((rpcReq, rpcRes, grpcStatus) -> {
        assertThat(rpcReq.params()).containsExactly(EMPTY);
        assertThat(rpcRes.get()).isEqualTo(EMPTY);
    });
}
 

public HealthStub getHealthStub() {
  HealthStub returnValue = null;
  if (this.channel != null) {
    returnValue = MetadataUtils.attachHeaders(HealthGrpc.newStub(this.channel), metadata);
  }
  return returnValue;
}
 

private static DgraphClient createDgraphClient(boolean withAuthHeader) {
  ManagedChannel channel =
      ManagedChannelBuilder.forAddress(TEST_HOSTNAME, TEST_PORT).usePlaintext().build();
  DgraphStub stub = DgraphGrpc.newStub(channel);

  if (withAuthHeader) {
    Metadata metadata = new Metadata();
    metadata.put(
        Metadata.Key.of("auth-token", Metadata.ASCII_STRING_MARSHALLER), "the-auth-token-value");
    stub = MetadataUtils.attachHeaders(stub, metadata);
  }

  return new DgraphClient(stub);
}
 

public HealthBlockingStub getHealthBlockingStub() {
  HealthBlockingStub returnValue = null;
  if (this.channel != null) {
    returnValue = MetadataUtils.attachHeaders(HealthGrpc.newBlockingStub(this.channel), metadata);
  }
  return returnValue;
}
 

public JwtAuthBlockingClient(Channel channel, String token) {
    JwtAuthTestServiceGrpc.JwtAuthTestServiceBlockingStub stub = JwtAuthTestServiceGrpc.newBlockingStub(channel);
    //add metadata
    Metadata metadata = new Metadata();
    metadata.put(Metadata.Key.of("Authorization", Metadata.ASCII_STRING_MARSHALLER), token);
    blockingStub = MetadataUtils.attachHeaders(stub,metadata);
}
 

public ThrottlingBlockingClient(Channel channel, String token) {
    ThrottlingTestServiceGrpc.ThrottlingTestServiceBlockingStub stub =
            ThrottlingTestServiceGrpc.newBlockingStub(channel);
    //add metadata
    Metadata metadata = new Metadata();
    metadata.put(Metadata.Key.of("Authorization", Metadata.ASCII_STRING_MARSHALLER), token);
    blockingStub = MetadataUtils.attachHeaders(stub,metadata);
}
 

@Test
public void exchangeMetadataStreamingCall() throws Exception {
  TestServiceGrpc.TestServiceStub stub = asyncStub;

  // Capture the metadata exchange
  Metadata fixedHeaders = new Metadata();
  // Send a context proto (as it's in the default extension registry)
  Messages.SimpleContext contextValue =
      Messages.SimpleContext.newBuilder().setValue("dog").build();
  fixedHeaders.put(Util.METADATA_KEY, contextValue);
  stub = MetadataUtils.attachHeaders(stub, fixedHeaders);
  // .. and expect it to be echoed back in trailers
  AtomicReference<Metadata> trailersCapture = new AtomicReference<>();
  AtomicReference<Metadata> headersCapture = new AtomicReference<>();
  stub = MetadataUtils.captureMetadata(stub, headersCapture, trailersCapture);

  List<Integer> responseSizes = Arrays.asList(50, 100, 150, 200);
  Messages.StreamingOutputCallRequest.Builder streamingOutputBuilder =
      Messages.StreamingOutputCallRequest.newBuilder();
  for (Integer size : responseSizes) {
    streamingOutputBuilder.addResponseParameters(
        ResponseParameters.newBuilder().setSize(size).setIntervalUs(0));
  }
  final Messages.StreamingOutputCallRequest request = streamingOutputBuilder.build();

  StreamRecorder<Messages.StreamingOutputCallResponse> recorder = StreamRecorder.create();
  StreamObserver<Messages.StreamingOutputCallRequest> requestStream =
      stub.fullDuplexCall(recorder);

  final int numRequests = 10;
  for (int ix = numRequests; ix > 0; --ix) {
    requestStream.onNext(request);
  }
  requestStream.onCompleted();
  recorder.awaitCompletion();
  assertSuccess(recorder);
  org.junit.Assert.assertEquals(responseSizes.size() * numRequests, recorder.getValues().size());

  // Assert that our side channel object is echoed back in both headers and trailers
  Assert.assertEquals(contextValue, headersCapture.get().get(Util.METADATA_KEY));
  Assert.assertEquals(contextValue, trailersCapture.get().get(Util.METADATA_KEY));
}
 

@Test
public void exchangeMetadataStreamingCall() throws Exception {
  TestServiceGrpc.TestServiceStub stub = asyncStub;

  // Capture the metadata exchange
  Metadata fixedHeaders = new Metadata();
  // Send a context proto (as it's in the default extension registry)
  Messages.SimpleContext contextValue =
      Messages.SimpleContext.newBuilder().setValue("dog").build();
  fixedHeaders.put(Util.METADATA_KEY, contextValue);
  stub = MetadataUtils.attachHeaders(stub, fixedHeaders);
  // .. and expect it to be echoed back in trailers
  AtomicReference<Metadata> trailersCapture = new AtomicReference<>();
  AtomicReference<Metadata> headersCapture = new AtomicReference<>();
  stub = MetadataUtils.captureMetadata(stub, headersCapture, trailersCapture);

  List<Integer> responseSizes = Arrays.asList(50, 100, 150, 200);
  Messages.StreamingOutputCallRequest.Builder streamingOutputBuilder =
      Messages.StreamingOutputCallRequest.newBuilder();
  for (Integer size : responseSizes) {
    streamingOutputBuilder.addResponseParameters(
        ResponseParameters.newBuilder().setSize(size).setIntervalUs(0));
  }
  final Messages.StreamingOutputCallRequest request = streamingOutputBuilder.build();

  StreamRecorder<Messages.StreamingOutputCallResponse> recorder = StreamRecorder.create();
  StreamObserver<Messages.StreamingOutputCallRequest> requestStream =
      stub.fullDuplexCall(recorder);

  final int numRequests = 10;
  for (int ix = numRequests; ix > 0; --ix) {
    requestStream.onNext(request);
  }
  requestStream.onCompleted();
  recorder.awaitCompletion();
  assertSuccess(recorder);
  org.junit.Assert.assertEquals(responseSizes.size() * numRequests, recorder.getValues().size());

  // Assert that our side channel object is echoed back in both headers and trailers
  Assert.assertEquals(contextValue, headersCapture.get().get(Util.METADATA_KEY));
  Assert.assertEquals(contextValue, trailersCapture.get().get(Util.METADATA_KEY));
}
 

private boolean uploadQueue(@NonNull final String tid) {
    List<Location> pack = new ArrayList<>();
    boolean completed = false;

    for (; ; ) {
        pack.clear();
        int packSize;
        synchronized (uploadQueue) {
            packSize = Math.min(uploadQueue.size(), PACK_SIZE);
            for (int i = 0; i < packSize; i++) {
                Location head = uploadQueue.peekFirst();
                if (head != null) {
                    pack.add(head);
                }
            }
        }
        if (pack.size() > 0) {
            TrackerOuterClass.AddLocationsRequest.Builder req =
                    TrackerOuterClass.AddLocationsRequest
                            .newBuilder();
            req.setTid(tid);
            for (Location location : pack) {
                req.addLocations(marshall(tid, location));
            }
            if (ch == null) ch = grpcChannelProvider.channel();
            TrackerGrpc.TrackerBlockingStub grpcStub = getGrpcStub();
            try {
                Metadata headers = new Metadata();
                headers.put(Keys.wsseKey, Wsse.getToken());

                grpcStub = MetadataUtils.attachHeaders(grpcStub, headers);
                TrackerOuterClass.AddLocationResponse resp = grpcStub.addLocations(req.build());
                if (resp.getOk()) {
                    for (int i = 0; i < packSize; i++) {
                        synchronized (uploadQueue) {
                            uploadQueue.pollFirst();
                        }
                    }
                } else {
                    ErrorsObservable.notify(new Exception("upload_failed"), false);
                    break;
                }
            } catch (Exception e) {
                ErrorsObservable.notify(e, true);
                break;
            }
        }
        synchronized (uploadQueue) {
            if (uploadQueue.size() == 0) {
                completed = true;
                break;
            }
        }
    }
    return completed;
}
 

public <T extends AbstractStub<T>, R> R withToken(Metadata meta, T stub, Function<T, R> f) {
    T newStub = MetadataUtils.attachHeaders(stub, meta);
    return f.apply(newStub);
}
 

public <T extends AbstractStub<T>, R> CompletableFuture<R> withToken(T stub, Function<T, ListenableFuture<R>> f) {
    T newStub = MetadataUtils.attachHeaders(stub, metadata);
    TaskManager<R> task = new TaskManager<>(FutureConverter.toCompletableFuture(f.apply(newStub)), config.getRetryOptions());
    return task.scheduleTask(0);
}
 

/**
 * Returns the gRPC-generated {@link ReleaseServiceFutureStub}
 * object that represents the capabilities of the Tiller server.
 *
 * <p>This method will never return {@code null}.</p>
 *
 * <p>Overrides of this method must never return {@code null}.</p>
 *
 * @return a non-{@code null} {@link ReleaseServiceFutureStub}
 *
 * @see ReleaseServiceFutureStub
 */
public ReleaseServiceFutureStub getReleaseServiceFutureStub() {
  ReleaseServiceFutureStub returnValue = null;
  if (this.channel != null) {
    returnValue = MetadataUtils.attachHeaders(ReleaseServiceGrpc.newFutureStub(this.channel), metadata);
  }
  return returnValue;
}
 

/**
 * Returns the gRPC-generated {@link ReleaseServiceBlockingStub}
 * object that represents the capabilities of the Tiller server.
 *
 * <p>This method will never return {@code null}.</p>
 *
 * <p>Overrides of this method must never return {@code null}.</p>
 *
 * @return a non-{@code null} {@link ReleaseServiceBlockingStub}
 *
 * @see ReleaseServiceBlockingStub
 */
public ReleaseServiceBlockingStub getReleaseServiceBlockingStub() {
  ReleaseServiceBlockingStub returnValue = null;
  if (this.channel != null) {
    returnValue = MetadataUtils.attachHeaders(ReleaseServiceGrpc.newBlockingStub(this.channel), metadata);
  }
  return returnValue;
}
 

/**
 * 自定义 token内容
 *
 * @param stub
 * @return: T
 * @author: fruiqi
 * @date: 19-5-15 上午11:02
 */
public static <T extends AbstractStub<T>> T blockingStubToken(T stub, String token) {
	Metadata metadata = new Metadata();
	metadata.put(AUTHORIZATION, String.format("%s %s", BEARER_TYPE, token));
	return MetadataUtils.attachHeaders(stub, metadata);
}