aes: add aes_cbc_with_ecdh_key for ephemeral request/response encryption

jgriffiths · jgriffiths · commit 06ce5e709fdc · 2023-09-21T10:52:14.000+12:00
diff --git a/include/wally.hpp b/include/wally.hpp
@@ -401,6 +401,19 @@ inline int aes_cbc_get_maximum_length(const KEY& key, const IV& iv, const BYTES&
     return ret;
 }
 
+template <class PRIV_KEY, class IV, class BYTES, class PUB_KEY, class LABEL, class BYTES_OUT>
+inline int aes_cbc_with_ecdh_key(const PRIV_KEY& priv_key, const IV& iv, const BYTES& bytes, const PUB_KEY& pub_key, const LABEL& label, uint32_t flags, BYTES_OUT& bytes_out, size_t* written = 0) {
+    size_t n;
+    int ret = ::wally_aes_cbc_with_ecdh_key(priv_key.data(), priv_key.size(), iv.data(), iv.size(), bytes.data(), bytes.size(), pub_key.data(), pub_key.size(), label.data(), label.size(), flags, bytes_out.data(), bytes_out.size(), written ? written : &n);
+    return written || ret != WALLY_OK ? ret : n == static_cast<size_t>(bytes_out.size()) ? WALLY_OK : WALLY_EINVAL;
+}
+
+template <class PRIV_KEY, class IV, class BYTES, class PUB_KEY, class LABEL>
+inline int aes_cbc_with_ecdh_key_get_maximum_length(const PRIV_KEY& priv_key, const IV& iv, const BYTES& bytes, const PUB_KEY& pub_key, const LABEL& label, uint32_t flags, size_t* written) {
+    int ret = ::wally_aes_cbc_with_ecdh_key_get_maximum_length(priv_key.data(), priv_key.size(), iv.data(), iv.size(), bytes.data(), bytes.size(), pub_key.data(), pub_key.size(), label.data(), label.size(), flags, written);
+    return ret;
+}
+
 template <class KEY, class BYTES>
 inline int aes_len(const KEY& key, const BYTES& bytes, uint32_t flags, size_t* written) {
     int ret = ::wally_aes_len(key.data(), key.size(), bytes.data(), bytes.size(), flags, written);
diff --git a/include/wally_crypto.h b/include/wally_crypto.h
@@ -891,6 +891,81 @@ WALLY_CORE_API int wally_s2c_commitment_verify(
     size_t s2c_opening_len,
     uint32_t flags);
 
+/**
+ * Get the maximum length of data encrypted/decrypted using `wally_aes_cbc_with_ecdh_key`.
+ *
+ * :param priv_key: The callers private key used for Diffie-Helman exchange.
+ * :param priv_key_len: The length of ``priv_key`` in bytes. Must be `EC_PRIVATE_KEY_LEN`.
+ * :param iv: Initialisation vector. Only required when encrypting, otherwise pass NULL.
+ * :param iv_len: Length of ``iv`` in bytes. Must be `AES_BLOCK_LEN`.
+ * :param bytes: Bytes to encrypt/decrypt.
+ * :param bytes_len: Length of ``bytes`` in bytes.
+ * :param pub_key: The other parties public key used for Diffie-Helman exchange.
+ * :param pub_key_len: Length of ``pub_key`` in bytes. Must be `EC_PUBLIC_KEY_LEN`.
+ * :param label: A non-empty array of bytes for internal key generation. Must
+ *|    be the same (fixed) value when encrypting and decrypting.
+ * :param label_len: Length of ``label`` in bytes.
+ * :param flags: :ref:`aes-operation-flag` indicating the desired behavior.
+ * :param written: Destination for the maximum length of the encrypted/decrypted data.
+ */
+WALLY_CORE_API int wally_aes_cbc_with_ecdh_key_get_maximum_length(
+    const unsigned char *priv_key,
+    size_t priv_key_len,
+    const unsigned char *iv,
+    size_t iv_len,
+    const unsigned char *bytes,
+    size_t bytes_len,
+    const unsigned char *pub_key,
+    size_t pub_key_len,
+    const unsigned char *label,
+    size_t label_len,
+    uint32_t flags,
+    size_t *written);
+
+/**
+ * Encrypt/decrypt data using AES-256 (CBC mode, PKCS#7 padding) and a shared Diffie-Helman secret.
+ *
+ * :param priv_key: The callers private key used for Diffie-Helman exchange.
+ * :param priv_key_len: The length of ``priv_key`` in bytes. Must be `EC_PRIVATE_KEY_LEN`.
+ * :param iv: Initialisation vector. Only required when encrypting, otherwise pass NULL.
+ * :param iv_len: Length of ``iv`` in bytes. Must be `AES_BLOCK_LEN` if encrypting otherwise 0.
+ * :param bytes: Bytes to encrypt/decrypt.
+ * :param bytes_len: Length of ``bytes`` in bytes.
+ * :param pub_key: The other parties public key used for Diffie-Helman exchange.
+ * :param pub_key_len: Length of ``pub_key`` in bytes. Must be `EC_PUBLIC_KEY_LEN`.
+ * :param label: A non-empty array of bytes for internal key generation. Must
+ *|    be the same (fixed) value when encrypting and decrypting.
+ * :param label_len: Length of ``label`` in bytes.
+ * :param flags: :ref:`aes-operation-flag` indicating the desired behavior.
+ * :param bytes_out: Destination for the encrypted/decrypted data.
+ * :param len: The length of ``bytes_out`` in bytes.
+ * :param written: Destination for the number of bytes written to ``bytes_out``.
+ *
+ * This function implements a scheme for sharing data using a derived secret.
+ * Alice creates an ephemeral key pair and sends her public key to Bob along
+ * with any request details. Bob creates an ephemeral key pair and calls this
+ * function with his private key and Alices public key to encrypt ``bytes``
+ * (the request payload). Bob returns his public key and the encrypted data to
+ * Alice, who calls this function with her private key and Bobs public key
+ * to decrypt and authenticate the payload. The ``label`` parameter must be
+ * be the same for both Alice and Bob for a given request/response.
+ */
+WALLY_CORE_API int wally_aes_cbc_with_ecdh_key(
+    const unsigned char *priv_key,
+    size_t priv_key_len,
+    const unsigned char *iv,
+    size_t iv_len,
+    const unsigned char *bytes,
+    size_t bytes_len,
+    const unsigned char *pub_key,
+    size_t pub_key_len,
+    const unsigned char *label,
+    size_t label_len,
+    uint32_t flags,
+    unsigned char *bytes_out,
+    size_t len,
+    size_t *written);
+
 #ifdef __cplusplus
 }
 #endif
