Java Code Examples for software.amazon.awssdk.core.interceptor.Context#BeforeTransmission

@Override
public void beforeTransmission(Context.BeforeTransmission context, ExecutionAttributes attrs) {
  logRetryAttempt(attrs);

  String serviceName = attrs.getAttribute(SdkExecutionAttribute.SERVICE_NAME);
  String opName = attrs.getAttribute(SdkExecutionAttribute.OPERATION_NAME);
  String endpoint = serviceName + "." + opName;

  SdkHttpRequest request = context.httpRequest();

  IpcLogEntry logEntry = logger.createClientEntry()
      .withOwner("aws-sdk-java-v2")
      .withProtocol(IpcProtocol.http_1)
      .withHttpMethod(request.method().name())
      .withUri(request.getUri())
      .withEndpoint(endpoint)
      .withAttempt(extractAttempt(request))
      .withAttemptFinal(false); // Don't know if it is the final attempt

  request.headers().forEach((k, vs) -> vs.forEach(v -> logEntry.addRequestHeader(k, v)));

  attrs.putAttribute(LOG_ENTRY, logEntry.markStart());
}
 

private Context.BeforeTransmission captureBeforeTransmissionArg(ExecutionInterceptor interceptor, boolean isAsync) {
    ArgumentCaptor<Context.BeforeTransmission> beforeTransmissionArg = ArgumentCaptor.forClass(Context.BeforeTransmission.class);

    InOrder inOrder = Mockito.inOrder(interceptor);
    inOrder.verify(interceptor).beforeExecution(any(), any());
    inOrder.verify(interceptor).modifyRequest(any(), any());
    inOrder.verify(interceptor).beforeMarshalling(any(), any());
    inOrder.verify(interceptor).afterMarshalling(any(), any());
    inOrder.verify(interceptor).modifyAsyncHttpContent(any(), any());
    inOrder.verify(interceptor).modifyHttpContent(any(), any());
    inOrder.verify(interceptor).modifyHttpRequest(any(), any());
    inOrder.verify(interceptor).beforeTransmission(beforeTransmissionArg.capture(), any());
    inOrder.verify(interceptor).afterTransmission(any(), any());
    inOrder.verify(interceptor).modifyHttpResponse(any(), any());
    inOrder.verify(interceptor).modifyHttpResponseContent(any(), any());
    inOrder.verify(interceptor).beforeUnmarshalling(any(), any());
    if (isAsync) {
        inOrder.verify(interceptor).modifyAsyncHttpResponseContent(any(), any());
    }
    inOrder.verify(interceptor).afterUnmarshalling(any(), any());
    inOrder.verify(interceptor).modifyResponse(any(), any());
    inOrder.verify(interceptor).afterExecution(any(), any());
    verifyNoMoreInteractions(interceptor);
    return beforeTransmissionArg.getValue();
}
 

@Test
public void async_streamingInput_success_allInterceptorMethodsCalled()
        throws ExecutionException, InterruptedException, TimeoutException, IOException {
    // Given
    ExecutionInterceptor interceptor = mock(NoOpInterceptor.class, CALLS_REAL_METHODS);
    ProtocolRestJsonAsyncClient client = asyncClient(interceptor);
    StreamingInputOperationRequest request = StreamingInputOperationRequest.builder().build();
    stubFor(post(urlPathEqualTo(STREAMING_INPUT_PATH)).willReturn(aResponse().withStatus(200).withBody("")));

    // When
    client.streamingInputOperation(request, new NoOpAsyncRequestBody()).get(10, TimeUnit.SECONDS);

    // Expect
    Context.BeforeTransmission beforeTransmissionArg = captureBeforeTransmissionArg(interceptor, true);

    // TODO: The content should actually be empty to match responses. We can fix this by updating the StructuredJsonGenerator
    // to use null for NO-OP marshalling of payloads. This will break streaming POST operations for JSON because of a hack in
    // the MoveParametersToBodyStage, but we can move the logic from there into the query marshallers (why the hack exists)
    // and then everything should be good for JSON.
    assertThat(beforeTransmissionArg.requestBody().get().contentStreamProvider().newStream().read()).isEqualTo(-1);
    assertThat(beforeTransmissionArg.httpRequest().firstMatchingHeader(Header.CONTENT_LENGTH).get())
        .contains(Long.toString(0L));
}
 

