software.amazon.awssdk.core.interceptor.ExecutionAttributes Java Examples

@Setup
public void setup(Blackhole bh) {
    dynamoDb = DynamoDbClient.builder()
            .credentialsProvider(StaticCredentialsProvider.create(
                    AwsBasicCredentials.create("akid", "skid")))
            .httpClient(new MockHttpClient(testItem.responseContent, "{}"))
            .overrideConfiguration(o -> o.addExecutionInterceptor(new ExecutionInterceptor() {
                @Override
                public void afterUnmarshalling(Context.AfterUnmarshalling context,
                                               ExecutionAttributes executionAttributes) {
                    bh.consume(context);
                    bh.consume(executionAttributes);
                }
            }))
            .build();

    DynamoDbEnhancedClient ddbEnh = DynamoDbEnhancedClient.builder()
            .dynamoDbClient(dynamoDb)
            .build();

    table = ddbEnh.table(testItem.name(), testItem.tableSchema);
}
 

/**
 * Validates that checksums should be enabled based on {@link ClientType} and the presence
 * or S3 specific headers.
 *
 * @param expectedClientType - The expected client type for enabling checksums
 * @param executionAttributes - {@link ExecutionAttributes} to determine the actual client type
 * @return If trailing checksums should be enabled for this request.
 */
public static boolean shouldRecordChecksum(SdkRequest sdkRequest,
                                           ClientType expectedClientType,
                                           ExecutionAttributes executionAttributes,
                                           SdkHttpRequest httpRequest) {
    if (!(sdkRequest instanceof PutObjectRequest)) {
        return false;
    }

    ClientType actualClientType = executionAttributes.getAttribute(SdkExecutionAttribute.CLIENT_TYPE);

    if (!expectedClientType.equals(actualClientType)) {
        return false;
    }


    if (hasServerSideEncryptionHeader(httpRequest)) {
        return false;
    }

    return checksumEnabledPerConfig(executionAttributes);
}
 

protected <InputT extends SdkRequest, OutputT extends SdkResponse> ExecutionContext createExecutionContext(
    ClientExecutionParams<InputT, OutputT> params, ExecutionAttributes executionAttributes) {

    SdkRequest originalRequest = params.getInput();

    executionAttributes
        .putAttribute(SdkExecutionAttribute.SERVICE_CONFIG,
                      clientConfiguration.option(SdkClientOption.SERVICE_CONFIGURATION))
        .putAttribute(SdkExecutionAttribute.SERVICE_NAME, clientConfiguration.option(SdkClientOption.SERVICE_NAME));

    ExecutionInterceptorChain interceptorChain =
            new ExecutionInterceptorChain(clientConfiguration.option(SdkClientOption.EXECUTION_INTERCEPTORS));

    return ExecutionContext.builder()
                           .interceptorChain(interceptorChain)
                           .interceptorContext(InterceptorContext.builder()
                                                                 .request(originalRequest)
                                                                 .build())
                           .executionAttributes(executionAttributes)
                           .signer(clientConfiguration.option(SdkAdvancedClientOption.SIGNER))
                           .build();
}
 

@Override
public SdkHttpRequest modifyHttpRequest(Context.ModifyHttpRequest context,
                                        ExecutionAttributes executionAttributes) {
    if (context.request() instanceof CreateMultipartUploadRequest) {
        SdkHttpRequest.Builder builder = context.httpRequest()
                                                .toBuilder()
                                                .putHeader(CONTENT_LENGTH, String.valueOf(0));

        if (!context.httpRequest().firstMatchingHeader(CONTENT_TYPE).isPresent()) {
            builder.putHeader(CONTENT_TYPE, "binary/octet-stream");
        }

        return builder.build();
    }

    return context.httpRequest();
}
 

