f de-DRY ChaChaPoly::encrypt_in_place and remove tag param

Author: valentinewallace

lightning/src/util/chacha20poly1305rfc.rs

@@ -61,31 +61,29 @@ mod real_chachapoly {
 		}
 
 		pub fn encrypt(&mut self, input: &[u8], output: &mut [u8], out_tag: &mut [u8]) {
+			assert!(input.len() == output.len());
+			assert!(self.finished == false);
 			self.cipher.process(input, output);
-			self.encrypt_inner(input, Some(output), Some(out_tag));
-		}
-
-		pub fn encrypt_in_place(&mut self, input_output: &mut [u8], out_tag: Option<&mut [u8]>) {
-			self.cipher.process_in_place(input_output);
-			self.encrypt_inner(input_output, None, out_tag);
+			self.data_len += input.len();
+			self.mac.input(output);
+			ChaCha20Poly1305RFC::pad_mac_16(&mut self.mac, self.data_len);
+			self.finished = true;
+			self.mac.input(&self.aad_len.to_le_bytes());
+			self.mac.input(&(self.data_len as u64).to_le_bytes());
+			self.mac.raw_result(out_tag);
 		}
 
-		// Encrypt in place, and fill in the tag if it's provided. If the tag is not provided, then
-		// `finish_and_get_tag` may be called to check it later.
-		fn encrypt_inner(&mut self, input: &[u8], output: Option<&mut [u8]>, out_tag: Option<&mut [u8]>) {
+		// ChaCha20Poly1305-encrypt `input_output` in-place. To finish and calculate the tag, use
+		// `finish_and_get_tag` below.
+		pub fn encrypt_in_place(&mut self, input_output: &mut [u8]) {
 			assert!(self.finished == false);
-			if let Some(output) = output {
-				assert!(input.len() == output.len());
-				self.mac.input(output);
-			} else {
-				self.mac.input(input);
-			}
-			self.data_len += input.len();
-			if let Some(tag) = out_tag {
-				self.finish_and_get_tag(tag);
-			}
+			self.data_len += input_output.len();
+			self.cipher.process_in_place(input_output);
+			self.mac.input(input_output);
 		}
 
+		// If we were previously encrypting with `encrypt_in_place`, this method can be used to finish
+		// encrypting and calculate the tag.
 		pub fn finish_and_get_tag(&mut self, out_tag: &mut [u8]) {
 			ChaCha20Poly1305RFC::pad_mac_16(&mut self.mac, self.data_len);
 			self.finished = true;
@@ -136,7 +134,7 @@ impl<'a, W: Writer> Writer for ChaChaPolyWriter<'a, W> {
 		for i in 0..num_writes {
 			let mut write_buffer = [0; 8192];
 			let bytes_written = (&mut write_buffer[..]).write(&src[i * 8192..])?;
-			self.chacha.encrypt_in_place(&mut write_buffer[..bytes_written], None);
+			self.chacha.encrypt_in_place(&mut write_buffer[..bytes_written]);
 			self.write.write_all(&write_buffer[..bytes_written])?;
 		}
 		Ok(())
@@ -204,7 +202,7 @@ mod fuzzy_chachapoly {
 			self.finished = true;
 		}
 
-		pub fn encrypt_in_place(&mut self, _input_output: &mut [u8], out_tag: Option<&mut [u8]>) {
+		pub fn encrypt_in_place(&mut self, _input_output: &mut [u8]) {
 			assert!(self.finished == false);
 			if let Some(tag) = out_tag {
 				tag.copy_from_slice(&self.tag);