1 /*
2  * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
3  *
4  * Licensed under the Apache License, Version 2.0 (the "License").
5  * You may not use this file except in compliance with the License.
6  * A copy of the License is located at
7  *
8  *  http://aws.amazon.com/apache2.0
9  *
10  * or in the "license" file accompanying this file. This file is distributed
11  * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
12  * express or implied. See the License for the specific language governing
13  * permissions and limitations under the License.
14  */
15
16 package software.amazon.awssdk.services.s3.internal.handlers;
17
18 import static software.amazon.awssdk.core.ClientType.SYNC;
19 import static software.amazon.awssdk.services.s3.checksums.ChecksumConstant.CONTENT_LENGTH_HEADER;
20 import static software.amazon.awssdk.services.s3.checksums.ChecksumsEnabledValidator.CHECKSUM;
21 import static software.amazon.awssdk.services.s3.checksums.ChecksumsEnabledValidator.getObjectChecksumEnabledPerResponse;
22 import static software.amazon.awssdk.services.s3.checksums.ChecksumsEnabledValidator.responseChecksumIsValid;
23 import static software.amazon.awssdk.services.s3.checksums.ChecksumsEnabledValidator.shouldRecordChecksum;
24 import static software.amazon.awssdk.services.s3.checksums.ChecksumsEnabledValidator.validatePutObjectChecksum;
25 import static software.amazon.awssdk.utils.FunctionalUtils.invokeSafely;
26
27 import java.io.InputStream;
28 import java.util.Optional;
29 import software.amazon.awssdk.annotations.SdkInternalApi;
30 import software.amazon.awssdk.core.checksums.Md5Checksum;
31 import software.amazon.awssdk.core.checksums.SdkChecksum;
32 import software.amazon.awssdk.core.interceptor.Context;
33 import software.amazon.awssdk.core.interceptor.ExecutionAttribute;
34 import software.amazon.awssdk.core.interceptor.ExecutionAttributes;
35 import software.amazon.awssdk.core.interceptor.ExecutionInterceptor;
36 import software.amazon.awssdk.core.sync.RequestBody;
37 import software.amazon.awssdk.http.ContentStreamProvider;
38 import software.amazon.awssdk.services.s3.checksums.ChecksumCalculatingInputStream;
39 import software.amazon.awssdk.services.s3.checksums.ChecksumValidatingInputStream;
40 import software.amazon.awssdk.services.s3.model.PutObjectResponse;
41
42 @SdkInternalApi
43 public final class SyncChecksumValidationInterceptor implements ExecutionInterceptor {
44     private static ExecutionAttribute<Boolean> SYNC_RECORDING_CHECKSUM = new ExecutionAttribute<>("syncRecordingChecksum");
45
46     @Override
47     public Optional<RequestBody> modifyHttpContent(Context.ModifyHttpRequest context,
48                                                    ExecutionAttributes executionAttributes) {
49         if (shouldRecordChecksum(context.request(), SYNC, executionAttributes, context.httpRequest())
50             && context.requestBody().isPresent()) {
51             SdkChecksum checksum = new Md5Checksum();
52             executionAttributes.putAttribute(CHECKSUM, checksum);
53             executionAttributes.putAttribute(SYNC_RECORDING_CHECKSUM, true);
54
55             RequestBody requestBody = context.requestBody().get();
56
57             ChecksumCalculatingStreamProvider streamProvider =
58                 new ChecksumCalculatingStreamProvider(requestBody.contentStreamProvider(), checksum);
59
60             return Optional.of(RequestBody.fromContentProvider(streamProvider,
61                                                                requestBody.contentLength(),
62                                                                requestBody.contentType()));
63         }
64
65         return context.requestBody();
66     }
67
68     @Override
69     public Optional<InputStream> modifyHttpResponseContent(Context.ModifyHttpResponse context,
70                                                            ExecutionAttributes executionAttributes) {
71         if (getObjectChecksumEnabledPerResponse(context.request(), context.httpResponse())
72             && context.responseBody().isPresent()) {
73
74             SdkChecksum checksum = new Md5Checksum();
75
76             long contentLength = context.httpResponse()
77                                         .firstMatchingHeader(CONTENT_LENGTH_HEADER)
78                                         .map(Long::parseLong)
79                                         .orElse(0L);
80
81             if (contentLength > 0) {
82                 return Optional.of(new ChecksumValidatingInputStream(context.responseBody().get(), checksum, contentLength));
83             }
84         }
85
86         return context.responseBody();
87     }
88
89     @Override
90     public void afterUnmarshalling(Context.AfterUnmarshalling context, ExecutionAttributes executionAttributes) {
91         boolean recordingChecksum = Boolean.TRUE.equals(executionAttributes.getAttribute(SYNC_RECORDING_CHECKSUM));
92         boolean responseChecksumIsValid = responseChecksumIsValid(context.httpResponse());
93
94         if (recordingChecksum && responseChecksumIsValid) {
95             validatePutObjectChecksum((PutObjectResponse) context.response(), executionAttributes);
96         }
97     }
98
99     static final class ChecksumCalculatingStreamProvider implements ContentStreamProvider {
100         private final SdkChecksum checksum;
101         private InputStream currentStream;
102         private final ContentStreamProvider underlyingInputStreamProvider;
103
104         ChecksumCalculatingStreamProvider(ContentStreamProvider underlyingInputStreamProvider, SdkChecksum checksum) {
105             this.underlyingInputStreamProvider = underlyingInputStreamProvider;
106             this.checksum = checksum;
107         }
108
109         @Override
110         public InputStream newStream() {
111             closeCurrentStream();
112             currentStream = invokeSafely(() -> new ChecksumCalculatingInputStream(underlyingInputStreamProvider.newStream(),
113                                                                                   checksum));
114             return currentStream;
115         }
116
117         private void closeCurrentStream() {
118             checksum.reset();
119             if (currentStream != null) {
120                 invokeSafely(currentStream::close);
121                 currentStream = null;
122             }
123         }
124     }
125
126 }
127