From e09c8dd51254abaa50c9bba8b428fd4ffef5f0bd Mon Sep 17 00:00:00 2001
From: rcombs <rcombs@rcombs.me>
Date: Wed, 22 May 2024 17:01:24 -0700
Subject: [PATCH] Support custom signing routines in SignaturePrivateKey

This enables support for signing with non-extractable private keys (e.g. TPM-bound).
---
 include/mls/crypto.h |  9 ++++++++-
 src/crypto.cpp       | 24 +++++++++++++++++++++++-
 2 files changed, 31 insertions(+), 2 deletions(-)

diff --git a/include/mls/crypto.h b/include/mls/crypto.h
index 16aefaab..d266f9b9 100644
--- a/include/mls/crypto.h
+++ b/include/mls/crypto.h
@@ -240,11 +240,14 @@ struct PublicJWK
 
 struct SignaturePrivateKey
 {
+  using SignerFunc = std::function<bytes(const std::vector<uint8_t>&)>;
+
   static SignaturePrivateKey generate(CipherSuite suite);
   static SignaturePrivateKey parse(CipherSuite suite, const bytes& data);
   static SignaturePrivateKey derive(CipherSuite suite, const bytes& secret);
   static SignaturePrivateKey from_jwk(CipherSuite suite,
                                       const std::string& json_str);
+  static SignaturePrivateKey from_func(SignerFunc func, bytes pub_data);
 
   SignaturePrivateKey() = default;
 
@@ -261,7 +264,11 @@ struct SignaturePrivateKey
   TLS_SERIALIZABLE(data)
 
 private:
-  SignaturePrivateKey(bytes priv_data, bytes pub_data);
+  SignerFunc _sign_func;
+
+  SignaturePrivateKey(bytes priv_data,
+                      bytes pub_data,
+                      SignerFunc func = SignerFunc());
 };
 
 } // namespace MLS_NAMESPACE
diff --git a/src/crypto.cpp b/src/crypto.cpp
index 8d859c0e..fcaf471a 100644
--- a/src/crypto.cpp
+++ b/src/crypto.cpp
@@ -461,19 +461,31 @@ SignaturePrivateKey::sign(const CipherSuite& suite,
 {
   auto label_plus = mls_1_0_plus(label);
   const auto content = tls::marshal(SignContent{ label_plus, message });
+
+  if (_sign_func) {
+    return _sign_func(content);
+  }
+
   const auto priv = suite.sig().deserialize_private(data);
   return suite.sig().sign(content, *priv);
 }
 
-SignaturePrivateKey::SignaturePrivateKey(bytes priv_data, bytes pub_data)
+SignaturePrivateKey::SignaturePrivateKey(bytes priv_data,
+                                         bytes pub_data,
+                                         SignerFunc func)
   : data(std::move(priv_data))
   , public_key{ std::move(pub_data) }
+  , _sign_func{ std::move(func) }
 {
 }
 
 void
 SignaturePrivateKey::set_public_key(CipherSuite suite)
 {
+  if (_sign_func) {
+    throw std::runtime_error("not implemented");
+  }
+
   const auto priv = suite.sig().deserialize_private(data);
   auto pub = priv->public_key();
   public_key.data = suite.sig().serialize(*pub);
@@ -489,9 +501,19 @@ SignaturePrivateKey::from_jwk(CipherSuite suite, const std::string& json_str)
   return { priv_data, pub_data };
 }
 
+SignaturePrivateKey
+SignaturePrivateKey::from_func(SignerFunc func, bytes pub_data)
+{
+  return { bytes(), std::move(pub_data), std::move(func) };
+}
+
 std::string
 SignaturePrivateKey::to_jwk(CipherSuite suite) const
 {
+  if (_sign_func) {
+    throw std::runtime_error("not implemented");
+  }
+
   const auto priv = suite.sig().deserialize_private(data);
   return suite.sig().export_jwk_private(*priv);
 }