@Test
public void sync_streamingInput_success_allInterceptorMethodsCalled() throws IOException {
    // Given
    ExecutionInterceptor interceptor = mock(NoOpInterceptor.class, CALLS_REAL_METHODS);
    ProtocolRestJsonClient client = client(interceptor);
    StreamingInputOperationRequest request = StreamingInputOperationRequest.builder().build();
    stubFor(post(urlPathEqualTo(STREAMING_INPUT_PATH)).willReturn(aResponse().withStatus(200).withBody("")));

    // When
    client.streamingInputOperation(request, RequestBody.fromBytes(new byte[] {0}));

    // Expect
    Context.BeforeTransmission beforeTransmissionArg = captureBeforeTransmissionArg(interceptor, false);
    assertThat(beforeTransmissionArg.requestBody().get().contentStreamProvider().newStream().read()).isEqualTo(0);
    assertThat(beforeTransmissionArg.httpRequest().firstMatchingHeader(Header.CONTENT_LENGTH).get())
        .contains(Long.toString(1L));
}
 

@Override
public void beforeTransmission(Context.BeforeTransmission context, ExecutionAttributes executionAttributes) {
    Subsegment subsegment = executionAttributes.getAttribute(entityKey);
    if (subsegment == null) {
        return;
    }

    Map<String, Object> awsProperties = subsegment.getAws();
    // beforeTransmission is run before every API call attempt
    // default value is set to -1 and will always be -1 on the first API call attempt
    // this value will be incremented by 1, so initial run will have a stored retryCount of 0
    int retryCount = (int) awsProperties.getOrDefault(EntityDataKeys.AWS.RETRIES_KEY, -1);
    awsProperties.put(EntityDataKeys.AWS.RETRIES_KEY, retryCount + 1);
}
 

/**
 * Before sending an http request. Will be called multiple times in the case of retries.
 */
@Override public void beforeTransmission(Context.BeforeTransmission context,
    ExecutionAttributes executionAttributes) {
  Span span = executionAttributes.getAttribute(SPAN);
  if (span == null) {
    // An evil interceptor deleted our attribute.
    return;
  }
  span.annotate("ws");
}
 

@Test
public void specifiedInOverrideConfig_shouldUse() {
    ExecutionInterceptor interceptor = Mockito.spy(AbstractExecutionInterceptor.class);

    String profileFileContent =
        "[default]\n" +
        "s3_use_arn_region = true\n";

    ProfileFile profileFile = ProfileFile.builder()
                                         .type(ProfileFile.Type.CONFIGURATION)
                                         .content(new StringInputStream(profileFileContent))
                                         .build();

    S3Client s3 = S3Client.builder()
                          .region(Region.US_WEST_2)
                          .credentialsProvider(AnonymousCredentialsProvider.create())
                          .overrideConfiguration(c -> c.defaultProfileFile(profileFile)
                                                       .defaultProfileName("default")
                                                       .addExecutionInterceptor(interceptor)
                                                       .retryPolicy(r -> r.numRetries(0)))
                          .build();

    String arn = "arn:aws:s3:us-banana-46:12345567890:accesspoint:foo";
    assertThatThrownBy(() -> s3.getObject(r -> r.bucket(arn).key("bar"))).isInstanceOf(SdkException.class);

    ArgumentCaptor<Context.BeforeTransmission> context = ArgumentCaptor.forClass(Context.BeforeTransmission.class);
    Mockito.verify(interceptor).beforeTransmission(context.capture(), any());

    String host = context.getValue().httpRequest().host();
    assertThat(host).contains("us-banana-46");
}
 

@Test(timeout = 10_000)
public void canBeEnabledViaProfileOnOverrideConfiguration() throws InterruptedException {
    ExecutionInterceptor interceptor = Mockito.spy(AbstractExecutionInterceptor.class);

    String profileFileContent =
        "[default]\n" +
        "aws_endpoint_discovery_enabled = true";

    ProfileFile profileFile = ProfileFile.builder()
                                         .type(ProfileFile.Type.CONFIGURATION)
                                         .content(new StringInputStream(profileFileContent))
                                         .build();

    DynamoDbClient dynamoDb = DynamoDbClient.builder()
                                            .region(Region.US_WEST_2)
                                            .credentialsProvider(AnonymousCredentialsProvider.create())
                                            .overrideConfiguration(c -> c.defaultProfileFile(profileFile)
                                                                         .defaultProfileName("default")
                                                                         .addExecutionInterceptor(interceptor)
                                                                         .retryPolicy(r -> r.numRetries(0)))
                                            .build();

    assertThatThrownBy(dynamoDb::listTables).isInstanceOf(SdkException.class);

    ArgumentCaptor<Context.BeforeTransmission> context;

    do {
        Thread.sleep(1);
        context = ArgumentCaptor.forClass(Context.BeforeTransmission.class);
        Mockito.verify(interceptor, atLeastOnce()).beforeTransmission(context.capture(), any());
    } while (context.getAllValues().size() < 2);

    assertThat(context.getAllValues()
                      .stream()
                      .anyMatch(v -> v.httpRequest()
                                      .firstMatchingHeader("X-Amz-Target")
                                      .map(h -> h.equals("DynamoDB_20120810.DescribeEndpoints"))
                                      .orElse(false)))
        .isTrue();
}
 

