Crypto changes

Signed-off-by: Kunal Kotwani <[email protected]>

libs/common/src/main/java/org/opensearch/common/crypto/CryptoHandler.java

@@ -115,4 +115,21 @@ public interface CryptoHandler<T, U> extends Closeable {
      * @param endPosOfRawContent ending position in the raw/decrypted content
      */
     DecryptedRangedStreamProvider createDecryptingStreamOfRange(U cryptoContext, long startPosOfRawContent, long endPosOfRawContent);
+
+    /**
+     * This method creates a {@link DecryptedRangedStreamProvider} which provides a wrapped stream to decrypt the
+     * underlying stream.
+     * If provided encrypted positions are not aligned with encryption mechanics then this method can throw
+     * {@link IllegalArgumentException}
+     *
+     * @param cryptoContext crypto metadata instance.
+     * @param startPosOfEncContent starting position in encrypted content
+     * @param endPosOfEncContent ending position in encrypted content
+     */
+    DecryptedRangedStreamProvider createDecryptingStreamFromEncryptedOffsets(
+        U cryptoContext,
+        long fullEncryptedLength,
+        long startPosOfEncContent,
+        long endPosOfEncContent
+    );
 }

libs/common/src/main/java/org/opensearch/common/crypto/DecryptedRangedStreamProvider.java

