Java Code Examples for software.amazon.awssdk.core.interceptor.ExecutionAttributes#putAttribute()

/**
 * Before an individual http request is finalized. This is only called once per operation, meaning
 * we can only have one span per operation.
 */
@Override public SdkHttpRequest modifyHttpRequest(
    Context.ModifyHttpRequest context,
    ExecutionAttributes executionAttributes
) {
  HttpClientRequest request = new HttpClientRequest(context.httpRequest());

  Span span = handler.handleSend(request);
  executionAttributes.putAttribute(SPAN, span);

  String serviceName = executionAttributes.getAttribute(SdkExecutionAttribute.SERVICE_NAME);
  String operation = getAwsOperationNameFromRequestClass(context.request());
  // TODO: This overwrites user configuration. We don't do this in other layered tools such
  // as WebMVC. Instead, we add tags (such as we do here) and neither overwrite the name, nor
  // remoteServiceName. Users can always remap in an span handler using tags!
  span.name(operation)
      .remoteServiceName(serviceName)
      .tag("aws.service_name", serviceName)
      .tag("aws.operation", operation);

  return request.build();
}
 

private void verifyVirtualStyleConvertDnsEndpoint(String protocol) {
    String bucketName = "test-bucket";
    String key = "test-key";
    URI customUri = URI.create(String.format("%s://s3-test.com/%s/%s", protocol, bucketName, key));
    URI expectedUri = URI.create(String.format("%s://%s.s3.dualstack.us-east-1.amazonaws.com/%s", protocol,
                                               bucketName, key));

    Context.ModifyHttpRequest ctx = context(ListObjectsV2Request.builder().bucket(bucketName).build(),
                                            sdkHttpRequest(customUri));
    ExecutionAttributes executionAttributes = new ExecutionAttributes();
    S3Configuration s3Configuration = S3Configuration.builder().dualstackEnabled(true).build();

    executionAttributes.putAttribute(SERVICE_CONFIG, s3Configuration);
    executionAttributes.putAttribute(AWS_REGION, Region.US_EAST_1);

    SdkHttpRequest sdkHttpFullRequest = interceptor.modifyHttpRequest(ctx, executionAttributes);

    assertThat(sdkHttpFullRequest.getUri()).isEqualTo(expectedUri);
}
 

@Override
public void beforeExecution(Context.BeforeExecution context, ExecutionAttributes executionAttributes) {
    AWSXRayRecorder recorder = getRecorder();
    Entity origin = recorder.getTraceEntity();

    Subsegment subsegment = recorder.beginSubsegment(executionAttributes.getAttribute(SdkExecutionAttribute.SERVICE_NAME));
    if (subsegment == null) {
        return;
    }
    subsegment.setNamespace(Namespace.AWS.toString());
    subsegment.putAws(EntityDataKeys.AWS.OPERATION_KEY,
                      executionAttributes.getAttribute(SdkExecutionAttribute.OPERATION_NAME));
    Region region = executionAttributes.getAttribute(AwsExecutionAttribute.AWS_REGION);
    if (region != null) {
        subsegment.putAws(EntityDataKeys.AWS.REGION_KEY, region.id());
    }
    subsegment.putAllAws(extractRequestParameters(context, executionAttributes));
    if (accountId != null) {
        subsegment.putAws(EntityDataKeys.AWS.ACCOUNT_ID_SUBSEGMENT_KEY, accountId);
    }

    recorder.setTraceEntity(origin);
    // store the subsegment in the AWS SDK's executionAttributes so it can be accessed across threads
    executionAttributes.putAttribute(entityKey, subsegment);
}
 

@Test
public void copySnapshotRequest_httpsProtocolAddedToEndpoint() {
    SdkHttpFullRequest request = SdkHttpFullRequest.builder()
            .uri(URI.create("https://ec2.us-west-2.amazonaws.com"))
            .method(SdkHttpMethod.POST)
            .build();

    CopySnapshotRequest ec2Request = CopySnapshotRequest.builder()
            .sourceRegion("us-west-2")
            .destinationRegion("us-east-2")
            .build();

    when(mockContext.httpRequest()).thenReturn(request);
    when(mockContext.request()).thenReturn(ec2Request);

    ExecutionAttributes attrs = new ExecutionAttributes();
    attrs.putAttribute(AWS_CREDENTIALS, AwsBasicCredentials.create("foo", "bar"));

    SdkHttpRequest modifiedRequest = INTERCEPTOR.modifyHttpRequest(mockContext, attrs);

    String presignedUrl = modifiedRequest.rawQueryParameters().get("PresignedUrl").get(0);

    assertThat(presignedUrl).startsWith("https://");
}
 

