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.

6 files changed, 93 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());
diff --git a/vespaclient-java/src/test/java/com/yahoo/vespa/security/tool/CryptoToolsTest.java b/vespaclient-java/src/test/java/com/yahoo/vespa/security/tool/CryptoToolsTest.java
index d7b8f1f09ae..f55278342e1 100644
--- a/vespaclient-java/src/test/java/com/yahoo/vespa/security/tool/CryptoToolsTest.java
+++ b/vespaclient-java/src/test/java/com/yahoo/vespa/security/tool/CryptoToolsTest.java
@@ -537,6 +537,50 @@ public class CryptoToolsTest {
         assertEquals("", procOut.stdOut());
     }
 
+    @Test
+    void can_transparently_compress_and_decompress_plaintext() throws IOException {
+        String substring = "here is some stuff that can be compressed!";
+        String compressibleSecret = substring.repeat(100);
+
+        var secretFile = pathInTemp("secret.txt");
+        Files.writeString(secretFile, compressibleSecret);
+
+        var privKeyFile   = pathInTemp("my-priv.txt");
+        writePrivateKeyFile(privKeyFile, TEST_PRIV_KEY);
+
+        var encryptedPath = pathInTemp("encrypted.bin");
+        var procOut = runMain(List.of(
+                "encrypt",
+                absPathOf(secretFile),
+                "--output-file",          absPathOf(encryptedPath),
+                "--recipient-public-key", TEST_PUB_KEY,
+                "--key-id",               "1234",
+                "--zstd-compress"));
+        assertEquals(0, procOut.exitCode());
+        assertEquals("", procOut.stdErr());
+
+        var token = procOut.stdOut();
+        assertFalse(token.isBlank());
+
+        assertTrue(Files.exists(encryptedPath));
+        assertTrue(Files.size(encryptedPath) < compressibleSecret.length());
+
+        var decryptedPath = pathInTemp("decrypted.txt");
+        procOut = runMain(List.of(
+                "decrypt",
+                absPathOf(encryptedPath),
+                "--output-file",      absPathOf(decryptedPath),
+                "--private-key-file", absPathOf(privKeyFile),
+                "--token",            token,
+                "--zstd-decompress"
+        ));
+        assertEquals(0, procOut.exitCode());
+        assertEquals("", procOut.stdOut());
+        assertEquals("", procOut.stdErr());
+
+        assertEquals(compressibleSecret, Files.readString(decryptedPath));
+    }
+
     private ProcessOutput runMain(List<String> args) {
         return runMain(args, EMPTY_BYTES);
     }
diff --git a/vespaclient-java/src/test/resources/expected-decrypt-help-output.txt b/vespaclient-java/src/test/resources/expected-decrypt-help-output.txt
index f00db2bb6b9..ab47d11c602 100644
--- a/vespaclient-java/src/test/resources/expected-decrypt-help-output.txt
+++ b/vespaclient-java/src/test/resources/expected-decrypt-help-output.txt
@@ -21,5 +21,7 @@ the quotes).
                                plaintext to STDOUT instead of a file.
  -t,--token <arg>              Token generated when the input file was
                                encrypted
+ -z,--zstd-decompress          Decrypted data will be transparently
+                               Zstd-decompressed before being output.
 Note: this is a BETA tool version; its interface may be changed at any
 time
diff --git a/vespaclient-java/src/test/resources/expected-encrypt-help-output.txt b/vespaclient-java/src/test/resources/expected-encrypt-help-output.txt
index 46185b29986..55cdfc73e9a 100644
--- a/vespaclient-java/src/test/resources/expected-encrypt-help-output.txt
+++ b/vespaclient-java/src/test/resources/expected-encrypt-help-output.txt
@@ -12,5 +12,7 @@ the quotes).
                                    already exists)
  -r,--recipient-public-key <arg>   Recipient X25519 public key in Base58
                                    encoded format
+ -z,--zstd-compress                Input data will be transparently
+                                   Zstd-compressed before being encrypted.
 Note: this is a BETA tool version; its interface may be changed at any
 time