@Before
public void testSetup() throws Exception {
    mockClient = mock(SdkAsyncHttpClient.class);

    executionParams = new ClientExecutionParams<SdkRequest, SdkResponse>()
            .withInput(request)
            .withMarshaller(marshaller)
            .withResponseHandler(responseHandler)
            .withErrorResponseHandler(errorResponseHandler);

    SdkClientConfiguration config = HttpTestUtils.testClientConfiguration().toBuilder()
            .option(SdkClientOption.ASYNC_HTTP_CLIENT, mockClient)
            .option(SdkClientOption.RETRY_POLICY, RETRY_POLICY)
            .option(SdkAdvancedAsyncClientOption.FUTURE_COMPLETION_EXECUTOR, Runnable::run)
            .build();

    clientHandler = new SdkAsyncClientHandler(config);

    when(request.overrideConfiguration()).thenReturn(Optional.empty());

    when(marshaller.marshall(eq(request))).thenReturn(ValidSdkObjects.sdkHttpFullRequest().build());

    when(responseHandler.handle(any(SdkHttpFullResponse.class), any(ExecutionAttributes.class)))
            .thenReturn(VoidSdkResponse.builder().build());
}
 

/**
 * For us-east-1 there must not be a location constraint (or containing CreateBucketConfiguration) sent.
 */
@Test
public void modifyRequest_UsEast1_UsesNullBucketConfiguration() {
    CreateBucketRequest request = CreateBucketRequest.builder()
                                                     .bucket("test-bucket")
                                                     .createBucketConfiguration(CreateBucketConfiguration.builder()
                                                                                                         .build())
                                                     .build();

    Context.ModifyRequest context = () -> request;

    ExecutionAttributes attributes = new ExecutionAttributes()
            .putAttribute(AwsExecutionAttribute.AWS_REGION, Region.US_EAST_1);

    SdkRequest modifiedRequest = new CreateBucketInterceptor().modifyRequest(context, attributes);
    assertThat(((CreateBucketRequest) modifiedRequest).createBucketConfiguration()).isNull();
}
 

@Override
public SdkHttpRequest modifyHttpRequest(Context.ModifyHttpRequest context, ExecutionAttributes executionAttributes) {
    SdkHttpRequest request = context.httpRequest();
    Object originalRequest = context.request();
    if (originalRequest instanceof PredictRequest) {
        PredictRequest pr = (PredictRequest) originalRequest;
        if (pr.predictEndpoint() == null) {
            throw SdkClientException.builder().message("PredictRequest.PredictEndpoint is required!").build();
        }

        try {
            URI endpoint = new URI(pr.predictEndpoint());
            return request.toBuilder().uri(endpoint).build();
        } catch (URISyntaxException e) {
            throw SdkClientException.builder()
                                    .message("Unable to parse PredictRequest.PredictEndpoint")
                                    .cause(e)
                                    .build();
        }
    }
    return request;
}
 

@Test
public void modifyRequest_UpdatesLocationConstraint_When_NullCreateBucketConfiguration() {
    CreateBucketRequest request = CreateBucketRequest.builder()
                                                     .bucket("test-bucket")
                                                     .build();

    Context.ModifyRequest context = () -> request;

    ExecutionAttributes attributes = new ExecutionAttributes()
            .putAttribute(AwsExecutionAttribute.AWS_REGION, Region.US_EAST_2);

    SdkRequest modifiedRequest = new CreateBucketInterceptor().modifyRequest(context, attributes);
    String locationConstraint = ((CreateBucketRequest) modifiedRequest).createBucketConfiguration().locationConstraintAsString();

    assertThat(locationConstraint).isEqualToIgnoringCase("us-east-2");
}
 

@Override
public void afterExecution(Context.AfterExecution context, ExecutionAttributes executionAttributes) {
    SdkResponse response = context.response();
    SdkRequest originalRequest = context.request();
    if (response != null) {
        if (originalRequest instanceof SendMessageRequest) {
            SendMessageRequest sendMessageRequest = (SendMessageRequest) originalRequest;
            SendMessageResponse sendMessageResult = (SendMessageResponse) response;
            sendMessageOperationMd5Check(sendMessageRequest, sendMessageResult);

        } else if (originalRequest instanceof ReceiveMessageRequest) {
            ReceiveMessageResponse receiveMessageResult = (ReceiveMessageResponse) response;
            receiveMessageResultMd5Check(receiveMessageResult);

        } else if (originalRequest instanceof SendMessageBatchRequest) {
            SendMessageBatchRequest sendMessageBatchRequest = (SendMessageBatchRequest) originalRequest;
            SendMessageBatchResponse sendMessageBatchResult = (SendMessageBatchResponse) response;
            sendMessageBatchOperationMd5Check(sendMessageBatchRequest, sendMessageBatchResult);
        }
    }
}
 

