fixup: Add OurPeerStorage for serialized Peer Storage backups

Author: adi2011

lightning/src/ln/our_peer_storage.rs

@@ -7,10 +7,9 @@
 // You may not use this file except in accordance with one or both of these
 // licenses.
 
-//! `OurPeerStorage` enables storage of encrypted serialized channel data.
-//! It provides encryption and decryption of data to maintain data integrity and
+//! `DecryptedOurPeerStorage` enables storage of encrypted serialized channel data.
+//! It provides encryption of data to maintain data integrity and
 //! security during transmission.
-//!
 
 use bitcoin::hashes::sha256::Hash as Sha256;
 use bitcoin::hashes::{Hash, HashEngine, Hmac, HmacEngine};
@@ -20,113 +19,111 @@ use crate::sign::PeerStorageKey;
 use crate::crypto::chacha20poly1305rfc::ChaCha20Poly1305RFC;
 use crate::prelude::*;
 
-/// [`OurPeerStorage`] is used to store encrypted channel information that allows for the creation of a
+/// [`DecryptedOurPeerStorage`] is used to store serialised channel information that allows for the creation of a
 /// `peer_storage` backup.
 ///
 /// This structure is designed to serialize channel data for backup and supports encryption
-/// and decryption using `ChaCha20Poly1305RFC` for transmission.
+/// using `ChaCha20Poly1305RFC` for transmission.
 ///
 /// # Key Methods
-/// - [`OurPeerStorage::new`]: Returns [`OurPeerStorage`] with the given encrypted_data.
-/// - [`OurPeerStorage::create_from_data`]: Returns [`OurPeerStorage`] created from encrypting the provided data.
-/// - [`OurPeerStorage::decrypt_our_peer_storage`]: Decrypts the [`OurPeerStorage::encrypted_data`] using the key and returns decrypted data.
+/// - [`DecryptedOurPeerStorage::new`]: Returns [`DecryptedOurPeerStorage`] with the given data.
+/// - [`DecryptedOurPeerStorage::encrypt`]: Returns [`EncryptedOurPeerStorage`] created from encrypting the provided data.
+/// - [`DecryptedOurPeerStorage::into_vec`]: Returns the data in [`Vec<u8>`] format.
 ///
 /// ## Example
 /// ```
-/// use lightning::ln::our_peer_storage::OurPeerStorage;
+/// use lightning::ln::our_peer_storage::DecryptedOurPeerStorage;
 /// use lightning::sign::PeerStorageKey;
 /// let key = PeerStorageKey{inner: [0u8; 32]};
-/// let decrypted_ops = OurPeerStorage::DecryptedPeerStorage { data: vec![1, 2, 3]};
-/// let our_peer_storage = decrypted_ops.encrypt_peer_storage(&key, [0u8; 32]);
-/// let decrypted_data = our_peer_storage.decrypt_peer_storage(&key).unwrap();
+/// let decrypted_ops = DecryptedOurPeerStorage::new(vec![1, 2, 3]);
+/// let our_peer_storage = decrypted_ops.encrypt(&key, [0u8; 32]);
+/// let decrypted_data = our_peer_storage.decrypt(&key).unwrap();
 /// assert_eq!(decrypted_data.into_vec(), vec![1, 2, 3]);
 /// ```
