(iceberg) branch main updated: Core: Minor updates to AES GCM streams (#9453)

[blue]

This is an automated email from the ASF dual-hosted git repository.

blue pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/iceberg.git


The following commit(s) were added to refs/heads/main by this push:
     new e76988b1aa Core: Minor updates to AES GCM streams (#9453)
e76988b1aa is described below

commit e76988b1aa29c3cffca64dc36ccebbfb1c670561
Author: ggershinsky <ggershin...@users.noreply.github.com>
AuthorDate: Sun Jan 14 23:29:59 2024 +0200

    Core: Minor updates to AES GCM streams (#9453)
---
 .../org/apache/iceberg/encryption/AesGcmInputStream.java  |  2 +-
 .../org/apache/iceberg/encryption/AesGcmOutputStream.java | 15 +++++++++++++++
 .../org/apache/iceberg/encryption/TestGcmStreams.java     |  6 +++---
 3 files changed, 19 insertions(+), 4 deletions(-)

diff --git 
a/core/src/main/java/org/apache/iceberg/encryption/AesGcmInputStream.java 
b/core/src/main/java/org/apache/iceberg/encryption/AesGcmInputStream.java
index 57fed69c01..1f52ab3682 100644
--- a/core/src/main/java/org/apache/iceberg/encryption/AesGcmInputStream.java
+++ b/core/src/main/java/org/apache/iceberg/encryption/AesGcmInputStream.java
@@ -112,7 +112,7 @@ public class AesGcmInputStream extends SeekableInputStream {
     }
 
     if (available() <= 0 && len > 0) {
-      throw new EOFException();
+      return -1;
     }
 
     if (len == 0) {
diff --git 
a/core/src/main/java/org/apache/iceberg/encryption/AesGcmOutputStream.java 
b/core/src/main/java/org/apache/iceberg/encryption/AesGcmOutputStream.java
index 4db0802ea1..b4f723cca3 100644
--- a/core/src/main/java/org/apache/iceberg/encryption/AesGcmOutputStream.java
+++ b/core/src/main/java/org/apache/iceberg/encryption/AesGcmOutputStream.java
@@ -44,6 +44,8 @@ public class AesGcmOutputStream extends PositionOutputStream {
   private int currentBlockIndex;
   private boolean isHeaderWritten;
   private boolean lastBlockWritten;
+  private boolean isClosed;
+  private long finalPosition;
 
   AesGcmOutputStream(PositionOutputStream targetStream, byte[] aesKey, byte[] 
fileAadPrefix) {
     this.targetStream = targetStream;
@@ -56,6 +58,8 @@ public class AesGcmOutputStream extends PositionOutputStream {
     this.currentBlockIndex = 0;
     this.isHeaderWritten = false;
     this.lastBlockWritten = false;
+    this.isClosed = false;
+    this.finalPosition = 0;
   }
 
   @Override
@@ -66,6 +70,10 @@ public class AesGcmOutputStream extends PositionOutputStream 
{
 
   @Override
   public void write(byte[] b, int off, int len) throws IOException {
+    if (isClosed) {
+      throw new IOException("Writing to closed stream");
+    }
+
     if (!isHeaderWritten) {
       writeHeader();
     }
@@ -95,6 +103,10 @@ public class AesGcmOutputStream extends 
PositionOutputStream {
 
   @Override
   public long getPos() throws IOException {
+    if (isClosed) {
+      return finalPosition;
+    }
+
     return (long) currentBlockIndex * Ciphers.PLAIN_BLOCK_SIZE + 
positionInPlainBlock;
   }
 
@@ -109,6 +121,9 @@ public class AesGcmOutputStream extends 
PositionOutputStream {
       writeHeader();
     }
 
+    finalPosition = getPos();
+    isClosed = true;
+
     encryptAndWriteBlock();
 
     targetStream.close();
diff --git 
a/core/src/test/java/org/apache/iceberg/encryption/TestGcmStreams.java 
b/core/src/test/java/org/apache/iceberg/encryption/TestGcmStreams.java
index 773d5f41af..a954cf760b 100644
--- a/core/src/test/java/org/apache/iceberg/encryption/TestGcmStreams.java
+++ b/core/src/test/java/org/apache/iceberg/encryption/TestGcmStreams.java
@@ -18,7 +18,6 @@
  */
 package org.apache.iceberg.encryption;
 
-import java.io.EOFException;
 import java.io.File;
 import java.io.IOException;
 import java.nio.ByteBuffer;
@@ -60,8 +59,7 @@ public class TestGcmStreams {
     Assert.assertEquals("File size", 0, decryptedFile.getLength());
 
     try (SeekableInputStream decryptedStream = decryptedFile.newStream()) {
-      Assertions.assertThatThrownBy(() -> decryptedStream.read(readBytes))
-          .isInstanceOf(EOFException.class);
+      Assert.assertEquals("Read empty stream", -1, 
decryptedStream.read(readBytes));
     }
 
     // check that the AAD is still verified, even for an empty file
@@ -284,6 +282,7 @@ public class TestGcmStreams {
         }
 
         encryptedStream.close();
+        Assert.assertEquals("Final position in closed stream", offset, 
encryptedStream.getPos());
 
         AesGcmInputFile decryptedFile =
             new AesGcmInputFile(Files.localInput(testFile), key, aadPrefix);
@@ -380,6 +379,7 @@ public class TestGcmStreams {
       }
 
       encryptedStream.close();
+      Assert.assertEquals("Final position in closed stream", offset, 
encryptedStream.getPos());
 
       AesGcmInputFile decryptedFile =
           new AesGcmInputFile(Files.localInput(testFile), key, aadPrefix);