diff --git a/src/Makefile.am b/src/Makefile.am
@@ -319,6 +319,7 @@ endif
 
 if USE_SWIG_PYTHON
 check-swig-python: $(SWIG_PYTHON_TEST_DEPS)
+	$(AM_V_at)$(PYTHON_SWIGTEST) swig_python/contrib/aes.py
 	$(AM_V_at)$(PYTHON_SWIGTEST) swig_python/contrib/bip32.py
 	$(AM_V_at)$(PYTHON_SWIGTEST) swig_python/contrib/coinselection.py
 	$(AM_V_at)$(PYTHON_SWIGTEST) swig_python/contrib/descriptor.py
diff --git a/src/aes.c b/src/aes.c
@@ -1,6 +1,7 @@
 #include "internal.h"
 #include <include/wally_crypto.h>
 
+#include "ccan/ccan/build_assert/build_assert.h"
 #include "ctaes/ctaes.h"
 #include "ctaes/ctaes.c"
 
@@ -208,3 +209,132 @@ int wally_aes_cbc(const unsigned char *key, size_t key_len,
     wally_clear_2(buf, sizeof(buf), &ctx, sizeof(ctx));
     return WALLY_OK;
 }
+
+int wally_aes_cbc_with_ecdh_key_get_maximum_length(
+    const unsigned char *priv_key, size_t priv_key_len,
+    const unsigned char *iv, size_t iv_len,
+    const unsigned char *bytes, size_t bytes_len,
+    const unsigned char *pub_key, size_t pub_key_len,
+    const unsigned char *label, size_t label_len, uint32_t flags,
+    size_t *written)
+{
+    if (written)
+        *written = 0;
+
+    if (!priv_key || priv_key_len != EC_PRIVATE_KEY_LEN || !bytes ||
+        !pub_key || pub_key_len != EC_PUBLIC_KEY_LEN || !label || !label_len ||
+        (flags != AES_FLAG_ENCRYPT && flags != AES_FLAG_DECRYPT) || !written)
+        return WALLY_EINVAL;
+
+    if (flags & AES_FLAG_ENCRYPT) {
+        /* Must provide IV + minimum 1 byte of payload for encryption */
+        if (!iv || iv_len != AES_BLOCK_LEN || !bytes_len)
+            return WALLY_EINVAL;
+        /* Output is IV + encrypted payload + HMAC */
+        *written = AES_BLOCK_LEN + HMAC_SHA256_LEN
+                   + ((bytes_len / AES_BLOCK_LEN) + 1) * AES_BLOCK_LEN;
+    } else {
+        /* Must not provide an IV for decryption */
+        if (iv || iv_len)
+            return WALLY_EINVAL;
+         /* Payload must contain IV, payload and the HMAC for decryption */
+        if (bytes_len < AES_BLOCK_LEN + AES_BLOCK_LEN + HMAC_SHA256_LEN)
+            return WALLY_EINVAL;
+        /* Output is the decrypted payload without the IV and HMAC */
+        bytes_len = bytes_len - AES_BLOCK_LEN - HMAC_SHA256_LEN;
+        if (bytes_len % AES_BLOCK_LEN)
+            return WALLY_EINVAL; /* Payload isn't a block size multiple */
+        /* Actual bytes written may be less due to padding, but the
+         * caller must pass a buffer of the padded size. */
+        *written = bytes_len;
+    }
+    return WALLY_OK;
+}
+
+int wally_aes_cbc_with_ecdh_key(
+    const unsigned char *priv_key, size_t priv_key_len,
+    const unsigned char *iv, size_t iv_len,
+    const unsigned char *bytes, size_t bytes_len,
+    const unsigned char *pub_key, size_t pub_key_len,
+    const unsigned char *label, size_t label_len, uint32_t flags,
+    unsigned char *bytes_out, size_t len, size_t *written)
+{
+    unsigned char secret[SHA256_LEN], keys[HMAC_SHA512_LEN];
+    const unsigned char *enc_key = keys, *hmac_key = keys + AES_KEY_LEN_256;
+    size_t expected_len;
+    const bool is_encrypt = flags & AES_FLAG_ENCRYPT;
+    int ret;
+
+    /* Derived key sizes must match the derived keys buffer */
+    BUILD_ASSERT(sizeof(keys) == AES_KEY_LEN_256 + SHA256_LEN);
+
+    if (written)
+        *written = 0;
+
+    if (!bytes_out || !len || !written)
+        return WALLY_EINVAL;
+    ret = wally_aes_cbc_with_ecdh_key_get_maximum_length(
+                  priv_key, priv_key_len, iv, iv_len, bytes, bytes_len,
+                  pub_key, pub_key_len, label, label_len, flags,
+                  &expected_len);
+    if (ret == WALLY_OK && expected_len > len) {
+        *written = expected_len;
+        return ret; /* Tell the caller how much space is needed */
+    }
+    if (ret != WALLY_OK)
+        return ret;
+
+    if (is_encrypt) {
+        /* Copy the IV to the start of the encrypted output */
+        memcpy(bytes_out, iv, iv_len);
+    } else {
+        /* The IV is the first AES_BLOCK_LEN bytes of the payload */
+        iv = bytes;
+        iv_len = AES_BLOCK_LEN;
+    }
+
+    /* Shared secret is ECDH(their pubkey, our private key) */
+    ret = wally_ecdh(pub_key, pub_key_len, priv_key, priv_key_len,
+                     secret, sizeof(secret));
+    /* Generate encryption/HMAC keys using the shared secret and label */
+    if (ret == WALLY_OK)
+        ret = wally_hmac_sha512(secret, sizeof(secret), label, label_len, keys, sizeof(keys));
+    if (ret == WALLY_OK && !is_encrypt) {
+        /* Verify the IV + encrypted data's HMAC */
+        unsigned char hmac[HMAC_SHA256_LEN];
+        ret = wally_hmac_sha256(hmac_key, SHA256_LEN,
+                                bytes, bytes_len - sizeof(hmac),
+                                hmac, sizeof(hmac));
+        if (ret == WALLY_OK &&
+            memcmp(hmac, bytes + bytes_len - sizeof(hmac), sizeof(hmac)))
+            ret = WALLY_EINVAL; /* Invalid HMAC */
+    }
+
+    /* Encrypt/decrypt the payload */
+    if (ret == WALLY_OK) {
+        /* Trim our output length to a block size multiple for aes_cbc */
+        size_t out_len = (len - (is_encrypt ? (iv_len + HMAC_SHA256_LEN) : 0)) & ~((size_t)0xf);
+        if (is_encrypt)
+            ret = wally_aes_cbc(enc_key, AES_KEY_LEN_256, iv, iv_len,
+                                bytes, bytes_len, flags,
+                                bytes_out + iv_len, out_len, written);
+        else
+            ret = wally_aes_cbc(enc_key, AES_KEY_LEN_256, iv, iv_len,
+                                bytes + iv_len, bytes_len - iv_len - HMAC_SHA256_LEN,
+                                flags, bytes_out, out_len, written);
+    }
+
+    if (ret == WALLY_OK && is_encrypt) {
+        /* append the HMAC of the IV + encrypted data */
+        *written += AES_BLOCK_LEN; /* Include the IV */
+        ret = wally_hmac_sha256(hmac_key, SHA256_LEN,
+                                bytes_out, *written,
+                                bytes_out + *written, HMAC_SHA256_LEN);
+        *written = ret == WALLY_OK ? *written + HMAC_SHA256_LEN : 0;
+    }
+
+    if (ret == WALLY_OK && *written > expected_len)
+        ret = WALLY_ERROR; /* Should never happen! */
+    wally_clear_2(secret, sizeof(secret), keys, sizeof(keys));
+    return ret;
+}
diff --git a/src/swig_java/swig.i b/src/swig_java/swig.i
@@ -523,6 +523,8 @@ static jobjectArray create_jstringArray(JNIEnv *jenv, char **p, size_t len) {
 %returns_size_t(wally_aes_len);
 %returns_size_t(wally_aes_cbc);
 %returns_size_t(wally_aes_cbc_get_maximum_length);
+%returns_size_t(wally_aes_cbc_with_ecdh_key);
+%returns_size_t(wally_aes_cbc_with_ecdh_key_get_maximum_length);
 %returns_array_(wally_asset_final_vbf, 8, 9, ASSET_TAG_LEN);
 %returns_array_(wally_asset_generator_from_bytes, 5, 6, ASSET_GENERATOR_LEN);
 %returns_size_t(wally_asset_rangeproof_get_maximum_len);
diff --git a/src/swig_python/contrib/aes.py b/src/swig_python/contrib/aes.py
@@ -0,0 +1,43 @@
+import unittest
+from wallycore import init, cleanup, \
+    ec_public_key_from_private_key, aes_cbc_with_ecdh_key, \
+    AES_FLAG_ENCRYPT, AES_FLAG_DECRYPT
+
+
+class AESTests(unittest.TestCase):
+
+    def test_aes_cbc_ecdh(self):
+        """Test python wrappers for aes_cbc_with_ecdh_key """
+        # Alice generates an ephemeral keypair for her request.
+        alice_priv = bytes.fromhex('1c6a837d1ac663fdc7f1002327ca38452766eaf4fe3b80ce620bf7cd3f584cf6')
+        alice_pub = ec_public_key_from_private_key(alice_priv)
+        # Bob generates an ephemeral keypair for his response.
+        bob_priv = bytes.fromhex('0b6b3dc90d203d854100110788ac87d43aa00620c9cdb361b281b09022ef4b53')
+        bob_pub = ec_public_key_from_private_key(bob_priv)
+        # Bob also generates a secure random IV for encrypting the response.
+        iv = bytes.fromhex('bd5d4724243880738e7e8b0c02658700')
+        # Both parties must agree on a shared label to use.
+        label = 'a sample label'.encode()
+        # This is the example payload we are using.
+        payload = 'This is an example response/payload to encrypt'.encode()
+
+        # Test we handle messages up to and over AES_BLOCK_LEN boundaries
+        for i in range(1, len(payload) + 1):
+            # The protocol:
+            # 1) Alice requests some data using her pubkey.
+            # 2) Bob encrypts the response (payload) with his private key,
+            #    a random IV, a shared label and Alice's pubkey.
+            encrypted = aes_cbc_with_ecdh_key(bob_priv, iv, payload[0:i],
+                                              alice_pub, label, AES_FLAG_ENCRYPT)
+            # 3) Bob sends the encrypted data and his pubkey to Alice.
+            # 4) Alice decrypts the payload with her private key and Bobs pubkey.
+            decrypted = aes_cbc_with_ecdh_key(alice_priv, None, encrypted,
+                                              bob_pub, label, AES_FLAG_DECRYPT)
+            # Alice now has the unencrypted payload.
+            self.assertEqual(decrypted.hex(), payload[0:i].hex())
+
+
+if __name__ == '__main__':
+    init(0)
+    unittest.main()
+    cleanup(0)
diff --git a/src/swig_python/python_extra.py_in b/src/swig_python/python_extra.py_in
@@ -116,6 +116,7 @@ ae_sig_from_bytes = _wrap_bin(ae_sig_from_bytes, EC_SIGNATURE_LEN)
 ae_signer_commit_from_bytes = _wrap_bin(ae_signer_commit_from_bytes, WALLY_S2C_OPENING_LEN)
 aes = _wrap_bin(aes, aes_len)
 aes_cbc = _wrap_bin(aes_cbc, aes_cbc_get_maximum_length, resize=True)
+aes_cbc_with_ecdh_key = _wrap_bin(aes_cbc_with_ecdh_key, aes_cbc_with_ecdh_key_get_maximum_length, resize=True)
 base58_n_to_bytes = _wrap_bin(base58_n_to_bytes, base58_n_to_bytes_len, resize=True)
 base58_to_bytes = _wrap_bin(base58_to_bytes, base58_to_bytes_len, resize=True)
 base64_to_bytes = _wrap_bin(base64_to_bytes, base64_get_maximum_length, resize=True)
diff --git a/src/test/test_aes.py b/src/test/test_aes.py
@@ -92,6 +92,85 @@ def test_aes_cbc(self):
                 self.assertEqual((ret, written), (0, len(o)))
                 self.assertEqual(h(out_buf), h(o))
 
+    def test_aes_cbc_with_ecdh_key(self):
+        ENCRYPT, DECRYPT, PUBKEY_LEN, _ = 1, 2, 33, True
+        a_priv = make_cbuffer('1c6a837d1ac663fdc7f1002327ca38452766eaf4fe3b80ce620bf7cd3f584cf6')[0]
+        a_pub = make_cbuffer('03e581be89d1ef8ce11d60746d08e4f8aedf934d1d861dd436042ee2e3b16db918')[0]
+        b_priv = make_cbuffer('0b6b3dc90d203d854100110788ac87d43aa00620c9cdb361b281b09022ef4b53')[0]
+        b_pub = make_cbuffer('03ff06999ad61c0f3a733b93fc1e6b75ecfb1439b326e840de590a56454f0eeb0d')[0]
+        iv = make_cbuffer('bd5d4724243880738e7e8b0c02658700')[0]
+        label = 'a sample label'.encode()
+        payload = 'This is an example response/payload to encrypt'.encode()
+        buf = make_cbuffer('00' * 256)[0]
+
+        # Encryption
+        good_args = [b_priv, len(b_priv), iv, len(iv), payload, len(payload),
+                     a_pub, len(a_pub), label, len(label), ENCRYPT, buf, len(buf)]
+
+        ret, written = wally_aes_cbc_with_ecdh_key(*good_args)
+        self.assertEqual(ret, WALLY_OK) # Make sure good args work
+        encrypted = make_cbuffer(buf[:written].hex())[0]
+
+        invalid_cases = [
+            (None, _, _,    _, _,    _, _,    _, _,    _, _,    _, _), # NULL privkey
+            (_,    0, _,    _, _,    _, _,    _, _,    _, _,    _, _), # Empty privkey
+            (_,    9, _,    _, _,    _, _,    _, _,    _, _,    _, _), # Wrong privkey length
+            (_,    _, None, _, _,    _, _,    _, _,    _, _,    _, _), # NULL IV (enc)
+            (_,    _, _,    0, _,    _, _,    _, _,    _, _,    _, _), # Empty IV (enc)
+            (_,    _, _,    9, _,    _, _,    _, _,    _, _,    _, _), # Wrong IV length (enc)
+            (_,    _, _,    _, None, _, _,    _, _,    _, _,    _, _), # NULL payload
+            (_,    _, _,    _, _,    0, _,    _, _,    _, _,    _, _), # Empty payload
+            (_,    _, _,    _, _,    _, None, _, _,    _, _,    _, _), # NULL pubkey
+            (_,    _, _,    _, _,    _, _,    0, _,    _, _,    _, _), # Empty pubkey
+            (_,    _, _,    _, _,    _, _,    9, _,    _, _,    _, _), # Wrong pubkey length
+            (_,    _, _,    _, _,    _, _,    _, None, _, _,    _, _), # NULL label
+            (_,    _, _,    _, _,    _, _,    _, _,    0, _,    _, _), # Empty label
+            (_,    _, _,    _, _,    _, _,    _, _,    _, 3,    _, _), # Encrypt+Decrypt flags
+            (_,    _, _,    _, _,    _, _,    _, _,    _, 5,    _, _), # Unknown flag
+            (_,    _, _,    _, _,    _, _,    _, _,    _, _, None, _), # NULL output
+            (_,    _, _,    _, _,    _, _,    _, _,    _, _, _,    0), # Zero-length output
+        ]
+        for case in invalid_cases:
+            args = [good_args[i] if a == _ else a for i, a in enumerate(case)]
+            self.assertEqual(wally_aes_cbc_with_ecdh_key(*args), (WALLY_EINVAL, 0))
+
+        # Test writing up to/beyond the output buffer size
+        for out_len in range(1, len(encrypted) + 16):
+            args = [a for a in good_args]
+            args[-1] = out_len
+            ret = wally_aes_cbc_with_ecdh_key(*args)
+            self.assertEqual(ret, (WALLY_OK, written)) # returns required length
+
+        # Decryption
+        good_args = [a_priv, len(a_priv), None, 0, encrypted, len(encrypted),
+                     b_pub, len(b_pub), label, len(label), DECRYPT, buf, len(buf)]
+
+        ret, written = wally_aes_cbc_with_ecdh_key(*good_args)
+        self.assertEqual(ret, WALLY_OK) # Make sure good args work
+        self.assertEqual(buf[:written], payload)
+
+        bad = make_cbuffer((encrypted[:-1] + b'?').hex())[0] # Corrupt the HMAC
+        bad_len = len(encrypted) - 1
+        invalid_cases = [
+            (_, _, iv, _,       _,    _,      _, _, _, _, _, _, _), # Non-NULL IV (dec)
+            (_, _, _,  len(iv), _,    _,      _, _, _, _, _, _, _), # Non-zero IV length (dec)
+            (_, _, _,  _,       bad, _,       _, _, _, _, _, _, _), # Corrupt HMAC
+            (_, _, _,  _,       _,   bad_len, _, _, _, _, _, _, _), # Truncated encrypted data
+        ]
+        for case in invalid_cases:
+            args = [good_args[i] if a == _ else a for i, a in enumerate(case)]
+            self.assertEqual(wally_aes_cbc_with_ecdh_key(*args), (WALLY_EINVAL, 0))
+
+        # Test writing up to/beyond the output buffer size
+        for out_len in range(1, len(payload) + 16):
+            args = [a for a in good_args]
+            args[-1] = out_len
+            ret, written = wally_aes_cbc_with_ecdh_key(*args)
+            self.assertEqual(ret, WALLY_OK)
+            # The output size required includes final padding which is
+            # stripped if the payload isn't a multiple of the AES block size.
+            self.assertLessEqual(len(payload), written)
+
 
 if __name__ == '__main__':
     unittest.main()
diff --git a/src/test/util.py b/src/test/util.py
diff --git a/src/wasm_package/src/functions.js b/src/wasm_package/src/functions.js
diff --git a/src/wasm_package/src/index.d.ts b/src/wasm_package/src/index.d.ts
diff --git a/tools/wasm_exports.sh b/tools/wasm_exports.sh
