software.amazon.awssdk.services.s3.checksums.ChecksumValidatingPublisher Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of s3 Show documentation
The AWS Java SDK for Amazon S3 module holds the client classes that are used for communicating with Amazon Simple Storage Service
There is a newer version: 2.31.2
Show newest version
/*
 * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License").
 * You may not use this file except in compliance with the License.
 * A copy of the License is located at
 *
 *  http://aws.amazon.com/apache2.0
 *
 * or in the "license" file accompanying this file. This file is distributed
 * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
 * express or implied. See the License for the specific language governing
 * permissions and limitations under the License.
 */

package software.amazon.awssdk.services.s3.checksums;

import static java.lang.Math.toIntExact;

import java.nio.ByteBuffer;
import java.util.Arrays;
import org.reactivestreams.Publisher;
import org.reactivestreams.Subscriber;
import org.reactivestreams.Subscription;
import software.amazon.awssdk.annotations.SdkInternalApi;
import software.amazon.awssdk.core.async.SdkPublisher;
import software.amazon.awssdk.core.checksums.SdkChecksum;
import software.amazon.awssdk.core.exception.RetryableException;
import software.amazon.awssdk.utils.BinaryUtils;

@SdkInternalApi
public final class ChecksumValidatingPublisher implements SdkPublisher {

    private final Publisher publisher;
    private final SdkChecksum sdkChecksum;
    private final long contentLength;

    public ChecksumValidatingPublisher(Publisher publisher,
                                       SdkChecksum sdkChecksum,
                                       long contentLength) {
        this.publisher = publisher;
        this.sdkChecksum = sdkChecksum;
        this.contentLength = contentLength;
    }

    @Override
    public void subscribe(Subscriber s) {
        if (contentLength > 0) {
            publisher.subscribe(new ChecksumValidatingSubscriber(s, sdkChecksum, contentLength));
        } else {
            publisher.subscribe(new ChecksumSkippingSubscriber(s));
        }
    }

    private static class ChecksumValidatingSubscriber implements Subscriber {

        private static final int CHECKSUM_SIZE = 16;

        private final Subscriber wrapped;
        private final SdkChecksum sdkChecksum;
        private final long strippedLength;

        private byte[] streamChecksum = new byte[CHECKSUM_SIZE];
        private long lengthRead = 0;

        ChecksumValidatingSubscriber(Subscriber wrapped,
                                     SdkChecksum sdkChecksum,
                                     long contentLength) {
            this.wrapped = wrapped;
            this.sdkChecksum = sdkChecksum;
            this.strippedLength = contentLength - CHECKSUM_SIZE;
        }

        @Override
        public void onSubscribe(Subscription s) {
            wrapped.onSubscribe(s);
        }

        @Override
        public void onNext(ByteBuffer byteBuffer) {
            byte[] buf = BinaryUtils.copyBytesFrom(byteBuffer);

            if (lengthRead < strippedLength) {
                int toUpdate = (int) Math.min(strippedLength - lengthRead, buf.length);

                sdkChecksum.update(buf, 0, toUpdate);
            }
            lengthRead += buf.length;

            if (lengthRead >= strippedLength) {
                // Incoming buffer contains at least a bit of the checksum
                // Code below covers both cases of the incoming buffer relative to checksum border
                // a) buffer starts before checksum border and extends into checksum
                //      |<------ data ------->|<--cksum-->|   <--- original data
                //                       |<---buffer--->|     <--- incoming buffer
                //                            |<------->|     <--- checksum bytes so far
                //                       |<-->|               <--- bufChecksumOffset
                //                            |               <--- streamChecksumOffset
                // b) buffer starts at or after checksum border
                //      |<------ data ------->|<--cksum-->|   <--- original data
                //                                |<-->|      <--- incoming buffer
                //                            |<------>|      <--- checksum bytes so far
                //                                |           <--- bufChecksumOffset
                //                            |<->|           <--- streamChecksumOffset
                int cksumBytesSoFar = toIntExact(lengthRead - strippedLength);
                int bufChecksumOffset = (buf.length > cksumBytesSoFar) ? (buf.length - cksumBytesSoFar) : 0;
                int streamChecksumOffset = (buf.length > cksumBytesSoFar) ? 0 : (cksumBytesSoFar - buf.length);
                int cksumBytes = Math.min(cksumBytesSoFar, buf.length);
                System.arraycopy(buf, bufChecksumOffset, streamChecksum, streamChecksumOffset, cksumBytes);
                if (buf.length > cksumBytesSoFar) {
                    wrapped.onNext(ByteBuffer.wrap(Arrays.copyOfRange(buf, 0, buf.length - cksumBytesSoFar)));
                } else {
                    // Always be sure to satisfy the wrapped publisher's demand.
                    wrapped.onNext(ByteBuffer.allocate(0));

                    // TODO: The most efficient implementation would request more from the upstream publisher instead of relying
                    //  on the downstream publisher to do that, but that's much more complicated: it requires tracking
                    //  outstanding demand from the downstream publisher. Long-term we should migrate to an RxJava publisher
                    //  implementation to reduce how error-prone our publisher implementations are.
                }
            } else {
                // Incoming buffer totally excludes the checksum
                wrapped.onNext(byteBuffer);
            }
        }

        @Override
        public void onError(Throwable t) {
            wrapped.onError(t);
        }

        @Override
        public void onComplete() {
            if (strippedLength > 0) {
                byte[] computedChecksum = sdkChecksum.getChecksumBytes();
                if (!Arrays.equals(computedChecksum, streamChecksum)) {
                    onError(RetryableException.create(
                        String.format("Data read has a different checksum than expected. Was 0x%s, but expected 0x%s. "
                                      + "Common causes: (1) You modified a request ByteBuffer before it could be "
                                      + "written to the service. Please ensure your data source does not modify the "
                                      + " byte buffers after you pass them to the SDK. (2) The data was corrupted between the "
                                      + "client and service. Note: Despite this error, the upload still completed and was "
                                      + "persisted in S3.",
                                      BinaryUtils.toHex(computedChecksum), BinaryUtils.toHex(streamChecksum))));
                    return; // Return after onError and not call onComplete below
                }
            }
            wrapped.onComplete();
        }
    }

    private static class ChecksumSkippingSubscriber implements Subscriber {
        private static final int CHECKSUM_SIZE = 16;

        private final Subscriber wrapped;

        ChecksumSkippingSubscriber(Subscriber wrapped) {
            this.wrapped = wrapped;
        }

        @Override
        public void onSubscribe(Subscription s) {
            wrapped.onSubscribe(s);
        }

        @Override
        public void onNext(ByteBuffer byteBuffer) {
            byte[] buf = BinaryUtils.copyBytesFrom(byteBuffer);
            wrapped.onNext(ByteBuffer.wrap(Arrays.copyOfRange(buf, 0, buf.length - CHECKSUM_SIZE)));
        }

        @Override
        public void onError(Throwable t) {
            wrapped.onError(t);
        }

        @Override
        public void onComplete() {
            wrapped.onComplete();
        }
    }

}