@Setup
public void setup(Blackhole bh) {
    ddb = DynamoDbClient.builder()
            .credentialsProvider(StaticCredentialsProvider.create(
                    AwsBasicCredentials.create("akid", "skid")))
            .httpClient(new MockHttpClient("{}", "{}"))
            .overrideConfiguration(c -> c.addExecutionInterceptor(new ExecutionInterceptor() {
                @Override
                public void afterUnmarshalling(Context.AfterUnmarshalling context,
                                               ExecutionAttributes executionAttributes) {
                    bh.consume(context);
                    bh.consume(executionAttributes);
                }
            }))
            .build();

    DynamoDbEnhancedClient ddbEnh = DynamoDbEnhancedClient.builder()
            .dynamoDbClient(ddb)
            .build();

    enhTable = ddbEnh.table(testItem.name(), testItem.tableSchema);
}
 

@Bean
@ConditionalOnProperty(prefix = "hedera.mirror.importer.downloader", name = "cloudProvider", havingValue = "GCP")
public S3AsyncClient gcpCloudStorageClient() {
    log.info("Configured to download from GCP with bucket name '{}'", downloaderProperties.getBucketName());
    // Any valid region for aws client. Ignored by GCP.
    S3AsyncClientBuilder clientBuilder = asyncClientBuilder("us-east-1")
            .endpointOverride(URI.create(downloaderProperties.getCloudProvider().getEndpoint()));
    String projectId = downloaderProperties.getGcpProjectId();
    if (StringUtils.isNotBlank(projectId)) {
        clientBuilder.overrideConfiguration(builder -> builder.addExecutionInterceptor(new ExecutionInterceptor() {
            @Override
            public SdkHttpRequest modifyHttpRequest(
                    Context.ModifyHttpRequest context, ExecutionAttributes executionAttributes) {
                return context.httpRequest().toBuilder()
                        .appendRawQueryParameter("userProject", projectId).build();
            }
        }));
    }
    return clientBuilder.build();
}
 

/**
 * Handle the event stream message according to it's type.
 *
 * @param m Decoded message.
 */
private void handleMessage(Message m) {
    try {
        if (isEvent(m)) {
            if (m.getHeaders().get(":event-type").getString().equals("initial-response")) {
                eventStreamResponseHandler.responseReceived(
                    initialResponseHandler.handle(adaptMessageToResponse(m, false),
                                                  EMPTY_EXECUTION_ATTRIBUTES));
            } else {
                // Add to queue to be delivered later by the executor
                eventsToDeliver.add(eventResponseHandler.handle(adaptMessageToResponse(m, false),
                                                                EMPTY_EXECUTION_ATTRIBUTES));
            }
        } else if (isError(m) || isException(m)) {
            SdkHttpFullResponse errorResponse = adaptMessageToResponse(m, true);
            Throwable exception = exceptionResponseHandler.handle(
                errorResponse, new ExecutionAttributes().putAttribute(SdkExecutionAttribute.SERVICE_NAME, serviceName));
            runAndLogError(log, "Error thrown from exceptionOccurred, ignoring.", () -> exceptionOccurred(exception));
        }
    } catch (Exception e) {
        throw SdkClientException.builder().cause(e).build();
    }
}
 

@Override
public SdkHttpRequest modifyHttpRequest(Context.ModifyHttpRequest context,
                                        ExecutionAttributes executionAttributes) {
    SdkHttpRequest request = context.httpRequest();

    if (!request.headers().containsKey(X_AMZ_ACCOUNT_ID)) {
        throw SdkClientException.create("Account ID must be specified for all requests");
    }

    String accountId = request.headers().get(X_AMZ_ACCOUNT_ID).get(0);

    S3ControlConfiguration config = (S3ControlConfiguration) executionAttributes.getAttribute(
        AwsSignerExecutionAttribute.SERVICE_CONFIG);

    String host = resolveHost(request, accountId, config);

    return request.toBuilder()
                  .host(host)
                  .build();
}
 