@Override
public void beforeTransmission(Context.BeforeTransmission context, ExecutionAttributes executionAttributes) {
    this.beforeTransmission = context.httpRequest();
}
 

@Override
public void beforeTransmission(Context.BeforeTransmission context, ExecutionAttributes executionAttributes) {
    assertThat(context.httpRequest().firstMatchingHeader("User-Agent").get()).containsIgnoringCase("io/sync");
    assertThat(context.httpRequest().firstMatchingHeader("User-Agent").get()).containsIgnoringCase("http/UrlConnection");
}
 

@Override
public void beforeTransmission(Context.BeforeTransmission context, ExecutionAttributes executionAttributes) {
    Optional<String> headers = context.httpRequest().firstMatchingHeader("User-agent");
    assertThat(headers).isPresent();
    assertThat(headers.get()).contains("hll/ddb-enh");
}
 

@Override
public void beforeTransmission(Context.BeforeTransmission context, ExecutionAttributes executionAttributes) {
    assertThat(context.httpRequest().firstMatchingHeader("User-Agent").get()).containsIgnoringCase("io/" + clientType.name());
    assertThat(context.httpRequest().firstMatchingHeader("User-Agent").get()).containsIgnoringCase("http/" + clientName);
}
 

@Override
public void beforeTransmission(Context.BeforeTransmission context, ExecutionAttributes executionAttributes) {
if (!(context.request() instanceof ListBucketsRequest)) {
    assertEquals(context.httpRequest().host(), US_BUCKET_NAME + ".s3-accelerate.amazonaws.com");
}
}
 

