Skip to content

Commit

Permalink
Add flexible checksum support for chunked-encoding
Browse files Browse the repository at this point in the history
  • Loading branch information
@haydenbaker
haydenbaker committed Aug 31, 2023
1 parent eb70caa commit f5a1a44
Show file tree
Hide file tree
Showing 5 changed files with 274 additions and 35 deletions.
10 changes: 10 additions & 0 deletions ...tware/amazon/awssdk/http/auth/aws/internal/chunkedencoding/ChunkedEncodedInputStream.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -201,6 +201,11 @@ public static Builder builder() {

public interface Builder {

/**
* Get the current backing input stream configured for the builder.
*/
InputStream inputStream();

/**
* Set the backing input stream.
*/
Expand Down Expand Up @@ -255,6 +260,11 @@ private static class BuilderImpl implements Builder {
private final List<ChunkExtensionProvider> extensions = new ArrayList<>();
private final List<TrailerProvider> trailers = new ArrayList<>();

@Override
public InputStream inputStream() {
return this.inputStream;
}

@Override
public Builder inputStream(InputStream inputStream) {
this.inputStream = inputStream;
Expand Down
110 changes: 95 additions & 15 deletions .../java/software/amazon/awssdk/http/auth/aws/internal/signer/AwsChunkedV4PayloadSigner.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -15,12 +15,16 @@

package software.amazon.awssdk.http.auth.aws.internal.signer;

import static software.amazon.awssdk.http.auth.aws.internal.util.ChecksumUtil.checksumHeaderName;
import static software.amazon.awssdk.http.auth.aws.internal.util.ChecksumUtil.fromChecksumAlgorithm;
import static software.amazon.awssdk.http.auth.aws.signer.V4CanonicalRequest.getCanonicalHeaders;
import static software.amazon.awssdk.http.auth.aws.signer.V4CanonicalRequest.getCanonicalHeadersString;
import static software.amazon.awssdk.http.auth.aws.util.SignerConstant.STREAMING_SIGNED_PAYLOAD;
import static software.amazon.awssdk.http.auth.aws.util.SignerConstant.STREAMING_SIGNED_PAYLOAD_TRAILER;
import static software.amazon.awssdk.http.auth.aws.util.SignerConstant.STREAMING_UNSIGNED_PAYLOAD_TRAILER;
import static software.amazon.awssdk.http.auth.aws.util.SignerConstant.X_AMZ_CONTENT_SHA256;
import static software.amazon.awssdk.http.auth.aws.util.SignerConstant.X_AMZ_DECODED_CONTENT_LENGTH;
import static software.amazon.awssdk.http.auth.aws.util.SignerConstant.X_AMZ_TRAILER;
import static software.amazon.awssdk.http.auth.aws.util.SignerUtils.hash;
import static software.amazon.awssdk.utils.BinaryUtils.toHex;

Expand All @@ -29,16 +33,20 @@
import java.nio.charset.StandardCharsets;
import java.util.Collections;
import java.util.List;
import java.util.Map;
import java.util.function.Function;
import java.util.function.Supplier;
import java.util.stream.Collectors;
import org.reactivestreams.Publisher;
import software.amazon.awssdk.annotations.SdkInternalApi;
import software.amazon.awssdk.checksums.spi.ChecksumAlgorithm;
import software.amazon.awssdk.http.ContentStreamProvider;
import software.amazon.awssdk.http.Header;
import software.amazon.awssdk.http.SdkHttpRequest;
import software.amazon.awssdk.http.auth.aws.internal.checksums.SdkChecksum;
import software.amazon.awssdk.http.auth.aws.internal.chunkedencoding.ChunkedEncodedInputStream;
import software.amazon.awssdk.http.auth.aws.internal.chunkedencoding.TrailerProvider;
import software.amazon.awssdk.http.auth.aws.internal.io.ChecksumInputStream;
import software.amazon.awssdk.http.auth.aws.signer.CredentialScope;
import software.amazon.awssdk.http.auth.aws.signer.V4Context;
import software.amazon.awssdk.http.auth.aws.signer.V4PayloadSigner;
Expand All @@ -56,10 +64,12 @@ public class AwsChunkedV4PayloadSigner implements V4PayloadSigner {

private final CredentialScope credentialScope;
private final int chunkSize;
private final ChecksumAlgorithm checksumAlgorithm;

public AwsChunkedV4PayloadSigner(CredentialScope credentialScope, int chunkSize) {
public AwsChunkedV4PayloadSigner(CredentialScope credentialScope, int chunkSize, ChecksumAlgorithm checksumAlgorithm) {
this.credentialScope = credentialScope;
this.chunkSize = chunkSize;
this.checksumAlgorithm = checksumAlgorithm;
}

/**
Expand Down Expand Up @@ -96,6 +106,7 @@ public ContentStreamProvider sign(ContentStreamProvider payload, V4Context v4Con
.inputStream(inputStream)
.chunkSize(chunkSize)
.header(chunk -> Integer.toHexString(chunk.length).getBytes(StandardCharsets.UTF_8));
setupPreExistingTrailers(chunkedEncodedInputStreamBuilder, request);

switch (checksum) {
case STREAMING_SIGNED_PAYLOAD: {
Expand All @@ -104,13 +115,17 @@ public ContentStreamProvider sign(ContentStreamProvider payload, V4Context v4Con
break;
}
case STREAMING_UNSIGNED_PAYLOAD_TRAILER:
setupChecksumTrailer(chunkedEncodedInputStreamBuilder);
if (checksumAlgorithm != null) {
setupChecksumTrailer(chunkedEncodedInputStreamBuilder, request);
}
break;
case STREAMING_SIGNED_PAYLOAD_TRAILER: {
RollingSigner rollingSigner = new RollingSigner(v4Context.getSigningKey(), v4Context.getSignature());
setupSigExt(chunkedEncodedInputStreamBuilder, rollingSigner);
if (checksumAlgorithm != null) {
setupChecksumTrailer(chunkedEncodedInputStreamBuilder, request);
}
setupSigTrailer(chunkedEncodedInputStreamBuilder, rollingSigner);
setupChecksumTrailer(chunkedEncodedInputStreamBuilder);
break;
}
default:
Expand All @@ -125,6 +140,12 @@ public Publisher<ByteBuffer> signAsync(Publisher<ByteBuffer> payload, V4Context
throw new UnsupportedOperationException();
}

/**
* Add the chunk signature as a chunk-extension.
* <p>
* An instance of a rolling-signer is required, since each chunk's signature is dependent on the last. The first
* chunk signature is dependent on the request signature ("seed" signature).
*/
private void setupSigExt(ChunkedEncodedInputStream.Builder builder, RollingSigner rollingSigner) {
Function<byte[], String> extTemplate = chunk -> rollingSigner.sign(
previousSignature ->
Expand All @@ -146,19 +167,36 @@ private void setupSigExt(ChunkedEncodedInputStream.Builder builder, RollingSigne
);
}

/**
* Add the trailer signature as a chunk-trailer.
* <p>
* In order for this to work, the instance of the rolling-signer MUST be the same instance used to provide a chunk
* signature chunk-extension. The trailer signature depends on the rolling calculation of all previous chunks.
*/
private void setupSigTrailer(ChunkedEncodedInputStream.Builder builder, RollingSigner rollingSigner) {
List<Pair<String, List<String>>> trailers =
builder.trailers().stream().map(TrailerProvider::get).collect(Collectors.toList());
List<TrailerProvider> trailers = builder.trailers();

Supplier<String> sigSupplier = () -> rollingSigner.sign(
previousSignature ->
String.join("\n",
"AWS4-HMAC-SHA256-TRAILER",
credentialScope.getDatetime(),
credentialScope.scope(),
previousSignature,
toHex(hash(getCanonicalHeadersString(trailers)))
)
previousSignature -> {
// Get the headers by calling get() on each of the trailers
Map<String, List<String>> headers =
trailers.stream().map(TrailerProvider::get).collect(
Collectors.toMap(
Pair::left,
Pair::right
)
);

// build the string-to-sign template for the rolling-signer to sign
return String.join("\n",
"AWS4-HMAC-SHA256-TRAILER",
credentialScope.getDatetime(),
credentialScope.scope(),
previousSignature,
toHex(hash(getCanonicalHeadersString(getCanonicalHeaders(headers))))
);
}

);

builder.addTrailer(
Expand All @@ -169,7 +207,49 @@ private void setupSigTrailer(ChunkedEncodedInputStream.Builder builder, RollingS
);
}

private void setupChecksumTrailer(ChunkedEncodedInputStream.Builder builder) {
// TODO(sra-identity-and-auth): Set up checksumming of chunks and add as a trailer
/**
* Add the checksum as a chunk-trailer and add it to the request's trailer header.
* <p>
* The checksum-algorithm MUST be set if this is called, otherwise it will throw.
*/
private void setupChecksumTrailer(ChunkedEncodedInputStream.Builder builder, SdkHttpRequest.Builder request) {
if (checksumAlgorithm == null) {
throw new IllegalArgumentException("A checksum-algorithm must be configured in order to add a checksum trailer!");
}
SdkChecksum sdkChecksum = fromChecksumAlgorithm(checksumAlgorithm);
ChecksumInputStream checksumInputStream = new ChecksumInputStream(
builder.inputStream(),
Collections.singleton(sdkChecksum)
);
String checksumHeaderName = checksumHeaderName(checksumAlgorithm);

TrailerProvider checksumTrailer = () -> Pair.of(
checksumHeaderName,
Collections.singletonList(sdkChecksum.getChecksum())
);

request.appendHeader(X_AMZ_TRAILER, checksumHeaderName);
builder.inputStream(checksumInputStream).addTrailer(checksumTrailer);
}

/**
* Create chunk-trailers for each pre-existing trailer given in the request.
* <p>
* However, we need to validate that these are valid trailers. Since aws-chunked encoding adds the checksum as a trailer, it
* isn't part of the request headers, but other trailers MUST be present in the request-headers.
*/
private void setupPreExistingTrailers(ChunkedEncodedInputStream.Builder builder, SdkHttpRequest.Builder request) {
List<String> trailerHeaders = request.matchingHeaders(X_AMZ_TRAILER);

for (String header : trailerHeaders) {
List<String> values = request.matchingHeaders(header);
if (values.isEmpty()) {
throw new IllegalArgumentException(header + " must be present in the request headers to be a valid trailer.");
}

// Add the trailer to the aws-chunked stream-builder, and remove it from the request headers
builder.addTrailer(() -> Pair.of(header, values));
request.removeHeader(header);
}
}
}
2 changes: 1 addition & 1 deletion ...ain/java/software/amazon/awssdk/http/auth/aws/internal/signer/DefaultAwsV4HttpSigner.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -157,7 +157,7 @@ private static V4PayloadSigner v4PayloadSigner(
}

if (isChunkEncoding) {
return new AwsChunkedV4PayloadSigner(properties.getCredentialScope(), DEFAULT_CHUNK_SIZE_IN_BYTES);
return new AwsChunkedV4PayloadSigner(properties.getCredentialScope(), DEFAULT_CHUNK_SIZE_IN_BYTES, null);
}

return V4PayloadSigner.create();
Expand Down
Loading

0 comments on commit f5a1a44

Please sign in to comment.