@Override
public SdkResponse modifyResponse(Context.ModifyResponse context, ExecutionAttributes executionAttributes) {
    return ((GetObjectResponse) context.response())
            .toBuilder()
            .metadata(ImmutableMap.of("x-amz-assert", "injected-value"))
            .build();
}
 

@Override
public Optional<Publisher<ByteBuffer>> modifyAsyncHttpResponseContent(Context.ModifyHttpResponse context,
                                                                      ExecutionAttributes executionAttributes) {
    if (INTERCEPTOR_CONTEXT_PREDICATE.test(context)) {
        return context.responsePublisher().map(
            body -> SdkPublishers.envelopeWrappedPublisher(body, XML_ENVELOPE_PREFIX, XML_ENVELOPE_SUFFIX));
    }

    return context.responsePublisher();
}
 

@Test
public void testUnmarshalling() throws Exception {
    HttpResponseHandler<SdkPojo> responseHandler = protocolFactory
            .createResponseHandler(JsonOperationMetadata.builder().build(), InputEvent::builder);

    InputEvent unmarshalled = (InputEvent) responseHandler.handle(SdkHttpFullResponse.builder()
                    .content(AbortableInputStream.create(SdkBytes.fromUtf8String(EXPLICIT_PAYLOAD).asInputStream()))
                    .putHeader(HEADER_MEMBER_NAME, HEADER_MEMBER)
                    .build(),
            new ExecutionAttributes());

    assertThat(unmarshalled.headerMember()).isEqualTo(HEADER_MEMBER);
    assertThat(unmarshalled.explicitPayloadMember().asUtf8String()).isEqualTo(EXPLICIT_PAYLOAD);
}
 

@Test
public void modifyHttpRequest_nonGetObjectRequest_shouldNotModifyHttpRequest() {
    Context.ModifyHttpRequest modifyHttpRequestContext =
        modifyHttpRequestContext(PutObjectRequest.builder().build());

    SdkHttpRequest sdkHttpRequest = interceptor.modifyHttpRequest(modifyHttpRequestContext, new ExecutionAttributes());
    assertThat(sdkHttpRequest).isEqualToComparingFieldByField(modifyHttpRequestContext.httpRequest());
}
 

@Test
public void successReleasesAcquiredCost() {
    ExecutionAttributes attributes = new ExecutionAttributes();

    TokenBucketRetryCondition condition = create(20, e -> 10);

    assertThat(condition.shouldRetry(context(EXCEPTION, attributes))).isTrue();
    assertThat(condition.tokensAvailable()).isEqualTo(10);

    condition.requestSucceeded(context(EXCEPTION, attributes));
    assertThat(condition.tokensAvailable()).isEqualTo(20);
}
 

@Test
public void putObjectChecksumEnabled_disabledFromConfig_false() {
    ExecutionAttributes executionAttributes = getExecutionAttributesWithChecksumDisabled();

    assertThat(shouldRecordChecksum(PutObjectRequest.builder().build(),
                                    ClientType.SYNC,
                                    executionAttributes,
                                    emptyHttpRequest().build())).isFalse();
}
 

@Test
public void modifyResponse_getObjectResponseNoChecksumHeader_shouldNotModifyResponse() {
    long contentLength = 50;
    GetObjectResponse response = GetObjectResponse.builder().contentLength(contentLength).build();
    Context.ModifyResponse modifyResponseContext = modifyResponseContext(
        GetObjectRequest.builder().build(),
        response,
        SdkHttpFullResponse.builder().build());

    GetObjectResponse actualResponse = (GetObjectResponse) interceptor.modifyResponse(modifyResponseContext,
                                                                                      new ExecutionAttributes());
    assertThat(actualResponse).isEqualTo(response);
}
 

