From 1b0bb351049644d683ddd5f0693eb5fd1003f031 Mon Sep 17 00:00:00 2001
From: Santiago Carmuega <santiago@carmuega.me>
Date: Tue, 8 Aug 2023 18:09:37 -0300
Subject: [PATCH] Allow external private key in v2

---
 src/common.rs   |  2 +-
 src/lib.rs      |  4 ++--
 src/version2.rs | 37 ++++++++++++++++++++++++-------------
 3 files changed, 27 insertions(+), 16 deletions(-)

diff --git a/src/common.rs b/src/common.rs
index b95583f..953bd72 100644
--- a/src/common.rs
+++ b/src/common.rs
@@ -8,7 +8,7 @@ use ct_codecs::{Base64UrlSafeNoPadding, Decoder, Encoder};
 use subtle::ConstantTimeEq;
 
 /// Encode bytes with Base64 URL-safe and no padding.
-pub(crate) fn encode_b64<T: AsRef<[u8]>>(bytes: T) -> Result<String, Error> {
+pub fn encode_b64<T: AsRef<[u8]>>(bytes: T) -> Result<String, Error> {
     let inlen = bytes.as_ref().len();
     let mut buf = vec![0u8; Base64UrlSafeNoPadding::encoded_len(inlen)?];
 
diff --git a/src/lib.rs b/src/lib.rs
index eb29a71..b36ff2d 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -189,12 +189,12 @@
 #[macro_use]
 extern crate alloc;
 
-mod pae;
+pub mod pae;
 
 /// Errors for token operations.
 pub mod errors;
 
-mod common;
+pub mod common;
 
 #[cfg(feature = "std")]
 /// Claims for tokens and validation thereof.
diff --git a/src/version2.rs b/src/version2.rs
index f2a5c9a..8105a83 100644
--- a/src/version2.rs
+++ b/src/version2.rs
@@ -108,33 +108,44 @@ impl PublicToken {
     /// The header and purpose for the public token: `v2.public.`.
     pub const HEADER: &'static str = "v2.public.";
 
-    /// Create a public token.
-    pub fn sign(
-        secret_key: &AsymmetricSecretKey<V2>,
-        message: &[u8],
-        footer: Option<&[u8]>,
-    ) -> Result<String, Error> {
+    /// Pre-authenticated encoding
+    pub fn pae(message: &[u8], footer: &[u8]) -> Result<Vec<u8>, Error> {
         if message.is_empty() {
             return Err(Error::EmptyPayload);
         }
 
-        let sk = SigningKey::from_slice(secret_key.as_bytes()).map_err(|_| Error::Key)?;
-        let f = footer.unwrap_or(&[]);
-        let m2 = pae::pae(&[Self::HEADER.as_bytes(), message, f])?;
-        let sig = sk.sign(m2, None);
+        pae::pae(&[Self::HEADER.as_bytes(), message, footer])
+    }
 
+    /// Concatenate message with signature
+    pub fn concatenate(message: &[u8], sig: &[u8], footer: &[u8]) -> Result<String, Error> {
         let mut m_sig: Vec<u8> = Vec::from(message);
-        m_sig.extend_from_slice(sig.as_ref());
+        m_sig.extend_from_slice(sig);
 
         let token_no_footer = format!("{}{}", Self::HEADER, encode_b64(m_sig)?);
 
-        if f.is_empty() {
+        if footer.is_empty() {
             Ok(token_no_footer)
         } else {
-            Ok(format!("{}.{}", token_no_footer, encode_b64(f)?))
+            Ok(format!("{}.{}", token_no_footer, encode_b64(footer)?))
         }
     }
 
+    /// Create a public token.
+    pub fn sign(
+        secret_key: &AsymmetricSecretKey<V2>,
+        message: &[u8],
+        footer: Option<&[u8]>,
+    ) -> Result<String, Error> {
+        let f = footer.unwrap_or(&[]);
+        let m2 = Self::pae(message, f)?;
+
+        let sk = SigningKey::from_slice(secret_key.as_bytes()).map_err(|_| Error::Key)?;
+        let sig = sk.sign(m2, None);
+
+        Self::concatenate(message, sig.as_ref(), f)
+    }
+
     /// Verify a public token.
     ///
     /// If `footer.is_none()`, then it will be validated but not compared to a known value.