@Override
public Optional<RequestBody> modifyHttpContent(Context.ModifyHttpRequest context,
                                               ExecutionAttributes executionAttributes) {

    if (!BLACKLIST_METHODS.contains(context.request().getClass()) && context.requestBody().isPresent()
        && !context.httpRequest().firstMatchingHeader(CONTENT_MD5).isPresent()) {

        try {
            ByteArrayOutputStream baos = new ByteArrayOutputStream();
            IoUtils.copy(context.requestBody().get().contentStreamProvider().newStream(), baos);
            executionAttributes.putAttribute(CONTENT_MD5_ATTRIBUTE, Md5Utils.md5AsBase64(baos.toByteArray()));
            return context.requestBody();
        } catch (IOException e) {
            throw new UncheckedIOException(e);
        }
    }

    return context.requestBody();
}
 

private void verifyEndpoint(String protocol, String expectedEndpoint,
                            S3Configuration.Builder builder) {
    String bucket = "test-bucket";
    String key = "test-key";
    URI customUri = URI.create(String.format("%s://s3-test.com/%s/%s", protocol, bucket, key));
    URI expectedUri = URI.create(String.format("%s/%s/%s", expectedEndpoint, bucket, key));
    Context.ModifyHttpRequest ctx = context(PutObjectRequest.builder().build(), sdkHttpRequest(customUri));
    ExecutionAttributes executionAttributes = new ExecutionAttributes();
    S3Configuration s3Configuration = builder.build();

    executionAttributes.putAttribute(SERVICE_CONFIG, s3Configuration);
    executionAttributes.putAttribute(AWS_REGION, Region.US_EAST_1);

    SdkHttpRequest sdkHttpFullRequest = interceptor.modifyHttpRequest(ctx, executionAttributes);

    assertThat(sdkHttpFullRequest.getUri()).isEqualTo(expectedUri);
}
 

@Override
public Optional<Publisher<ByteBuffer>> modifyAsyncHttpResponseContent(Context.ModifyHttpResponse context,
                                                                      ExecutionAttributes executionAttributes) {
    if (getObjectChecksumEnabledPerResponse(context.request(), context.httpResponse())
        && context.responsePublisher().isPresent()) {
        long contentLength = context.httpResponse()
                                    .firstMatchingHeader(CONTENT_LENGTH_HEADER)
                                    .map(Long::parseLong)
                                    .orElse(0L);

        SdkChecksum checksum = new Md5Checksum();
        executionAttributes.putAttribute(CHECKSUM, checksum);
        if (contentLength > 0) {
            return Optional.of(new ChecksumValidatingPublisher(context.responsePublisher().get(), checksum, contentLength));
        }
    }

    return context.responsePublisher();
}
 

@Override
public void afterExecution(final AfterExecution context, final ExecutionAttributes executionAttributes) {
  final Span span = executionAttributes.getAttribute(SPAN_ATTRIBUTE);
  if (span == null)
    return;

  executionAttributes.putAttribute(SPAN_ATTRIBUTE, null);
  span.setTag(Tags.HTTP_STATUS, context.httpResponse().statusCode());
  span.finish();
}
 

@Override
public Optional<AsyncRequestBody> modifyAsyncHttpContent(Context.ModifyHttpRequest context,
                                                         ExecutionAttributes executionAttributes) {
    boolean shouldRecordChecksum = shouldRecordChecksum(context.request(), ASYNC, executionAttributes, context.httpRequest());

    if (shouldRecordChecksum && context.asyncRequestBody().isPresent()) {
        SdkChecksum checksum = new Md5Checksum();
        executionAttributes.putAttribute(ASYNC_RECORDING_CHECKSUM, true);
        executionAttributes.putAttribute(CHECKSUM, checksum);
        return Optional.of(new ChecksumCalculatingAsyncRequestBody(context.asyncRequestBody().get(), checksum));
    }

    return context.asyncRequestBody();
}
 

@Test(expected = SdkClientException.class)
public void modifyHttpRequest_ThrowsException_NoAccountId() {
    EndpointAddressInterceptor interceptor = new EndpointAddressInterceptor();

    S3ControlConfiguration controlConfiguration = S3ControlConfiguration.builder()
                                                                        .dualstackEnabled(true)
                                                                        .build();
    ExecutionAttributes executionAttributes = new ExecutionAttributes();
    executionAttributes.putAttribute(SdkExecutionAttribute.SERVICE_CONFIG, controlConfiguration);

    interceptor.modifyHttpRequest(new Context(request.toBuilder().removeHeader(X_AMZ_ACCOUNT_ID).build()),
                                  executionAttributes);
}
 

@Test(expected = SdkClientException.class)
public void modifyHttpRequest_ThrowsException_NonStandardEndpoint() {
    EndpointAddressInterceptor interceptor = new EndpointAddressInterceptor();

    S3ControlConfiguration controlConfiguration = S3ControlConfiguration.builder()
                                                                        .dualstackEnabled(true)
                                                                        .build();
    ExecutionAttributes executionAttributes = new ExecutionAttributes();
    executionAttributes.putAttribute(SdkExecutionAttribute.SERVICE_CONFIG, controlConfiguration);

    interceptor.modifyHttpRequest(new Context(request.toBuilder().host("some-garbage").build()),
                                  executionAttributes);
}
 

