Add basic tooling for public key encryption and decryption

Adds support for:
 * X25519 key pair generation
 * HPKE stream encryption with public key and token generation
 * HPKE stream decryption with private key

1 files changed, 296 insertions, 0 deletions

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
new file mode 100644
index 00000000000..bb8024544cb
--- /dev/null
+++ b/vespaclient-java/src/test/java/com/yahoo/vespa/security/tool/CryptoToolsTest.java
@@ -0,0 +1,296 @@
+// Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+package com.yahoo.vespa.security.tool;
+
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.io.TempDir;
+
+import java.io.ByteArrayOutputStream;
+import java.io.File;
+import java.io.IOException;
+import java.io.PrintStream;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.nio.file.Paths;
+import java.nio.file.attribute.PosixFilePermissions;
+import java.util.List;
+import java.util.Map;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertNotEquals;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+/**
+ * @author vekterli
+ */
+public class CryptoToolsTest {
+
+    private static record ProcessOutput(int exitCode, String stdOut, String stdErr) {}
+
+    @TempDir
+    public File tmpFolder;
+
+    private void verifyStdoutMatchesFile(List<String> args, String expectedFile) throws IOException {
+        var procOut = runMain(args, Map.of());
+        assertEquals(0, procOut.exitCode());
+        assertEquals(readTestResource(expectedFile), procOut.stdOut());
+    }
+
+    private void verifyStderrEquals(List<String> args, String expectedMessage) throws IOException {
+        var procOut = runMain(args, Map.of());
+        assertEquals(1, procOut.exitCode()); // Assume checking stderr is because of a failure.
+        assertEquals(expectedMessage, procOut.stdErr());
+    }
+
+    @Test
+    void top_level_help_page_printed_if_help_option_given() throws IOException {
+        verifyStdoutMatchesFile(List.of("--help"), "expected-help-output.txt");
+    }
+
+    @Test
+    void top_level_help_page_printed_if_no_option_given() throws IOException {
+        verifyStdoutMatchesFile(List.of(), "expected-help-output.txt");
+    }
+
+    @Test
+    void keygen_help_printed_if_help_option_given_to_subtool() throws IOException {
+        verifyStdoutMatchesFile(List.of("keygen", "--help"), "expected-keygen-help-output.txt");
+    }
+
+    @Test
+    void encrypt_help_printed_if_help_option_given_to_subtool() throws IOException {
+        verifyStdoutMatchesFile(List.of("encrypt", "--help"), "expected-encrypt-help-output.txt");
+    }
+
+    @Test
+    void decrypt_help_printed_if_help_option_given_to_subtool() throws IOException {
+        verifyStdoutMatchesFile(List.of("decrypt", "--help"), "expected-decrypt-help-output.txt");
+    }
+
+    @Test
+    void missing_required_parameter_prints_error_message() throws IOException {
+        // We don't test all possible input arguments to all tools, since it'd be too closely
+        // bound to the order in which the implementation checks for argument presence.
+        // This primarily verifies that IllegalArgumentExceptions thrown by a tool will be caught
+        // and printed to stderr as expected.
+        verifyStderrEquals(List.of("keygen"),
+                "Invalid command line arguments: Required argument '--private-out-file' must be provided\n");
+        verifyStderrEquals(List.of("keygen", "--private-out-file", "foo.txt"),
+                "Invalid command line arguments: Required argument '--public-out-file' must be provided\n");
+    }
+
+    // We don't want to casually overwrite key material if someone runs a command twice by accident.
+    @Test
+    void keygen_fails_by_default_if_output_file_exists() throws IOException {
+        Path privKeyFile = pathInTemp("priv.txt");
+        Path pubKeyFile  = pathInTemp("pub.txt");
+        Files.writeString(privKeyFile, TEST_PRIV_KEY);
+
+        verifyStderrEquals(List.of("keygen",
+                                   "--private-out-file", absPathOf(privKeyFile),
+                                   "--public-out-file",  absPathOf(pubKeyFile)),
+                ("Invalid command line arguments: Output file '%s' already exists. No keys written. " +
+                 "If you want to overwrite existing files, specify --overwrite-existing.\n")
+                .formatted(absPathOf(privKeyFile)));
+
+        Files.delete(privKeyFile);
+        Files.writeString(pubKeyFile, TEST_PUB_KEY);
+
+        verifyStderrEquals(List.of("keygen",
+                                   "--private-out-file", absPathOf(privKeyFile),
+                                   "--public-out-file",  absPathOf(pubKeyFile)),
+                ("Invalid command line arguments: Output file '%s' already exists. No keys written. " +
+                 "If you want to overwrite existing files, specify --overwrite-existing.\n")
+                 .formatted(absPathOf(pubKeyFile)));
+    }
+
+    // ... but we'll allow it if someone enables the foot-gun option.
+    @Test
+    void keygen_allowed_if_output_file_exists_and_explicit_overwrite_option_specified() throws IOException {
+        Path privKeyFile = pathInTemp("priv.txt");
+        Path pubKeyFile  = pathInTemp("pub.txt");
+        Files.writeString(privKeyFile, TEST_PRIV_KEY);
+        Files.writeString(pubKeyFile,  TEST_PUB_KEY);
+
+        var procOut = runMain(List.of("keygen",
+                                     "--private-out-file", absPathOf(privKeyFile),
+                                     "--public-out-file",  absPathOf(pubKeyFile),
+                                     "--overwrite-existing"));
+        assertEquals(0, procOut.exitCode());
+
+        // Keys are random, so we don't know what they'll end up being. But the likelihood of them
+        // exactly matching the test keys is effectively and realistically zero.
+        assertNotEquals(TEST_PRIV_KEY, Files.readString(privKeyFile));
+        assertNotEquals(TEST_PUB_KEY,  Files.readString(pubKeyFile));
+    }
+
+    @Test
+    void keygen_writes_private_key_with_user_only_rw_permissions() throws IOException {
+        Path privKeyFile = pathInTemp("priv.txt");
+        Path pubKeyFile  = pathInTemp("pub.txt");
+
+        var procOut = runMain(List.of("keygen",
+                                      "--private-out-file", absPathOf(privKeyFile),
+                                      "--public-out-file",  absPathOf(pubKeyFile)));
+        assertEquals(0, procOut.exitCode());
+        var privKeyPerms  = Files.getPosixFilePermissions(privKeyFile);
+        var expectedPerms = PosixFilePermissions.fromString("rw-------");
+        assertEquals(expectedPerms, privKeyPerms);
+    }
+
+    private static final String TEST_PRIV_KEY     = "4qGcntygFn_a3uqeBa1PbDlygQ-cpOuNznTPIz9ftWE";
+    private static final String TEST_PUB_KEY      = "ROAH_S862tNMpbJ49lu1dPXFCPHFIXZK30pSrMZEmEg";
+    // Token created for the above public key (matching the above private key), using key id 1
+    private static final String TEST_TOKEN        = "AQAAAQAgwyxd7bFNQB_2LdL3bw-xFlvrxXhs7WWNVCKZ4" +
+                                                    "EFeNVtu42JMwM74bMN4E46v6mYcfQNPzcMGaP22Wl2cTnji0A";
+    private static final int    TEST_TOKEN_KEY_ID = 1;
+
+    @Test
+    void encrypt_fails_with_error_message_if_no_input_file_is_given() throws IOException {
+        verifyStderrEquals(List.of("encrypt",
+                                   "--output-file",          "foo",
+                                   "--recipient-public-key", TEST_PUB_KEY,
+                                   "--key-id",               "1234"),
+                "Invalid command line arguments: Expected exactly 1 file argument to encrypt\n");
+    }
+
+    @Test
+    void encrypt_fails_with_error_message_if_input_file_does_not_exist() throws IOException {
+        verifyStderrEquals(List.of("encrypt",
+                                   "no-such-file",
+                                   "--output-file",          "foo",
+                                   "--recipient-public-key", TEST_PUB_KEY,
+                                   "--key-id",               "1234"),
+                "Invalid command line arguments: Cannot encrypt file 'no-such-file' as it does not exist\n");
+    }
+
+    @Test
+    void decrypt_fails_with_error_message_if_no_input_file_is_given() throws IOException {
+        Path privKeyFile = pathInTemp("priv.txt");
+        Files.writeString(privKeyFile, TEST_PRIV_KEY);
+
+        verifyStderrEquals(List.of("decrypt",
+                                   "--output-file",                "foo",
+                                   "--recipient-private-key-file", absPathOf(privKeyFile),
+                                   "--token",                      TEST_TOKEN,
+                                   "--key-id",                     Integer.toString(TEST_TOKEN_KEY_ID)),
+                "Invalid command line arguments: Expected exactly 1 file argument to decrypt\n");
+    }
+
+    @Test
+    void decrypt_fails_with_error_message_if_input_file_does_not_exist() throws IOException {
+        Path privKeyFile = pathInTemp("priv.txt");
+        Files.writeString(privKeyFile, TEST_PRIV_KEY);
+
+        verifyStderrEquals(List.of("decrypt",
+                                   "no-such-file",
+                                   "--output-file",                "foo",
+                                   "--recipient-private-key-file", absPathOf(privKeyFile),
+                                   "--token",                      TEST_TOKEN,
+                                   "--key-id",                     Integer.toString(TEST_TOKEN_KEY_ID)),
+                "Invalid command line arguments: Cannot decrypt file 'no-such-file' as it does not exist\n");
+    }
+
+    @Test
+    void decrypt_fails_with_error_message_if_expected_key_id_does_not_match_key_id_in_token() throws IOException {
+        Path privKeyFile = pathInTemp("priv.txt");
+        Files.writeString(privKeyFile, TEST_PRIV_KEY);
+
+        Path inputFile = pathInTemp("input.txt");
+        Files.writeString(inputFile, "dummy-not-actually-encrypted-data");
+
+        verifyStderrEquals(List.of("decrypt",
+                                   absPathOf(inputFile),
+                                   "--output-file",                "foo",
+                                   "--recipient-private-key-file", absPathOf(privKeyFile),
+                                   "--token",                      TEST_TOKEN,
+                                   "--key-id",                     Integer.toString(TEST_TOKEN_KEY_ID + 1)),
+                "Invalid command line arguments: Key ID specified with --key-id (2) does not match " +
+                        "key ID used when generating the supplied token (1)\n");
+    }
+
+    @Test
+    void can_end_to_end_keygen_encrypt_and_decrypt() throws IOException {
+        String greatSecret = "Dogs can't look up";
+
+        Path secretFile = pathInTemp("secret.txt");
+        Files.writeString(secretFile, greatSecret);
+
+        var privPath = pathInTemp("priv.txt");
+        var pubPath = pathInTemp("pub.txt");
+        var procOut = runMain(List.of(
+                "keygen",
+                "--private-out-file", absPathOf(privPath),
+                "--public-out-file",  absPathOf(pubPath)));
+        assertEquals(0, procOut.exitCode());
+        assertEquals("", procOut.stdOut());
+        assertEquals("", procOut.stdErr());
+
+        assertTrue(privPath.toFile().exists());
+        assertTrue(pubPath.toFile().exists());
+
+        var encryptedPath = pathInTemp("encrypted.bin");
+        // TODO support (and test) public key via file
+        procOut = runMain(List.of(
+                "encrypt",
+                absPathOf(secretFile),
+                "--output-file",          absPathOf(encryptedPath),
+                "--recipient-public-key", Files.readString(pubPath),
+                "--key-id",               "1234"));
+        assertEquals(0, procOut.exitCode());
+        assertEquals("", procOut.stdErr());
+
+        var token = procOut.stdOut();
+        assertFalse(token.isBlank());
+
+        assertTrue(encryptedPath.toFile().exists());
+
+        var decryptedPath = pathInTemp("decrypted.txt");
+        procOut = runMain(List.of(
+                "decrypt",
+                absPathOf(encryptedPath),
+                "--output-file",                absPathOf(decryptedPath),
+                "--recipient-private-key-file", absPathOf(privPath),
+                "--key-id",                     "1234",
+                "--token",                      token
+                ));
+        assertEquals(0, procOut.exitCode());
+        assertEquals("", procOut.stdOut());
+        assertEquals("", procOut.stdErr());
+
+        assertEquals(greatSecret, Files.readString(decryptedPath));
+    }
+
+    private ProcessOutput runMain(List<String> args) {
+        // Expect that this is used for running a command that is not supposed to fail. But if it does,
+        // include exception trace in stderr to make it easier to debug.
+        return runMain(args, Map.of("VESPA_DEBUG", "true"));
+    }
+
+    private ProcessOutput runMain(List<String> args, Map<String, String> env) {
+        var stdOutBytes = new ByteArrayOutputStream();
+        var stdErrBytes = new ByteArrayOutputStream();
+        var stdOut      = new PrintStream(stdOutBytes);
+        var stdError    = new PrintStream(stdErrBytes);
+
+        int exitCode = new Main(stdOut, stdError).execute(args.toArray(new String[0]), env);
+
+        stdOut.flush();
+        stdError.flush();
+
+        return new ProcessOutput(exitCode, stdOutBytes.toString(), stdErrBytes.toString());
+    }
+
+    private static String readTestResource(String fileName) throws IOException {
+        return Files.readString(Paths.get(CryptoToolsTest.class.getResource('/' + fileName).getFile()));
+    }
+
+    private Path pathInTemp(String fileName) {
+        return tmpFolder.toPath().resolve(fileName);
+    }
+
+    private static String absPathOf(Path path) {
+        return path.toAbsolutePath().toString();
+    }
+
+}