private void expectAllMethodsCalled(ExecutionInterceptor interceptor,
                                    SdkRequest inputRequest,
                                    Exception expectedException,
                                    boolean isAsync) {
    ArgumentCaptor<ExecutionAttributes> attributes = ArgumentCaptor.forClass(ExecutionAttributes.class);

    ArgumentCaptor<Context.BeforeExecution> beforeExecutionArg = ArgumentCaptor.forClass(Context.BeforeExecution.class);
    ArgumentCaptor<Context.BeforeMarshalling> modifyRequestArg = ArgumentCaptor.forClass(Context.BeforeMarshalling.class);
    ArgumentCaptor<Context.BeforeMarshalling> beforeMarshallingArg = ArgumentCaptor.forClass(Context.BeforeMarshalling.class);
    ArgumentCaptor<Context.AfterMarshalling> afterMarshallingArg = ArgumentCaptor.forClass(Context.AfterMarshalling.class);
    ArgumentCaptor<Context.BeforeTransmission> modifyHttpRequestArg = ArgumentCaptor.forClass(Context.BeforeTransmission.class);
    ArgumentCaptor<Context.BeforeTransmission> modifyHttpContentArg = ArgumentCaptor.forClass(Context.BeforeTransmission.class);
    ArgumentCaptor<Context.BeforeTransmission> modifyHttpContentAsyncArg = ArgumentCaptor.forClass(Context.BeforeTransmission.class);
    ArgumentCaptor<Context.BeforeTransmission> beforeTransmissionArg = ArgumentCaptor.forClass(Context.BeforeTransmission.class);
    ArgumentCaptor<Context.AfterTransmission> afterTransmissionArg = ArgumentCaptor.forClass(Context.AfterTransmission.class);
    ArgumentCaptor<Context.BeforeUnmarshalling> modifyHttpResponseArg = ArgumentCaptor.forClass(Context.BeforeUnmarshalling.class);
    ArgumentCaptor<Context.BeforeUnmarshalling> modifyHttpResponseContentArg = ArgumentCaptor.forClass(Context.BeforeUnmarshalling.class);
    ArgumentCaptor<Context.BeforeUnmarshalling> beforeUnmarshallingArg = ArgumentCaptor.forClass(Context.BeforeUnmarshalling.class);
    ArgumentCaptor<Context.BeforeUnmarshalling> modifyAsyncHttpResponseContent = ArgumentCaptor.forClass(Context.BeforeUnmarshalling.class);
    ArgumentCaptor<Context.AfterUnmarshalling> afterUnmarshallingArg = ArgumentCaptor.forClass(Context.AfterUnmarshalling.class);
    ArgumentCaptor<Context.AfterExecution> modifyResponseArg = ArgumentCaptor.forClass(Context.AfterExecution.class);
    ArgumentCaptor<Context.AfterExecution> afterExecutionArg = ArgumentCaptor.forClass(Context.AfterExecution.class);

    // Verify methods are called in the right order
    InOrder inOrder = Mockito.inOrder(interceptor);
    inOrder.verify(interceptor).beforeExecution(beforeExecutionArg.capture(), attributes.capture());
    inOrder.verify(interceptor).modifyRequest(modifyRequestArg.capture(), attributes.capture());
    inOrder.verify(interceptor).beforeMarshalling(beforeMarshallingArg.capture(), attributes.capture());
    inOrder.verify(interceptor).afterMarshalling(afterMarshallingArg.capture(), attributes.capture());
    inOrder.verify(interceptor).modifyAsyncHttpContent(modifyHttpContentAsyncArg.capture(), attributes.capture());
    inOrder.verify(interceptor).modifyHttpContent(modifyHttpContentArg.capture(), attributes.capture());
    inOrder.verify(interceptor).modifyHttpRequest(modifyHttpRequestArg.capture(), attributes.capture());
    inOrder.verify(interceptor).beforeTransmission(beforeTransmissionArg.capture(), attributes.capture());
    inOrder.verify(interceptor).afterTransmission(afterTransmissionArg.capture(), attributes.capture());
    inOrder.verify(interceptor).modifyHttpResponse(modifyHttpResponseArg.capture(), attributes.capture());
    inOrder.verify(interceptor).modifyHttpResponseContent(modifyHttpResponseContentArg.capture(), attributes.capture());
    inOrder.verify(interceptor).beforeUnmarshalling(beforeUnmarshallingArg.capture(), attributes.capture());
    if (isAsync) {
        inOrder.verify(interceptor).modifyAsyncHttpResponseContent(modifyAsyncHttpResponseContent.capture(), attributes.capture());
    }
    inOrder.verify(interceptor).afterUnmarshalling(afterUnmarshallingArg.capture(), attributes.capture());
    inOrder.verify(interceptor).modifyResponse(modifyResponseArg.capture(), attributes.capture());
    inOrder.verify(interceptor).afterExecution(afterExecutionArg.capture(), attributes.capture());
    if (expectedException != null) {
        ArgumentCaptor<Context.FailedExecution> failedExecutionArg = ArgumentCaptor.forClass(Context.FailedExecution.class);
        inOrder.verify(interceptor).modifyException(failedExecutionArg.capture(), attributes.capture());
        inOrder.verify(interceptor).onExecutionFailure(failedExecutionArg.capture(), attributes.capture());
        verifyFailedExecutionMethodCalled(failedExecutionArg, true);
        assertThat(failedExecutionArg.getValue().exception()).isEqualTo(expectedException);
    }
    verifyNoMoreInteractions(interceptor);

    // Verify beforeExecution gets untouched request
    assertThat(beforeExecutionArg.getValue().request()).isSameAs(inputRequest);

    // Verify methods were given correct parameters
    validateArgs(beforeExecutionArg.getValue(), null);
    validateArgs(modifyRequestArg.getValue(), null);
    validateArgs(beforeMarshallingArg.getValue(), "1");
    validateArgs(afterMarshallingArg.getValue(), "1", null);
    validateArgs(modifyHttpRequestArg.getValue(), "1", null);
    validateArgs(beforeTransmissionArg.getValue(), "1", "2");
    validateArgs(afterTransmissionArg.getValue(), "1", "2", null);
    validateArgs(modifyHttpResponseArg.getValue(), "1", "2", null);
    validateArgs(beforeUnmarshallingArg.getValue(), "1", "2", "3");
    validateArgs(afterUnmarshallingArg.getValue(), "1", "2", "3", null);
    validateArgs(modifyResponseArg.getValue(), "1", "2", "3", null);
    validateArgs(afterExecutionArg.getValue(), "1", "2", "3", "4");

    // Verify same execution attributes were used for all method calls
    assertThat(attributes.getAllValues()).containsOnly(attributes.getAllValues().get(0));
}
 

@Override
public void beforeTransmission(Context.BeforeTransmission context, ExecutionAttributes executionAttributes) {
    wait(beforeTransmissionWait);
}
 

@Override
public void beforeTransmission(Context.BeforeTransmission context, ExecutionAttributes executionAttributes) {
    List<String> contentTypeHeader = context.httpRequest().headers().get(CONTENT_TYPE);
    assertThat(contentTypeHeader.size()).isEqualTo(1);
    assertThat(contentTypeHeader.get(0)).isEqualTo(Mimetype.MIMETYPE_EVENT_STREAM);
}
 

@Override
public void beforeTransmission(Context.BeforeTransmission context, ExecutionAttributes executionAttributes) {
    executionAttributes.putAttributeIfAbsent(START_TIME, Instant.now());
}