-#[derive(PartialEq, Eq, Debug)]
-pub enum OurPeerStorage {
-	/// Contains ciphertext for transmission.
-	EncryptedPeerStorage {
-		/// Stores ciphertext.
-		cipher: Vec<u8>,
-	},
-	/// Contains serialised decrypted peer storage backup.
-	DecryptedPeerStorage {
-		/// Stores decrypted serialised data.
-		data: Vec<u8>,
-	},
+pub struct DecryptedOurPeerStorage {
+	data: Vec<u8>,
 }
 
-impl OurPeerStorage {
-	/// Return [`OurPeerStorage::EncryptedPeerStorage`] or [`OurPeerStorage::DecryptedPeerStorage`]
-	/// in [`Vec<u8>`] form.
+impl DecryptedOurPeerStorage {
+	/// Returns [`DecryptedOurPeerStorage`] with the given data.
+	pub fn new(data: Vec<u8>) -> Self {
+		Self { data }
+	}
+
+	/// Returns data stored in [`Vec<u8>`] format.
 	pub fn into_vec(self) -> Vec<u8> {
-		match self {
-			OurPeerStorage::DecryptedPeerStorage { data } => data,
-			OurPeerStorage::EncryptedPeerStorage { cipher } => cipher,
-		}
+		self.data
 	}
 
-	/// Returns [`OurPeerStorage::EncryptedPeerStorage`] with encrypted `cipher`.
-	///
-	/// This function takes a `key` (for encryption) and random_bytes (to derive nonce for encryption)
-	/// and returns a [`OurPeerStorage::EncryptedPeerStorage`] with encrypted data inside.
-	/// It should only be called on [`OurPeerStorage::DecryptedPeerStorage`] otherwise it would panic.
-	///
-	/// The resulting serialised data is intended to be directly used for transmission to the peers.
-	pub fn encrypt_peer_storage(
-		self, key: &PeerStorageKey, random_bytes: [u8; 32],
-	) -> OurPeerStorage {
-		match self {
-			OurPeerStorage::DecryptedPeerStorage { mut data } => {
-				let plaintext_len = data.len();
-				let nonce = derive_nonce(key, &random_bytes);
-
-				let mut chacha = ChaCha20Poly1305RFC::new(&key.inner, &nonce, b"");
-				let mut tag = [0; 16];
-				chacha.encrypt_full_message_in_place(&mut data[0..plaintext_len], &mut tag);
-
-				data.extend_from_slice(&tag);
-
-				// Prepend `random_bytes` in front of the encrypted_blob.
-				data.splice(0..0, random_bytes);
-				Self::EncryptedPeerStorage { cipher: data }
-			},
-			OurPeerStorage::EncryptedPeerStorage { cipher: _ } => {
-				panic!("Expected OurPeerStorage::DecryptedPeerStorage!");
-			},
-		}
+	/// Encrypts the data inside [`DecryptedOurPeerStorage`] using [`PeerStorageKey`] and `random_bytes`
+	/// and returns [`EncryptedOurPeerStorage`].
+	pub fn encrypt(self, key: &PeerStorageKey, random_bytes: [u8; 32]) -> EncryptedOurPeerStorage {
+		let mut data = self.data;
+		let plaintext_len = data.len();
+		let nonce = derive_nonce(key, &random_bytes);
+
+		let mut chacha = ChaCha20Poly1305RFC::new(&key.inner, &nonce, b"");
+		let mut tag = [0; 16];
+		chacha.encrypt_full_message_in_place(&mut data[0..plaintext_len], &mut tag);
+
+		data.extend_from_slice(&tag);
+
+		// Prepend `random_bytes` in front of the encrypted_blob.
+		data.splice(0..0, random_bytes);
+
+		EncryptedOurPeerStorage { cipher: data }
 	}
+}
+
+/// [`EncryptedOurPeerStorage`] represents encrypted state of the corresponding [`DecryptedOurPeerStorage`].
+///
+/// # Key Methods
+/// - [`EncryptedOurPeerStorage::new`]: Returns [`EncryptedOurPeerStorage`] with the given encrypted cipher.
+/// - [`EncryptedOurPeerStorage::decrypt`]: Returns [`DecryptedOurPeerStorage`] created from decrypting the cipher.
+/// - [`EncryptedOurPeerStorage::into_vec`]: Returns the cipher in [`Vec<u8>`] format.
+pub struct EncryptedOurPeerStorage {
+	cipher: Vec<u8>,
+}
 
-	/// This expects a [`OurPeerStorage::EncryptedPeerStorage`] which would be consumed
-	/// to return a [`OurPeerStorage::DecryptedPeerStorage`] in case of successful decryption.
-	///
-	/// It would return error if the ciphertext inside [`OurPeerStorage::EncryptedPeerStorage`] is
-	/// not correct and panic if it is called on [`OurPeerStorage::DecryptedPeerStorage`].
-	pub fn decrypt_peer_storage(self, key: &PeerStorageKey) -> Result<Self, ()> {
-		match self {
-			OurPeerStorage::EncryptedPeerStorage { mut cipher } => {
-				let cyphertext_len = cipher.len();
+impl EncryptedOurPeerStorage {
+	/// Returns [`EncryptedOurPeerStorage`] if cipher is of appropriate length, else returns error.
+	pub fn new(cipher: Vec<u8>) -> Result<Self, ()> {
+		// Ciphertext is of the form: random_bytes(32 bytes) + encrypted_data + tag(16 bytes).
+		const MIN_CIPHERTEXT_LEN: usize = 32 + 16;
+
+		if cipher.len() < MIN_CIPHERTEXT_LEN {
+			return Err(());
+		}
+		return Ok(Self { cipher });
+	}
 
-				// Ciphertext is of the form: random_bytes(32 bytes) + encrypted_data + tag(16 bytes).
-				let (data_mut, tag) = cipher.split_at_mut(cyphertext_len - 16);
-				let (random_bytes, encrypted_data) = data_mut.split_at_mut(32);
+	/// Returns cipher in the format [`Vec<u8>`].
+	pub fn into_vec(self) -> Vec<u8> {
+		self.cipher
+	}
 
-				let nonce = derive_nonce(key, random_bytes);
+	/// Returns [`DecryptedOurPeerStorage`] if it successfully decrypts the ciphertext with the `key`,
+	/// else returns error.
+	pub fn decrypt(self, key: &PeerStorageKey) -> Result<DecryptedOurPeerStorage, ()> {
+		let mut cipher = self.cipher;
+		let cyphertext_len = cipher.len();
 
-				let mut chacha = ChaCha20Poly1305RFC::new(&key.inner, &nonce, b"");
+		// Ciphertext is of the form: random_bytes(32 bytes) + encrypted_data + tag(16 bytes).
+		let (data_mut, tag) = cipher.split_at_mut(cyphertext_len - 16);
+		let (random_bytes, encrypted_data) = data_mut.split_at_mut(32);
 
-				if chacha.check_decrypt_in_place(encrypted_data, tag).is_err() {
-					return Err(());
-				}
+		let nonce = derive_nonce(key, random_bytes);
 
-				cipher.truncate(cyphertext_len - 16);
-				cipher.drain(0..32);
+		let mut chacha = ChaCha20Poly1305RFC::new(&key.inner, &nonce, b"");
 
-				Ok(Self::DecryptedPeerStorage { data: cipher })
-			},
-			OurPeerStorage::DecryptedPeerStorage { data: _ } => {
-				panic!("Expected OurPeerStorage::EncryptedPeerStorage!");
-			},
+		if chacha.check_decrypt_in_place(encrypted_data, tag).is_err() {
+			return Err(());
 		}
+
+		cipher.truncate(cyphertext_len - 16);
+		cipher.drain(0..32);
+
+		Ok(DecryptedOurPeerStorage { data: cipher })
 	}
 }
 
@@ -145,7 +142,7 @@ fn derive_nonce(key: &PeerStorageKey, random_bytes: &[u8]) -> [u8; 12] {
 
 #[cfg(test)]
 mod tests {
-	use crate::ln::our_peer_storage::{derive_nonce, OurPeerStorage};
+	use crate::ln::our_peer_storage::{derive_nonce, DecryptedOurPeerStorage};
 	use crate::sign::PeerStorageKey;
 
 	#[test]
@@ -156,18 +153,15 @@ mod tests {
 		let random_bytes2 = [201; 32];
 
 		// Happy Path
-		let decrypted_ops = OurPeerStorage::DecryptedPeerStorage { data: vec![42u8; 32] };
-		let decrypted_ops_res = decrypted_ops
-			.encrypt_peer_storage(&key1, random_bytes1)
-			.decrypt_peer_storage(&key1)
-			.unwrap();
+		let decrypted_ops = DecryptedOurPeerStorage::new(vec![42u8; 32]);
+		let decrypted_ops_res: DecryptedOurPeerStorage =
+			decrypted_ops.encrypt(&key1, random_bytes1).decrypt(&key1).unwrap();
 		assert_eq!(decrypted_ops_res.into_vec(), vec![42u8; 32]);
 
 		// Changing Key
-		let decrypted_ops_wrong_key = OurPeerStorage::DecryptedPeerStorage { data: vec![42u8; 32] };
-		let decrypted_ops_wrong_key_res = decrypted_ops_wrong_key
-			.encrypt_peer_storage(&key2, random_bytes2)
-			.decrypt_peer_storage(&key1);
+		let decrypted_ops_wrong_key = DecryptedOurPeerStorage::new(vec![42u8; 32]);
+		let decrypted_ops_wrong_key_res =
+			decrypted_ops_wrong_key.encrypt(&key2, random_bytes2).decrypt(&key1);
 		assert!(decrypted_ops_wrong_key_res.is_err());
 
 		// Nonce derivation happy path