@@ -17,26 +17,26 @@
 public class DecryptedRangedStreamProvider {
 
     /** Adjusted range of partial encrypted content which needs to be used for decryption. */
-    private final long[] adjustedRange;
+    private final long[] adjustedEncryptedRange;
     /** Stream provider for decryption and range re-adjustment. */
     private final UnaryOperator<InputStream> decryptedStreamProvider;
 
     /**
      * To construct adjusted encrypted range.
-     * @param adjustedRange range of partial encrypted content which needs to be used for decryption.
+     * @param adjustedEncryptedRange range of partial encrypted content which needs to be used for decryption.
      * @param decryptedStreamProvider stream provider for decryption and range re-adjustment.
      */
-    public DecryptedRangedStreamProvider(long[] adjustedRange, UnaryOperator<InputStream> decryptedStreamProvider) {
-        this.adjustedRange = adjustedRange;
+    public DecryptedRangedStreamProvider(long[] adjustedEncryptedRange, UnaryOperator<InputStream> decryptedStreamProvider) {
+        this.adjustedEncryptedRange = adjustedEncryptedRange;
         this.decryptedStreamProvider = decryptedStreamProvider;
     }
 
     /**
      * Adjusted range of partial encrypted content which needs to be used for decryption.
      * @return adjusted range
      */
-    public long[] getAdjustedRange() {
-        return adjustedRange;
+    public long[] getAdjustedEncryptedRange() {
+        return adjustedEncryptedRange;
     }
 
     /**

modules/crypto/src/main/java/org/opensearch/encryption/NoOpCryptoHandler.java

@@ -125,6 +125,17 @@ public DecryptedRangedStreamProvider createDecryptingStreamOfRange(
         return new DecryptedRangedStreamProvider(range, (encryptedStream) -> encryptedStream);
     }
 
+    @Override
+    public DecryptedRangedStreamProvider createDecryptingStreamFromEncryptedOffsets(
+        Object cryptoContext,
+        long fullEncryptedLength,
+        long startPosOfEncContent,
+        long endPosOfEncContent
+    ) {
+        long[] range = { startPosOfEncContent, endPosOfEncContent };
+        return new DecryptedRangedStreamProvider(range, (encryptedStream) -> encryptedStream);
+    }
+
     @Override
     public void close() {
         // Nothing to close.

modules/crypto/src/main/java/org/opensearch/encryption/frame/AwsCrypto.java

@@ -111,7 +111,28 @@ public long estimateOutputSizeWithFooter(int frameLen, int nonceLen, int tagLen,
 
     public long estimateDecryptedSize(int frameLen, int nonceLen, int tagLen, long contentLength, CryptoAlgorithm cryptoAlgorithm) {
         long contentLenWithoutTrailingSig = contentLength - getTrailingSignatureSize(cryptoAlgorithm);
-        return FrameDecryptionHandler.estimateDecryptedSize(contentLenWithoutTrailingSig, frameLen, nonceLen, tagLen);
+        return FrameDecryptionHandler.estimateDecryptedSize(contentLenWithoutTrailingSig, frameLen, nonceLen, tagLen, true);
+    }
+
+    public long estimatePartialDecryptedSize(
+        long fullEncryptedSize,
+        long partialEncryptedSize,
+        int frameLen,
+        int nonceLen,
+        int tagLen,
+        CryptoAlgorithm cryptoAlgorithm
+    ) {
+        long fullEncWithoutTrailingSig = fullEncryptedSize - getTrailingSignatureSize(cryptoAlgorithm);
+        if (fullEncWithoutTrailingSig <= partialEncryptedSize) {
+            partialEncryptedSize = fullEncWithoutTrailingSig;
+        }
+        return FrameDecryptionHandler.estimatePartialDecryptedSize(
+            fullEncWithoutTrailingSig,
+            partialEncryptedSize,
+            frameLen,
+            nonceLen,
+            tagLen
+        );
     }
 
     public int getTrailingSignatureSize(CryptoAlgorithm cryptoAlgorithm) {

modules/crypto/src/main/java/org/opensearch/encryption/frame/FrameCryptoHandler.java

@@ -105,6 +105,28 @@ public long estimateDecryptedLength(ParsedCiphertext parsedCiphertext, long cont
         );
     }
 
+    /**
+     * Estimate length of the decrypted stream.
+     *
+     * @param parsedCiphertext crypto metadata instance
+     * @param fullEncryptedSize Size of the entire encrypted content
+     * @param partialEncryptedSize Size of partial encrypted content
+     * @return Calculated size of the encrypted stream for the provided raw stream.
+     */
+    public long estimatePartialDecryptedLength(ParsedCiphertext parsedCiphertext, long fullEncryptedSize, long partialEncryptedSize) {
+        if (partialEncryptedSize <= parsedCiphertext.getOffset()) {
+            return 0;
+        }
+        return awsCrypto.estimatePartialDecryptedSize(
+            fullEncryptedSize - parsedCiphertext.getOffset(),
+            partialEncryptedSize - parsedCiphertext.getOffset(),
+            parsedCiphertext.getFrameLength(),
+            parsedCiphertext.getNonceLength(),
+            parsedCiphertext.getCryptoAlgoId().getTagLen(),
+            parsedCiphertext.getCryptoAlgoId()
+        );
+    }
+
     /**
      * Wraps a raw InputStream with encrypting stream
      * @param encryptionMetadata consists encryption metadata.
@@ -209,19 +231,56 @@ public DecryptedRangedStreamProvider createDecryptingStreamOfRange(
         long adjustedEndPos = endPosOverhead == 0
             ? endPosOfRawContent
             : (endPosOfRawContent - endPosOverhead + encryptionMetadata.getFrameLength());
-        long[] encryptedRange = transformToEncryptedRange(encryptionMetadata, adjustedStartPos, adjustedEndPos);
-        return new DecryptedRangedStreamProvider(encryptedRange, (encryptedStream) -> {
+        long[] adjustedEncryptedRange = transformToEncryptedRange(encryptionMetadata, adjustedStartPos, adjustedEndPos);
+        long[] adjustedRange = new long[4];
+        adjustedRange[0] = adjustedEncryptedRange[0];
+        adjustedRange[1] = adjustedEncryptedRange[1];
+        adjustedRange[2] = startPosOfRawContent;
+        adjustedRange[3] = endPosOfRawContent;
+
+        return new DecryptedRangedStreamProvider(adjustedRange, (encryptedStream) -> {
             InputStream decryptedStream = createBlockDecryptionStream(
                 encryptionMetadata,
                 encryptedStream,
                 adjustedStartPos,
                 adjustedEndPos,
-                encryptedRange
+                adjustedEncryptedRange
             );
             return new TrimmingStream(adjustedStartPos, adjustedEndPos, startPosOfRawContent, endPosOfRawContent, decryptedStream);
         });
     }
 
+    /**
+     * This method creates a {@link DecryptedRangedStreamProvider} which provides a wrapped stream to decrypt the
+     * underlying stream.
+     * If provided startPosOfEncContent does not fall on frame start position or endPosOfEncContent does not
+     * fall on frame end position then this method will throw an {@link IllegalArgumentException}.
+     *
+     * @param encryptionMetadata crypto metadata instance.
+     * @param fullEncryptedLength length of entire encrypted content.
+     * @param startPosOfEncContent starting position in encrypted content
+     * @param endPosOfEncContent ending position in encrypted content
+     */
+    @Override
+    public DecryptedRangedStreamProvider createDecryptingStreamFromEncryptedOffsets(
+        ParsedCiphertext encryptionMetadata,
+        long fullEncryptedLength,
+        long startPosOfEncContent,
+        long endPosOfEncContent
+    ) {
+        long startPosOfRawContent = estimatePartialDecryptedLength(encryptionMetadata, fullEncryptedLength, startPosOfEncContent);
+        long endPosOfRawContent = estimatePartialDecryptedLength(encryptionMetadata, fullEncryptedLength, endPosOfEncContent);
+        if (startPosOfRawContent % encryptionMetadata.getFrameLength() != 0
+            || (endPosOfRawContent + 1) % encryptionMetadata.getFrameLength() != 0 && endPosOfEncContent < fullEncryptedLength) {
+            // We can't estimate decryption sizes for random encryption lengths because frame positions within a frame metadata
+            // map to the same index and therefore, for different metadata positions in a frame, previous end position and current
+            // start position map to the same index which can result into consistent decrypted bytes.
+            throw new IllegalArgumentException("Computed start and end positions of raw content do not align with frame boundaries.");
+        }
+
+        return createDecryptingStreamOfRange(encryptionMetadata, startPosOfRawContent, endPosOfRawContent);
+    }
+
     private long[] transformToEncryptedRange(ParsedCiphertext parsedCiphertext, long startPosOfRawContent, long endPosOfRawContent) {
 
         long startPos = awsCrypto.estimatePartialOutputSize(

modules/crypto/src/main/java/org/opensearch/encryption/frame/FrameDecryptionHandler.java

@@ -56,6 +56,7 @@ class FrameDecryptionHandler implements CryptoHandler {
 
     boolean complete_ = false;
     private byte[] unparsedBytes_ = new byte[0];
+    private static final long lastFrameHeaderOverhead = (Integer.SIZE / Byte.SIZE) * 2;
 
     /**
      * Construct a decryption handler for decrypting bytes stored in frames.
@@ -241,18 +242,16 @@ public int estimateOutputSize(final int inLen) {
         return outSize;
     }
 
-    public static long estimateDecryptedSize(long encryptedSize, int frameSize, int nonceLen, int tagLenBytes) {
-        // Calculate the size of sequence number for the last frame
-        long lastFrameSeqNumberSize = (Integer.SIZE / Byte.SIZE);
-
-        // Calculate the size of the final frame size
-        long finalFrameSizeSize = (Integer.SIZE / Byte.SIZE);
-
-        // Calculate the total size of header overhead for the last frame
-        long lastFrameHeaderOverhead = lastFrameSeqNumberSize + finalFrameSizeSize;
-
+    public static long estimateDecryptedSize(
+        long encryptedSize,
+        int frameSize,
+        int nonceLen,
+        int tagLenBytes,
+        boolean includeLastFrameHeaderOverhead
+    ) {
+        long curLastFrameHeaderOverhead = includeLastFrameHeaderOverhead ? lastFrameHeaderOverhead : 0;
         // Calculate the number of frames
-        long frames = (encryptedSize - lastFrameHeaderOverhead) / (frameSize + nonceLen + tagLenBytes + (Integer.SIZE / Byte.SIZE)) + 1;
+        long frames = (encryptedSize - curLastFrameHeaderOverhead) / (frameSize + nonceLen + tagLenBytes + (Integer.SIZE / Byte.SIZE)) + 1;
 
         // Calculate the size of the actual content in frames
         long contentSizeWithoutLastFrame = (frames - 1) * frameSize;
@@ -264,11 +263,38 @@ public static long estimateDecryptedSize(long encryptedSize, int frameSize, int
         long headerOverhead = (nonceLen + tagLenBytes) * frames + seqNumberSize;
 
         // Calculate the size of the last frame content
-        long lastFrameSize = encryptedSize - contentSizeWithoutLastFrame - headerOverhead - lastFrameHeaderOverhead;
+        long lastFrameSize = encryptedSize - contentSizeWithoutLastFrame - headerOverhead - curLastFrameHeaderOverhead;
+
+        // Case where index in last frame data is somewhere in the frame metadata.
+        lastFrameSize = Math.max(0, lastFrameSize);
 
         return contentSizeWithoutLastFrame + lastFrameSize;
     }
 
+    public static long estimatePartialDecryptedSize(
+        long fullEncryptedSize,
+        long partialEncryptedSize,
+        int frameSize,
+        int nonceLen,
+        int tagLenBytes
+    ) {
+        long encryptedFrameSize = frameSize + nonceLen + tagLenBytes + (Integer.SIZE / Byte.SIZE);
+
+        if (partialEncryptedSize % encryptedFrameSize == 0) {
+            return partialEncryptedSize / encryptedFrameSize * frameSize;
+        }
+
+        boolean includeLastFrameHeaderOverhead;
+        if (fullEncryptedSize - lastFrameHeaderOverhead < partialEncryptedSize) {
+            partialEncryptedSize = fullEncryptedSize;
+            includeLastFrameHeaderOverhead = true;
+        } else {
+            includeLastFrameHeaderOverhead = false;
+        }
+
+        return estimateDecryptedSize(partialEncryptedSize, frameSize, nonceLen, tagLenBytes, includeLastFrameHeaderOverhead);
+    }
+
     @Override
     public int estimatePartialOutputSize(int inLen) {
         return estimateOutputSize(inLen);

modules/crypto/src/test/java/org/opensearch/encryption/frame/CryptoTests.java

@@ -434,7 +434,7 @@ private void validateBlockDownload(EncryptedStoreTest encryptedStoreTest, int st
             endPos
         );
 
-        long[] transformedRange = decryptedStreamProvider.getAdjustedRange();
+        long[] transformedRange = decryptedStreamProvider.getAdjustedEncryptedRange();
         int encryptedBlockSize = (int) (transformedRange[1] - transformedRange[0] + 1);
         byte[] encryptedBlockBytes = new byte[encryptedBlockSize];
         System.arraycopy(encryptedStoreTest.encryptedContent, (int) transformedRange[0], encryptedBlockBytes, 0, encryptedBlockSize);

server/src/main/java/org/opensearch/common/blobstore/AsyncMultiStreamEncryptedBlobContainer.java

@@ -164,8 +164,8 @@ private InputStreamContainer decryptInputStreamContainer(InputStreamContainer in
                 endOfStream
             );
 
-            long adjustedPos = decryptedStreamProvider.getAdjustedRange()[0];
-            long adjustedLength = decryptedStreamProvider.getAdjustedRange()[1] - adjustedPos + 1;
+            long adjustedPos = decryptedStreamProvider.getAdjustedEncryptedRange()[0];
+            long adjustedLength = decryptedStreamProvider.getAdjustedEncryptedRange()[1] - adjustedPos + 1;
             final InputStream decryptedStream = decryptedStreamProvider.getDecryptedStreamProvider()
                 .apply(inputStreamContainer.getInputStream());
             return new InputStreamContainer(decryptedStream, adjustedLength, adjustedPos);

server/src/main/java/org/opensearch/common/blobstore/EncryptedBlobContainer.java

@@ -70,8 +70,8 @@ public InputStream readBlob(String blobName, long position, long length) throws
             position,
             position + length - 1
         );
-        long adjustedPos = decryptedStreamProvider.getAdjustedRange()[0];
-        long adjustedLength = decryptedStreamProvider.getAdjustedRange()[1] - adjustedPos + 1;
+        long adjustedPos = decryptedStreamProvider.getAdjustedEncryptedRange()[0];
+        long adjustedLength = decryptedStreamProvider.getAdjustedEncryptedRange()[1] - adjustedPos + 1;
         InputStream encryptedStream = blobContainer.readBlob(blobName, adjustedPos, adjustedLength);
         return decryptedStreamProvider.getDecryptedStreamProvider().apply(encryptedStream);
     }