fix: add strict client side response validation for gRPC chunked resumable uploads (#2527)

* Rename JsonResumableSessionFailureScenario to ResumableSessionFailureScenario. The failure scenarios themselves are not json specific, and the methods which are json specific can have grpc overloads * Add more tests to validate GapicUnbufferedChunkedResumableWritableByteChannel is able to properly detect and handle various success responses from GCS which are not success for the client.

Author: BenWhitehead

google-cloud-storage/src/main/java/com/google/cloud/storage/GapicUnbufferedChunkedResumableWritableByteChannel.java

@@ -23,6 +23,7 @@
 import com.google.api.gax.retrying.ResultRetryAlgorithm;
 import com.google.api.gax.rpc.ApiStreamObserver;
 import com.google.api.gax.rpc.ClientStreamingCallable;
+import com.google.api.gax.rpc.OutOfRangeException;
 import com.google.cloud.storage.ChunkSegmenter.ChunkSegment;
 import com.google.cloud.storage.Conversions.Decoder;
 import com.google.cloud.storage.Crc32cValue.Crc32cLengthKnown;
@@ -41,10 +42,9 @@
 import java.util.ArrayList;
 import java.util.List;
 import java.util.concurrent.ExecutionException;
-import java.util.function.Consumer;
-import java.util.function.LongConsumer;
 import java.util.function.Supplier;
 import org.checkerframework.checker.nullness.qual.NonNull;
+import org.checkerframework.checker.nullness.qual.Nullable;
 
 final class GapicUnbufferedChunkedResumableWritableByteChannel
  implements UnbufferedWritableByteChannel {
@@ -58,30 +58,26 @@ final class GapicUnbufferedChunkedResumableWritableByteChannel
  private final RetryingDependencies deps;
  private final ResultRetryAlgorithm<?> alg;
  private final Supplier<GrpcCallContext> baseContextSupplier;
- private final LongConsumer sizeCallback;
- private final Consumer<WriteObjectResponse> completeCallback;
 
- private boolean open = true;
+ private volatile boolean open = true;
  private boolean finished = false;
 
  GapicUnbufferedChunkedResumableWritableByteChannel(
  SettableApiFuture<WriteObjectResponse> resultFuture,
  @NonNull ChunkSegmenter chunkSegmenter,
  ClientStreamingCallable<WriteObjectRequest, WriteObjectResponse> write,
- ResumableWrite requestFactory,
+ WriteCtx<ResumableWrite> writeCtx,
  RetryingDependencies deps,
  ResultRetryAlgorithm<?> alg,
  Supplier<GrpcCallContext> baseContextSupplier) {
  this.resultFuture = resultFuture;
  this.chunkSegmenter = chunkSegmenter;
  this.write = write;
- this.bucketName = requestFactory.bucketName();
- this.writeCtx = new WriteCtx<>(requestFactory);
+ this.bucketName = writeCtx.getRequestFactory().bucketName();
+ this.writeCtx = writeCtx;
  this.deps = deps;
  this.alg = alg;
  this.baseContextSupplier = baseContextSupplier;
- this.sizeCallback = writeCtx.getConfirmedBytes()::set;
- this.completeCallback = resultFuture::set;
  }
 
  @Override
@@ -106,7 +102,7 @@ public void close() throws IOException {
  if (open && !finished) {
  WriteObjectRequest message = finishMessage(true);
  try {
- flush(ImmutableList.of(message));
+ flush(ImmutableList.of(message), null, true);
  finished = true;
  } catch (RuntimeException e) {
  resultFuture.setException(e);
@@ -122,12 +118,13 @@ private long internalWrite(ByteBuffer[] srcs, int srcsOffset, int srcsLength, bo
  throw new ClosedChannelException();
  }
 
+ long begin = writeCtx.getConfirmedBytes().get();
+ RewindableContent content = RewindableContent.of(srcs, srcsOffset, srcsLength);
  ChunkSegment[] data = chunkSegmenter.segmentBuffers(srcs, srcsOffset, srcsLength);
 
  List<WriteObjectRequest> messages = new ArrayList<>();
 
  boolean first = true;
- int bytesConsumed = 0;
  for (ChunkSegment datum : data) {
  Crc32cLengthKnown crc32c = datum.getCrc32c();
  ByteString b = datum.getB();
@@ -144,8 +141,13 @@ private long internalWrite(ByteBuffer[] srcs, int srcsOffset, int srcsLength, bo
  WriteObjectRequest.Builder builder =
  writeCtx
  .newRequestBuilder()
+ .clearWriteObjectSpec()
+ .clearObjectChecksums()
  .setWriteOffset(offset)
  .setChecksummedData(checksummedData.build());
+ if (!first) {
+ builder.clearUploadId();
+ }
  if (!datum.isOnlyFullBlocks()) {
  builder.setFinishWrite(true);
  if (cumulative != null) {
@@ -155,23 +157,25 @@ private long internalWrite(ByteBuffer[] srcs, int srcsOffset, int srcsLength, bo
  finished = true;
  }
 
- WriteObjectRequest build = possiblyPairDownRequest(builder, first).build();
+ WriteObjectRequest build = builder.build();
  first = false;
  messages.add(build);
- bytesConsumed += contentSize;
  }
  if (finalize && !finished) {
  messages.add(finishMessage(first));
  finished = true;
  }
 
  try {
- flush(messages);
+ flush(messages, content, finalize);
  } catch (RuntimeException e) {
  resultFuture.setException(e);
  throw e;
  }
 
+ long end = writeCtx.getConfirmedBytes().get();
+
+ long bytesConsumed = end - begin;
  return bytesConsumed;
  }
 
@@ -182,14 +186,20 @@ private WriteObjectRequest finishMessage(boolean first) {
 
  WriteObjectRequest.Builder b =
  writeCtx.newRequestBuilder().setFinishWrite(true).setWriteOffset(offset);
+ if (!first) {
+ b.clearUploadId();
+ }
  if (crc32cValue != null) {
  b.setObjectChecksums(ObjectChecksums.newBuilder().setCrc32C(crc32cValue.getValue()).build());
  }
- WriteObjectRequest message = possiblyPairDownRequest(b, first).build();
+ WriteObjectRequest message = b.build();
  return message;
  }
 
- private void flush(@NonNull List<WriteObjectRequest> segments) {
+ private void flush(
+ @NonNull List<WriteObjectRequest> segments,
+ @Nullable RewindableContent content,
+ boolean finalizing) {
  GrpcCallContext internalContext = contextWithBucketName(bucketName, baseContextSupplier.get());
  ClientStreamingCallable<WriteObjectRequest, WriteObjectResponse> callable =
  write.withDefaultCallContext(internalContext);
@@ -198,7 +208,7 @@ private void flush(@NonNull List<WriteObjectRequest> segments) {
  deps,
  alg,
  () -> {
- Observer observer = new Observer(sizeCallback, completeCallback);
+ Observer observer = new Observer(content, finalizing);
  ApiStreamObserver<WriteObjectRequest> write = callable.clientStreamingCall(observer);
 
  for (WriteObjectRequest message : segments) {
@@ -211,81 +221,93 @@ private void flush(@NonNull List<WriteObjectRequest> segments) {
  Decoder.identity());
  }
 
- /**
- * Several fields of a WriteObjectRequest are only allowed on the "first" message sent to gcs,
- * this utility method centralizes the logic necessary to clear those fields for use by subsequent
- * messages.
- */
- private static WriteObjectRequest.Builder possiblyPairDownRequest(
- WriteObjectRequest.Builder b, boolean firstMessageOfStream) {
- if (firstMessageOfStream && b.getWriteOffset() == 0) {
- return b;
- }
-
- if (!firstMessageOfStream) {
- b.clearUploadId();
- }
-
- if (b.getWriteOffset() > 0) {
- b.clearWriteObjectSpec();
- }
-
- if (b.getWriteOffset() > 0 && !b.getFinishWrite()) {
- b.clearObjectChecksums();
- }
- return b;
- }
-
  @VisibleForTesting
  WriteCtx<?> getWriteCtx() {
  return writeCtx;
  }
 
- static class Observer implements ApiStreamObserver<WriteObjectResponse> {
+ class Observer implements ApiStreamObserver<WriteObjectResponse> {
 
- private final LongConsumer sizeCallback;
- private final Consumer<WriteObjectResponse> completeCallback;
+ private final RewindableContent content;
+ private final boolean finalizing;
 
  private final SettableApiFuture<Void> invocationHandle;
  private volatile WriteObjectResponse last;
 
- Observer(LongConsumer sizeCallback, Consumer<WriteObjectResponse> completeCallback) {
- this.sizeCallback = sizeCallback;
- this.completeCallback = completeCallback;
+ Observer(@Nullable RewindableContent content, boolean finalizing) {
+ this.content = content;
+ this.finalizing = finalizing;
  this.invocationHandle = SettableApiFuture.create();
  }
 
  @Override
  public void onNext(WriteObjectResponse value) {
- // incremental update
- if (value.hasPersistedSize()) {
- sizeCallback.accept(value.getPersistedSize());
- } else if (value.hasResource()) {
- sizeCallback.accept(value.getResource().getSize());
- }
  last = value;
  }
 
- /**
- * observed exceptions so far
- *
- * <ol>
- * <li>{@link com.google.api.gax.rpc.OutOfRangeException}
- * <li>{@link com.google.api.gax.rpc.AlreadyExistsException}
- * <li>{@link io.grpc.StatusRuntimeException}
- * </ol>
- */
  @Override
  public void onError(Throwable t) {
- invocationHandle.setException(t);
+ if (t instanceof OutOfRangeException) {
+ OutOfRangeException oore = (OutOfRangeException) t;
+ open = false;
+ invocationHandle.setException(
+ ResumableSessionFailureScenario.SCENARIO_5.toStorageException());
+ } else {
+ invocationHandle.setException(t);
+ }
  }
 
  @Override
  public void onCompleted() {
- if (last != null && last.hasResource()) {
- completeCallback.accept(last);
+ try {
+ if (last == null) {
+ throw new StorageException(
+ 0, "onComplete without preceding onNext, unable to determine success.");
+ } else if (!finalizing && last.hasPersistedSize()) { // incremental
+ long totalSentBytes = writeCtx.getTotalSentBytes().get();
+ long persistedSize = last.getPersistedSize();
+
+ if (totalSentBytes == persistedSize) {
+ writeCtx.getConfirmedBytes().set(persistedSize);
+ } else if (persistedSize < totalSentBytes) {
+ long delta = totalSentBytes - persistedSize;
+ // rewind our content and any state that my have run ahead of the actual ack'd bytes
+ content.rewindTo(delta);
+ writeCtx.getTotalSentBytes().set(persistedSize);
+ writeCtx.getConfirmedBytes().set(persistedSize);
+ } else {
+ throw ResumableSessionFailureScenario.SCENARIO_7.toStorageException();
+ }
+ } else if (finalizing && last.hasResource()) {
+ long totalSentBytes = writeCtx.getTotalSentBytes().get();
+ long finalSize = last.getResource().getSize();
+ if (totalSentBytes == finalSize) {
+ writeCtx.getConfirmedBytes().set(finalSize);
+ resultFuture.set(last);
+ } else if (finalSize < totalSentBytes) {
+ throw ResumableSessionFailureScenario.SCENARIO_4_1.toStorageException();
+ } else {
+ throw ResumableSessionFailureScenario.SCENARIO_4_2.toStorageException();
+ }
+ } else if (!finalizing && last.hasResource()) {
+ throw ResumableSessionFailureScenario.SCENARIO_1.toStorageException();
+ } else if (finalizing && last.hasPersistedSize()) {
+ long totalSentBytes = writeCtx.getTotalSentBytes().get();
+ long persistedSize = last.getPersistedSize();
+ if (persistedSize < totalSentBytes) {
+ throw ResumableSessionFailureScenario.SCENARIO_3.toStorageException();
+ } else {
+ throw ResumableSessionFailureScenario.SCENARIO_2.toStorageException();
+ }
+ } else {
+ throw ResumableSessionFailureScenario.SCENARIO_0.toStorageException();
+ }
+ } catch (Throwable se) {
+ open = false;
+ invocationHandle.setException(se);
+ } finally {
+ invocationHandle.set(null);
  }
- invocationHandle.set(null);
  }
 
  void await() {

google-cloud-storage/src/main/java/com/google/cloud/storage/GapicWritableByteChannelSessionBuilder.java

@@ -299,13 +299,13 @@ UnbufferedWritableByteChannelSession<WriteObjectResponse> build() {
  result,
  getChunkSegmenter(),
  write,
- ResumableWrite.identity(start),
+ new WriteCtx<>(start),
  deps,
  alg,
  Retrying::newCallContext);
  } else {
  return new GapicUnbufferedFinalizeOnCloseResumableWritableByteChannel(
- result, getChunkSegmenter(), write, ResumableWrite.identity(start));
+ result, getChunkSegmenter(), write, start);
  }
  })
  .andThen(StorageByteChannels.writable()::createSynchronized));
@@ -340,13 +340,13 @@ BufferedWritableByteChannelSession<WriteObjectResponse> build() {
  result,
  getChunkSegmenter(),
  write,
- ResumableWrite.identity(start),
+ new WriteCtx<>(start),
  deps,
  alg,
  Retrying::newCallContext);
  } else {
  return new GapicUnbufferedFinalizeOnCloseResumableWritableByteChannel(
- result, getChunkSegmenter(), write, ResumableWrite.identity(start));
+ result, getChunkSegmenter(), write, start);
  }
  })
  .andThen(c -> new DefaultBufferedWritableByteChannel(bufferHandle, c))