From 88cab67cf83d2e969a85feab8e871bc25bb52d1f Mon Sep 17 00:00:00 2001
From: Alex Parrill <alex.parrill@col32.net>
Date: Thu, 7 Apr 2022 17:44:20 -0400
Subject: [PATCH] Add StatelessTransportState::initiator_key and responder_key

Allows for saving and restoring of keys across re-keys.

This can be useful, for example, if you are working with UDP and
need to decrypt a packet that arrived after a rekey but was sent
before the rekey.
---
 src/cipherstate.rs              |  4 ++++
 src/resolvers/default.rs        |  8 ++++++++
 src/stateless_transportstate.rs | 18 ++++++++++++++++++
 src/types.rs                    |  3 +++
 4 files changed, 33 insertions(+)

diff --git a/src/cipherstate.rs b/src/cipherstate.rs
index 6ef741a..5597233 100644
--- a/src/cipherstate.rs
+++ b/src/cipherstate.rs
@@ -167,6 +167,10 @@ impl StatelessCipherState {
     pub fn rekey_manually(&mut self, key: &[u8]) {
         self.cipher.set(key);
     }
+
+    pub fn key(&self) -> &[u8] {
+        self.cipher.get()
+    }
 }
 
 impl From<CipherState> for StatelessCipherState {
diff --git a/src/resolvers/default.rs b/src/resolvers/default.rs
index 8c52a26..d5f8154 100644
--- a/src/resolvers/default.rs
+++ b/src/resolvers/default.rs
@@ -195,6 +195,10 @@ impl Cipher for CipherAesGcm {
         copy_slices!(key, &mut self.key)
     }
 
+    fn get(&self) -> &[u8] {
+        &self.key
+    }
+
     fn encrypt(&self, nonce: u64, authtext: &[u8], plaintext: &[u8], out: &mut [u8]) -> usize {
         let aead = aes_gcm::Aes256Gcm::new(&self.key.into());
 
@@ -248,6 +252,10 @@ impl Cipher for CipherChaChaPoly {
         copy_slices!(key, &mut self.key);
     }
 
+    fn get(&self) -> &[u8] {
+        &self.key
+    }
+
     fn encrypt(&self, nonce: u64, authtext: &[u8], plaintext: &[u8], out: &mut [u8]) -> usize {
         let mut nonce_bytes = [0u8; 12];
         copy_slices!(nonce.to_le_bytes(), &mut nonce_bytes[4..]);
diff --git a/src/stateless_transportstate.rs b/src/stateless_transportstate.rs
index 6ad7f7c..230fb4e 100644
--- a/src/stateless_transportstate.rs
+++ b/src/stateless_transportstate.rs
@@ -140,6 +140,24 @@ impl StatelessTransportState {
     pub fn is_initiator(&self) -> bool {
         self.initiator
     }
+
+    /// Gets the initiator's current symmetric key.
+    ///
+    /// This can be passed to `rekey_initiator_manually` to restore the state to
+    /// before a reykeying, in case you need to work with messages that have come in
+    /// before then.
+    pub fn initiator_key(&self) -> &[u8] {
+        self.cipherstates.0.key()
+    }
+
+    /// Gets the responder's current symmetric key.
+    ///
+    /// This can be passed to `rekey_responder_manually` to restore the state to
+    /// before a reykeying, in case you need to work with messages that have come in
+    /// before then.
+    pub fn responder_key(&self) -> &[u8] {
+        self.cipherstates.1.key()
+    }
 }
 
 impl fmt::Debug for StatelessTransportState {
diff --git a/src/types.rs b/src/types.rs
index 3aa3bfb..e2f5720 100644
--- a/src/types.rs
+++ b/src/types.rs
@@ -44,6 +44,9 @@ pub trait Cipher: Send + Sync {
     /// Set the key
     fn set(&mut self, key: &[u8]);
 
+    /// Get the key
+    fn get(&self) -> &[u8];
+
     /// Encrypt (with associated data) a given plaintext.
     fn encrypt(&self, nonce: u64, authtext: &[u8], plaintext: &[u8], out: &mut [u8]) -> usize;