@Test
public void modifyHttpResponseContent_nonGetObjectRequest_shouldNotModify() {
    ExecutionAttributes executionAttributes = getExecutionAttributes();
    SdkHttpResponse sdkHttpResponse = getSdkHttpResponseWithChecksumHeader();
    sdkHttpResponse.toBuilder().clearHeaders();
    Context.ModifyHttpResponse modifyHttpResponse =
        InterceptorTestUtils.modifyHttpResponse(PutObjectRequest.builder().build(), sdkHttpResponse);
    Optional<InputStream> publisher = interceptor.modifyHttpResponseContent(modifyHttpResponse,
                                                                            executionAttributes);
    assertThat(publisher).isEqualTo(modifyHttpResponse.responseBody());

}
 

@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");
}
 

private AwsS3V4SignerParams constructAwsS3SignerParams(ExecutionAttributes executionAttributes) {
    AwsS3V4SignerParams.Builder signerParams = extractSignerParams(AwsS3V4SignerParams.builder(),
                                                                         executionAttributes);

    Optional.ofNullable(executionAttributes.getAttribute(S3SignerExecutionAttribute.ENABLE_CHUNKED_ENCODING))
            .ifPresent(signerParams::enableChunkedEncoding);

    Optional.ofNullable(executionAttributes.getAttribute(S3SignerExecutionAttribute.ENABLE_PAYLOAD_SIGNING))
            .ifPresent(signerParams::enablePayloadSigning);

    return signerParams.build();
}
 

/**
 * Retrieve all global partitions for the AWS service that we're interacting with.
 */
private List<ServicePartitionMetadata> globalPartitionsForService(ExecutionAttributes executionAttributes) {
    return ServiceMetadata.of(executionAttributes.getAttribute(AwsExecutionAttribute.ENDPOINT_PREFIX))
                          .servicePartitions()
                          .stream()
                          .filter(sp -> sp.globalRegion().isPresent())
                          .collect(Collectors.toList());
}
 

@Test
public void modifyHttpContent_putObjectRequest_checksumDisabled_shouldNotModify() {
    ExecutionAttributes executionAttributes = getExecutionAttributesWithChecksumDisabled();
    Context.ModifyHttpRequest modifyHttpRequest =
        InterceptorTestUtils.modifyHttpRequestContext(GetObjectRequest.builder().build());

    Optional<RequestBody> requestBody = interceptor.modifyHttpContent(modifyHttpRequest,
                                                                      executionAttributes);
    assertThat(requestBody).isEqualTo(modifyHttpRequest.requestBody());
}
 

private void verifyAccelerateDisabledOperationsEndpointNotConverted(SdkRequest request) {
    URI customUri = URI.create("http://s3-test.com");
    Context.ModifyHttpRequest ctx = context(request, sdkHttpRequest(customUri));
    ExecutionAttributes executionAttributes = new ExecutionAttributes();
    S3Configuration s3Configuration = S3Configuration.builder().accelerateModeEnabled(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(customUri);
}
 

@Test
public void errorResponseShouldNotModifyResponse() {
    GetBucketPolicyRequest request = GetBucketPolicyRequest.builder().build();
    Context.ModifyHttpResponse context = modifyHttpResponseContent(request, SdkHttpResponse.builder().statusCode(404).build());
    Optional<InputStream> inputStream = interceptor.modifyHttpResponseContent(context, new ExecutionAttributes());
    assertThat(inputStream).isEqualTo(context.responseBody());
}
 

@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");
}
 

/**
 * If there is a retry, then {@code beforeTransmission} will be called with the previous
 * attributes. This method will look for an existing entry and write out the log message.
 * The log entry may not have been filled in with a status if no response was received,
 * e.g., a connection exception. Since we do not have access to the failure, the status
 * will get set to {@code unexpected_error}.
 */
private void logRetryAttempt(ExecutionAttributes attrs) {
  IpcLogEntry logEntry = attrs.getAttribute(LOG_ENTRY);
  if (logEntry != null) {
    if (!isStatusSet(attrs)) {
      logEntry.markEnd().withStatus(IpcStatus.unexpected_error);
    }
    logEntry.log();
  }
}