@Test(expected = SdkClientException.class)
public void modifyHttpRequest_ThrowsException_FipsAndDualstack() {
    EndpointAddressInterceptor interceptor = new EndpointAddressInterceptor();

    S3ControlConfiguration controlConfiguration = S3ControlConfiguration.builder()
                                                                        .fipsModeEnabled(true)
                                                                        .dualstackEnabled(true)
                                                                        .build();
    ExecutionAttributes executionAttributes = new ExecutionAttributes();
    executionAttributes.putAttribute(SdkExecutionAttribute.SERVICE_CONFIG, controlConfiguration);

    interceptor.modifyHttpRequest(new Context(request), executionAttributes);
}
 

@Test
public void modifyHttpRequest_ResolvesCorrectHost_Fips() {
    EndpointAddressInterceptor interceptor = new EndpointAddressInterceptor();

    S3ControlConfiguration controlConfiguration = S3ControlConfiguration.builder().fipsModeEnabled(true).build();
    ExecutionAttributes executionAttributes = new ExecutionAttributes();
    executionAttributes.putAttribute(SdkExecutionAttribute.SERVICE_CONFIG, controlConfiguration);

    SdkHttpRequest modified = interceptor.modifyHttpRequest(new Context(request), executionAttributes);
    assertThat(modified.host()).isEqualTo(ACCOUNT_ID + ".s3-control-fips.us-east-1.amazonaws.com");
}
 

@Test
public void modifyHttpRequest_ResolvesCorrectHost_Dualstack() {
    EndpointAddressInterceptor interceptor = new EndpointAddressInterceptor();

    S3ControlConfiguration controlConfiguration = S3ControlConfiguration.builder().dualstackEnabled(true).build();
    ExecutionAttributes executionAttributes = new ExecutionAttributes();
    executionAttributes.putAttribute(SdkExecutionAttribute.SERVICE_CONFIG, controlConfiguration);

    SdkHttpRequest modified = interceptor.modifyHttpRequest(new Context(request), executionAttributes);
    assertThat(modified.host()).isEqualTo(ACCOUNT_ID + ".s3-control.dualstack.us-east-1.amazonaws.com");
}
 

@Override
public void beforeExecution(BeforeExecution context, ExecutionAttributes executionAttributes) {
  final Span span = GlobalTracer.get().buildSpan(context.request().getClass().getSimpleName())
    .withTag(Tags.SPAN_KIND.getKey(), Tags.SPAN_KIND_CLIENT)
    .withTag(Tags.PEER_SERVICE, executionAttributes.getAttribute(SdkExecutionAttribute.SERVICE_NAME))
    .withTag(Tags.COMPONENT, COMPONENT_NAME).start();

  executionAttributes.putAttribute(SPAN_ATTRIBUTE, span);
}
 

private ExecutionAttributes getExecutionAttributes() {
    ExecutionAttributes executionAttributes = new ExecutionAttributes();
    executionAttributes.putAttribute(CLIENT_TYPE, SYNC);
    return executionAttributes;
}
 

/**
 * Always use the client level timeOffset.
 */
private void adjustForClockSkew(ExecutionAttributes attributes) {
    attributes.putAttribute(SdkExecutionAttribute.TIME_OFFSET, dependencies.timeOffset());
}
 

private ExecutionAttributes getSyncExecutionAttributes() {
    ExecutionAttributes executionAttributes = new ExecutionAttributes();
    executionAttributes.putAttribute(CLIENT_TYPE, ClientType.SYNC);
    return executionAttributes;
}
 

private ExecutionAttributes getExecutionAttributesWithChecksumDisabled() {
    ExecutionAttributes executionAttributes = getExecutionAttributes();
    executionAttributes.putAttribute(SERVICE_CONFIG, S3Configuration.builder().checksumValidationEnabled(false).build());
    return executionAttributes;
}
 

@Test
public void modifyRequest_DoesNotOverwriteExistingAttributeValue() {

    ExecutionAttributes executionAttributes = new ExecutionAttributes();

    interceptor.modifyRequest(context(PutObjectRequest.builder().build()), executionAttributes);

    boolean newValue = !executionAttributes.getAttribute(ENABLE_CHUNKED_ENCODING);

    executionAttributes.putAttribute(ENABLE_CHUNKED_ENCODING, newValue);

    interceptor.modifyRequest(context(PutObjectRequest.builder().build()), executionAttributes);

    assertThat(executionAttributes.getAttribute(ENABLE_CHUNKED_ENCODING)).isEqualTo(newValue);
}