From 44e1884eacee53bcd208317d45da1f92bd561b72 Mon Sep 17 00:00:00 2001
From: Jack Grigg <thestr4d@gmail.com>
Date: Sat, 16 May 2020 14:26:57 +1200
Subject: [PATCH] Read nonce for v1 format in Decryptor::new

A short input is now detected earlier, and no input reading is necessary
during the decryption process.
---
 age/src/protocol.rs           |  7 +++++--
 age/src/protocol/decryptor.rs | 35 ++++++++++++++++++++---------------
 2 files changed, 25 insertions(+), 17 deletions(-)

diff --git a/age/src/protocol.rs b/age/src/protocol.rs
index 5d36332..26db600 100644
--- a/age/src/protocol.rs
+++ b/age/src/protocol.rs
@@ -151,6 +151,9 @@ impl<R: Read> Decryptor<R> {
 
         match &header {
             Header::V1(v1_header) => {
+                let mut nonce = [0; 16];
+                input.read_exact(&mut nonce)?;
+
                 // Enforce structural requirements on the v1 header.
                 let any_scrypt = v1_header.recipients.iter().any(|r| {
                     if let RecipientStanza::Scrypt(_) = r {
@@ -161,9 +164,9 @@ impl<R: Read> Decryptor<R> {
                 });
 
                 if any_scrypt && v1_header.recipients.len() == 1 {
-                    Ok(decryptor::PassphraseDecryptor::new(input, header).into())
+                    Ok(decryptor::PassphraseDecryptor::new(input, header, nonce).into())
                 } else if !any_scrypt {
-                    Ok(decryptor::RecipientsDecryptor::new(input, header).into())
+                    Ok(decryptor::RecipientsDecryptor::new(input, header, nonce).into())
                 } else {
                     Err(Error::InvalidHeader)
                 }
diff --git a/age/src/protocol/decryptor.rs b/age/src/protocol/decryptor.rs
index 7f086c8..1733d3f 100644
--- a/age/src/protocol/decryptor.rs
+++ b/age/src/protocol/decryptor.rs
@@ -19,6 +19,8 @@ struct BaseDecryptor<R: Read> {
     input: ArmoredReader<R>,
     /// The age file's header.
     header: Header,
+    /// The age file's AEAD nonce
+    nonce: [u8; 16],
 }
 
 impl<R: Read> BaseDecryptor<R> {
@@ -27,17 +29,12 @@ impl<R: Read> BaseDecryptor<R> {
         F: FnMut(&RecipientStanza) -> Option<Result<FileKey, Error>>,
     {
         match &self.header {
-            Header::V1(header) => {
-                let mut nonce = [0; 16];
-                self.input.read_exact(&mut nonce)?;
-
-                header
-                    .recipients
-                    .iter()
-                    .find_map(filter)
-                    .unwrap_or(Err(Error::NoMatchingKeys))
-                    .and_then(|file_key| v1_payload_key(header, file_key, nonce))
-            }
+            Header::V1(header) => header
+                .recipients
+                .iter()
+                .find_map(filter)
+                .unwrap_or(Err(Error::NoMatchingKeys))
+                .and_then(|file_key| v1_payload_key(header, file_key, self.nonce)),
             Header::Unknown(_) => unreachable!(),
         }
     }
@@ -47,8 +44,12 @@ impl<R: Read> BaseDecryptor<R> {
 pub struct RecipientsDecryptor<R: Read>(BaseDecryptor<R>);
 
 impl<R: Read> RecipientsDecryptor<R> {
-    pub(super) fn new(input: ArmoredReader<R>, header: Header) -> Self {
-        RecipientsDecryptor(BaseDecryptor { input, header })
+    pub(super) fn new(input: ArmoredReader<R>, header: Header, nonce: [u8; 16]) -> Self {
+        RecipientsDecryptor(BaseDecryptor {
+            input,
+            header,
+            nonce,
+        })
     }
 
     /// Attempts to decrypt the age file.
@@ -83,8 +84,12 @@ impl<R: Read> RecipientsDecryptor<R> {
 pub struct PassphraseDecryptor<R: Read>(BaseDecryptor<R>);
 
 impl<R: Read> PassphraseDecryptor<R> {
-    pub(super) fn new(input: ArmoredReader<R>, header: Header) -> Self {
-        PassphraseDecryptor(BaseDecryptor { input, header })
+    pub(super) fn new(input: ArmoredReader<R>, header: Header, nonce: [u8; 16]) -> Self {
+        PassphraseDecryptor(BaseDecryptor {
+            input,
+            header,
+            nonce,
+        })
     }
 
     /// Attempts to decrypt the age file.