fix: fix CipherSubscriber to only call onNext once per request

src/main/java/software/amazon/encryption/s3/internal/CipherSubscriber.java

@@ -11,22 +11,26 @@
 import javax.crypto.Cipher;
 import java.nio.ByteBuffer;
 import java.security.GeneralSecurityException;
+import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicLong;
 
 public class CipherSubscriber implements Subscriber<ByteBuffer> {
  private final AtomicLong contentRead = new AtomicLong(0);
  private final Subscriber<? super ByteBuffer> wrappedSubscriber;
- private Cipher cipher;
+ private final Cipher cipher;
  private final Long contentLength;
- private boolean isLastPart;
+ private final boolean isLastPart;
+ private final int tagLength;
+ private final AtomicBoolean finalBytesCalled = new AtomicBoolean(false);
 
  private byte[] outputBuffer;
 
  CipherSubscriber(Subscriber<? super ByteBuffer> wrappedSubscriber, Long contentLength, CryptographicMaterials materials, byte[] iv, boolean isLastPart) {
  this.wrappedSubscriber = wrappedSubscriber;
  this.contentLength = contentLength;
- cipher = materials.getCipher(iv);
+ this.cipher = materials.getCipher(iv);
  this.isLastPart = isLastPart;
+ this.tagLength = materials.algorithmSuite().cipherTagLengthBytes();
  }
 
  CipherSubscriber(Subscriber<? super ByteBuffer> wrappedSubscriber, Long contentLength, CryptographicMaterials materials, byte[] iv) {
@@ -46,20 +50,50 @@ public void onNext(ByteBuffer byteBuffer) {
  if (amountToReadFromByteBuffer > 0) {
  byte[] buf = BinaryUtils.copyBytesFrom(byteBuffer, amountToReadFromByteBuffer);
  outputBuffer = cipher.update(buf, 0, amountToReadFromByteBuffer);
+
  if (outputBuffer == null || outputBuffer.length == 0) {
  // The underlying data is too short to fill in the block cipher.
  // Note that while the JCE Javadoc specifies that the outputBuffer is null in this case,
  // in practice SunJCE and ACCP return an empty buffer instead, hence checks for
  // null OR length == 0.
- if (contentRead.get() == contentLength) {
+ if (contentRead.get() + tagLength >= contentLength) {
  // All content has been read, so complete to get the final bytes
- this.onComplete();
+ finalBytes();
+ return;
  }
  // Otherwise, wait for more bytes. To avoid blocking,
  // send an empty buffer to the wrapped subscriber.
  wrappedSubscriber.onNext(ByteBuffer.allocate(0));
  } else {
- wrappedSubscriber.onNext(ByteBuffer.wrap(outputBuffer));
+ /*
+ Check if stream has read all expected content.
+ Once all content has been read, call onComplete.
+
+ This determines that all content has been read by checking if
+ the amount of data read so far plus the tag length is at least the content length.
+ Once this is true, downstream will never call `request` again
+ (beyond the current request that is being responded to in this onNext invocation.)
+ As a result, this class can only call `wrappedSubscriber.onNext` one more time.
+ (Reactive streams require that downstream sends a `request(n)`
+ to indicate it is ready for more data, and upstream responds to that request by calling `onNext`.
+ The `n` in request is the maximum number of `onNext` calls that downstream
+ will allow upstream to make, and seems to always be 1 for the AsyncBodySubscriber.)
+ Since this class can only call `wrappedSubscriber.onNext` once,
+ it must send all remaining data in the next onNext call,
+ including the result of cipher.doFinal(), if applicable.
+ Calling `wrappedSubscriber.onNext` more than once for `request(1)`
+ violates the Reactive Streams specification and can cause exceptions downstream.
+ */
+ if (contentRead.get() + tagLength >= contentLength) {
+ // All content has been read; complete the stream.
+ // (Signalling onComplete from here is Reactive Streams-spec compliant;
+ // this class is allowed to call onComplete, even if upstream has not yet signaled onComplete.)
+ finalBytes();
+ } else {
+ // Needs to read more data, so send the data downstream,
+ // expecting that downstream will continue to request more data.
+ wrappedSubscriber.onNext(ByteBuffer.wrap(outputBuffer));
+ }
  }
  } else {
  // Do nothing
@@ -91,21 +125,63 @@ public void onError(Throwable t) {
 
  @Override
  public void onComplete() {
+ // In rare cases, e.g. when the last part of a low-level MPU has 0 length,
+ // onComplete will be called before onNext is called once.
+ if (contentRead.get() < contentLength) {
+ finalBytes();
+ }
+ wrappedSubscriber.onComplete();
+ }
+
+ /**
+ * Finalize encryption, including calculating the auth tag for AES-GCM.
+ * As such this method MUST only be called once, which is enforced using
+ * `finalBytesCalled`.
+ */
+ private void finalBytes() {
+ if (!finalBytesCalled.compareAndSet(false, true)) {
+ // already called, don't repeat
+ return;
+ }
+
+ // If this isn't the last part, skip doFinal and just send outputBuffer downstream.
+ // doFinal requires that all parts have been processed to compute the tag,
+ // so the tag will only be computed when the last part is processed.
  if (!isLastPart) {
- // If this isn't the last part, skip doFinal, we aren't done
- wrappedSubscriber.onComplete();
+ wrappedSubscriber.onNext(ByteBuffer.wrap(outputBuffer));
  return;
  }
+
+ // If this is the last part, compute doFinal and include its result in the value sent downstream.
+ // The result of doFinal MUST be included with the bytes that were in outputBuffer in the final onNext call.
+ byte[] finalBytes;
  try {
- outputBuffer = cipher.doFinal();
- // Send the final bytes to the wrapped subscriber
- wrappedSubscriber.onNext(ByteBuffer.wrap(outputBuffer));
+ finalBytes = cipher.doFinal();
  } catch (final GeneralSecurityException exception) {
+ // Even if doFinal fails, downstream still expects to receive the bytes that were in outputBuffer
+ wrappedSubscriber.onNext(ByteBuffer.wrap(outputBuffer));
  // Forward error, else the wrapped subscriber waits indefinitely
  wrappedSubscriber.onError(exception);
  throw new S3EncryptionClientSecurityException(exception.getMessage(), exception);
  }
- wrappedSubscriber.onComplete();
+
+ // Combine the bytes from outputBuffer and finalBytes into one onNext call.
+ // Downstream has requested one item in its request method, so this class can only call onNext once.
+ // This single onNext call must contain both the bytes from outputBuffer and the tag.
+ byte[] combinedBytes;
+ if (outputBuffer != null && outputBuffer.length > 0 && finalBytes != null && finalBytes.length > 0) {
+ combinedBytes = new byte[outputBuffer.length + finalBytes.length];
+ System.arraycopy(outputBuffer, 0, combinedBytes, 0, outputBuffer.length);
+ System.arraycopy(finalBytes, 0, combinedBytes, outputBuffer.length, finalBytes.length);
+ } else if (outputBuffer != null && outputBuffer.length > 0) {
+ combinedBytes = outputBuffer;
+ } else if (finalBytes != null && finalBytes.length > 0) {
+ combinedBytes = finalBytes;
+ } else {
+ combinedBytes = new byte[0];
+ }
+
+ wrappedSubscriber.onNext(ByteBuffer.wrap(combinedBytes));
  }
 
 }

src/test/java/software/amazon/encryption/s3/S3AsyncEncryptionClientTest.java

@@ -729,6 +729,8 @@ public void copyObjectTransparentlyAsync() {
  * Test which artificially limits the size of buffers using {@link TinyBufferAsyncRequestBody}.
  * This tests edge cases where network conditions result in buffers with length shorter than
  * the cipher's block size.
+ * Note that TinyAsyncRequestBody is not fully spec-compliant, and will cause IllegalStateExceptions
+ * to be logged when debug logging is enabled.
  * @throws IOException
  */
  @Test

src/test/java/software/amazon/encryption/s3/S3EncryptionClientRangedGetCompatibilityTest.java

@@ -264,6 +264,7 @@ public void AsyncAesCbcV1toV3RangedGet(Object keyMaterial) {
  assertEquals("klmnopqrst0", output);
 
  // Valid start index within input and end index out of range, returns object from start index to End of Stream
+ // This causes a spurious NPE to be logged when debug logging is enabled.
  objectResponse = v3Client.getObject(builder -> builder
  .bucket(BUCKET)
  .range("bytes=190-300")
@@ -288,6 +289,7 @@ public void AsyncAesCbcV1toV3RangedGet(Object keyMaterial) {
  assertEquals(input, output);
 
  // Invalid range starting index and ending index greater than object length but within Cipher Block size, returns empty object
+ // This causes a spurious NPE to be logged when debug logging is enabled.
  objectResponse = v3Client.getObject(builder -> builder
  .bucket(BUCKET)
  .range("bytes=216-217")

src/test/java/software/amazon/encryption/s3/utils/TinyBufferAsyncRequestBody.java

@@ -10,6 +10,9 @@
  * AsyncRequestBody which wraps another AsyncRequestBody with a {@link TinyBufferSubscriber}.
  * This is useful for testing poor network conditions where buffers may not be larger than
  * the cipher's block size.
+ * DO NOT USE THIS IN PRODUCTION. In addition to degraded performance,
+ * it will cause IllegalStateExceptions in the base Subscriber as it does not comply
+ * with the Reactive Streaming spec.
  */
 public class TinyBufferAsyncRequestBody implements AsyncRequestBody {