Support Zstd (un)compression during crypto CLI encryption and decryption

Simplifies working with compressed plaintext, as it removes the need
for piping via `unzstd` or using a temporary file.

3 files changed, 45 insertions, 5 deletions

diff --git a/vespaclient-java/src/main/java/com/yahoo/vespa/security/tool/crypto/CipherUtils.java b/vespaclient-java/src/main/java/com/yahoo/vespa/security/tool/crypto/CipherUtils.java
index 5cb40aa8f3b..5834a166fb6 100644
--- a/vespaclient-java/src/main/java/com/yahoo/vespa/security/tool/crypto/CipherUtils.java
+++ b/vespaclient-java/src/main/java/com/yahoo/vespa/security/tool/crypto/CipherUtils.java
@@ -1,6 +1,8 @@
 // Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
 package com.yahoo.vespa.security.tool.crypto;
 
+import ai.vespa.airlift.zstd.ZstdInputStream;
+import com.yahoo.compress.ZstdOutputStream;
 import com.yahoo.security.AeadCipher;
 
 import java.io.IOException;
@@ -29,4 +31,26 @@ public class CipherUtils {
         }
     }
 
+    private static OutputStream maybeWrapCompress(OutputStream out, boolean compressZstd) throws IOException {
+        return compressZstd ? new ZstdOutputStream(out) : out;
+    }
+
+    public static void streamEncrypt(InputStream input, OutputStream output, AeadCipher cipher, boolean compressZstd) throws IOException {
+        try (var out = maybeWrapCompress(cipher.wrapOutputStream(output), compressZstd)) {
+            input.transferTo(out);
+            out.flush();
+        }
+    }
+
+    private static InputStream maybeWrapDecompress(InputStream in, boolean decompressZstd) throws IOException {
+        return decompressZstd ? new ZstdInputStream(in) : in;
+    }
+
+    public static void streamDecrypt(InputStream input, OutputStream output, AeadCipher cipher, boolean decompressZstd) throws IOException {
+        try (var in = maybeWrapDecompress(cipher.wrapInputStream(input), decompressZstd)) {
+            in.transferTo(output);
+            output.flush();
+        }
+    }
+
 }
diff --git a/vespaclient-java/src/main/java/com/yahoo/vespa/security/tool/crypto/DecryptTool.java b/vespaclient-java/src/main/java/com/yahoo/vespa/security/tool/crypto/DecryptTool.java
index ea79fe12c3d..ce3f5a89cd5 100644
--- a/vespaclient-java/src/main/java/com/yahoo/vespa/security/tool/crypto/DecryptTool.java
+++ b/vespaclient-java/src/main/java/com/yahoo/vespa/security/tool/crypto/DecryptTool.java
@@ -29,6 +29,7 @@ public class DecryptTool implements Tool {
 
     static final String OUTPUT_FILE_OPTION      = "output-file";
     static final String EXPECTED_KEY_ID_OPTION  = "expected-key-id";
+    static final String ZSTD_DECOMPRESS_OPTION  = "zstd-decompress";
     static final String TOKEN_OPTION            = "token";
 
     private static final List<Option> OPTIONS = List.of(
@@ -65,6 +66,12 @@ public class DecryptTool implements Tool {
                     .required(false)
                     .desc("Expected key ID in token. If this is not provided, the key ID is not verified.")
                     .build(),
+            Option.builder("z")
+                    .longOpt(ZSTD_DECOMPRESS_OPTION)
+                    .hasArg(false)
+                    .required(false)
+                    .desc("Decrypted data will be transparently Zstd-decompressed before being output.")
+                    .build(),
             Option.builder("t")
                     .longOpt(TOKEN_OPTION)
                     .hasArg(true)
@@ -107,10 +114,11 @@ public class DecryptTool implements Tool {
                                                                          !CliUtils.useStdIo(inputArg) && !CliUtils.useStdIo(outputArg));
             var secretShared = SharedKeyGenerator.fromSealedKey(sealedSharedKey, privateKey);
             var cipher       = SharedKeyGenerator.makeAesGcmDecryptionCipher(secretShared);
+            boolean unZstd   = arguments.hasOption(ZSTD_DECOMPRESS_OPTION);
 
             try (var inStream  = CliUtils.inputStreamFromFileOrStream(inputArg, invocation.stdIn());
                  var outStream = CliUtils.outputStreamToFileOrStream(outputArg, invocation.stdOut())) {
-                CipherUtils.streamEncipher(inStream, outStream, cipher);
+                CipherUtils.streamDecrypt(inStream, outStream, cipher, unZstd);
             }
         } catch (IOException e) {
             throw new RuntimeException(e);
diff --git a/vespaclient-java/src/main/java/com/yahoo/vespa/security/tool/crypto/EncryptTool.java b/vespaclient-java/src/main/java/com/yahoo/vespa/security/tool/crypto/EncryptTool.java
index 962b42f4c22..81a3eecce6b 100644
--- a/vespaclient-java/src/main/java/com/yahoo/vespa/security/tool/crypto/EncryptTool.java
+++ b/vespaclient-java/src/main/java/com/yahoo/vespa/security/tool/crypto/EncryptTool.java
@@ -28,6 +28,7 @@ public class EncryptTool implements Tool {
     static final String OUTPUT_FILE_OPTION          = "output-file";
     static final String KEY_ID_OPTION               = "key-id";
     static final String RECIPIENT_PUBLIC_KEY_OPTION = "recipient-public-key";
+    static final String ZSTD_COMPRESS_OPTION        = "zstd-compress";
 
     private static final List<Option> OPTIONS = List.of(
             Option.builder("o")
@@ -47,6 +48,12 @@ public class EncryptTool implements Tool {
                     .hasArg(true)
                     .required(false)
                     .desc("ID of recipient key")
+                    .build(),
+            Option.builder("z")
+                    .longOpt(ZSTD_COMPRESS_OPTION)
+                    .hasArg(false)
+                    .required(false)
+                    .desc("Input data will be transparently Zstd-compressed before being encrypted.")
                     .build());
 
     @Override
@@ -78,13 +85,14 @@ public class EncryptTool implements Tool {
             var outputPath = Paths.get(CliUtils.optionOrThrow(arguments, OUTPUT_FILE_OPTION));
 
             var recipientPubKey = KeyUtils.fromBase58EncodedX25519PublicKey(CliUtils.optionOrThrow(arguments, RECIPIENT_PUBLIC_KEY_OPTION).strip());
-            var keyId  = KeyId.ofString(CliUtils.optionOrThrow(arguments, KEY_ID_OPTION));
-            var shared = SharedKeyGenerator.generateForReceiverPublicKey(recipientPubKey, keyId);
-            var cipher = SharedKeyGenerator.makeAesGcmEncryptionCipher(shared);
+            var keyId    = KeyId.ofString(CliUtils.optionOrThrow(arguments, KEY_ID_OPTION));
+            var shared   = SharedKeyGenerator.generateForReceiverPublicKey(recipientPubKey, keyId);
+            var cipher   = SharedKeyGenerator.makeAesGcmEncryptionCipher(shared);
+            boolean zstd = arguments.hasOption(ZSTD_COMPRESS_OPTION);
 
             try (var inStream  = CliUtils.inputStreamFromFileOrStream(inputArg, invocation.stdIn());
                  var outStream = Files.newOutputStream(outputPath)) {
-                CipherUtils.streamEncipher(inStream, outStream, cipher);
+                CipherUtils.streamEncrypt(inStream, outStream, cipher, zstd);
             }
 
             invocation.stdOut().println(shared.sealedSharedKey().toTokenString());