From f76597faf6a93bbe641a4ffe92dee760018cacb5 Mon Sep 17 00:00:00 2001
From: Jack Lloyd <jack@signal.org>
Date: Tue, 7 Jul 2020 19:52:17 -0400
Subject: [PATCH] Have ChainKey/RootKey::new take inputs by slice instead of
 array

Since much of the time we are actually using a Vec due to protobuf
and it is inconvenient to have to array_ref! everywhere
---
 src/ratchet/keys.rs | 44 +++++++++++++++++++++++++++++++++++++++-----
 1 file changed, 39 insertions(+), 5 deletions(-)

diff --git a/src/ratchet/keys.rs b/src/ratchet/keys.rs
index 539d62a2..73d1169d 100644
--- a/src/ratchet/keys.rs
+++ b/src/ratchet/keys.rs
@@ -4,7 +4,7 @@ use sha2::Sha256;
 
 use crate::curve;
 use crate::kdf::HKDF;
-use crate::error::Result;
+use crate::error::{SignalProtocolError, Result};
 
 pub struct MessageKeys {
     cipher_key: [u8; 32],
@@ -24,6 +24,25 @@ impl MessageKeys {
         }
     }
 
+    pub fn new(cipher_key: &[u8], mac_key: &[u8], iv: &[u8], counter: u32) -> Result<Self> {
+        if cipher_key.len() != 32 {
+            return Err(SignalProtocolError::InvalidCipherKeyLength(cipher_key.len()));
+        }
+        if mac_key.len() != 32 {
+            return Err(SignalProtocolError::InvalidMacKeyLength(mac_key.len()));
+        }
+        if iv.len() != 16 {
+            return Err(SignalProtocolError::InvalidCipherNonceLength(iv.len()));
+        }
+
+        Ok(MessageKeys {
+            cipher_key: *array_ref![cipher_key, 0, 32],
+            mac_key: *array_ref![mac_key, 0, 32],
+            iv: *array_ref![iv, 0, 16],
+            counter,
+        })
+    }
+
     #[inline]
     pub fn cipher_key(&self) -> &[u8; 32] {
         &self.cipher_key
@@ -44,6 +63,7 @@ impl MessageKeys {
         self.counter
     }
 }
+
 pub struct ChainKey {
     kdf: HKDF,
     key: [u8; 32],
@@ -54,8 +74,16 @@ impl ChainKey {
     const MESSAGE_KEY_SEED: [u8; 1] = [0x01u8];
     const CHAIN_KEY_SEED: [u8; 1] = [0x02u8];
 
-    pub fn new(kdf: HKDF, key: [u8; 32], index: u32) -> Self {
-        Self { kdf, key, index }
+    pub fn new(kdf: HKDF, key: &[u8], index: u32) -> Result<Self> {
+        if key.len() != 32 {
+            return Err(SignalProtocolError::InvalidChainKeyLength(key.len()));
+        }
+
+        Ok(Self {
+            kdf,
+            key: *array_ref![key, 0, 32],
+            index,
+        })
     }
 
     #[inline]
@@ -98,8 +126,14 @@ pub struct RootKey {
 }
 
 impl RootKey {
-    pub fn new(kdf: HKDF, key: [u8; 32]) -> Self {
-        Self { kdf, key }
+    pub fn new(kdf: HKDF, key: &[u8]) -> Result<Self> {
+        if key.len() != 32 {
+            return Err(SignalProtocolError::InvalidRootKeyLength(key.len()));
+        }
+        Ok(Self {
+            kdf,
+            key: *array_ref![key, 0, 32],
+        })
     }
 
     pub fn key(&self) -> &[u8; 32] {