Add attached functions for v5. (#11)

* Add attached fns for v5.

* Expose magic for edoc for checking

* Code review

src/v3/mod.rs

@@ -14,7 +14,10 @@ use protobuf::Message;
 const IV_LEN: usize = 12;
 const GCM_TAG_LEN: usize = 16;
 
-pub const V3: u8 = 3u8;
+const V3: u8 = 3u8;
+
+/// For external users to check the first bytes of an edoc.
+pub const VERSION_AND_MAGIC: [u8; 5] = [V3, MAGIC[0], MAGIC[1], MAGIC[2], MAGIC[3]];
 
 // [3, b"IRON]
 const MAGIC_HEADER_LEN: usize = 5;

src/v5/attached.rs

@@ -0,0 +1,160 @@
+use bytes::{Buf, Bytes};
+use protobuf::Message;
+
+use crate::{aes::IvAndCiphertext, icl_header_v4::V4DocumentHeader, Error};
+
+use super::key_id_header::{self, KeyIdHeader};
+
+type Result<A> = std::result::Result<A, Error>;
+
+#[derive(Debug, PartialEq)]
+pub struct AttachedDocument {
+    pub key_id_header: KeyIdHeader,
+    pub edek: V4DocumentHeader,
+    pub edoc: IvAndCiphertext,
+}
+
+impl TryFrom<Vec<u8>> for AttachedDocument {
+    type Error = Error;
+
+    /// Breaks apart an attached edoc into its parts.
+    fn try_from(value: Vec<u8>) -> Result<Self> {
+        Bytes::from(value).try_into()
+    }
+}
+
+impl TryFrom<Bytes> for AttachedDocument {
+    type Error = Error;
+
+    /// Breaks apart an attached edoc into its parts.
+    fn try_from(value: Bytes) -> Result<Self> {
+        let (key_id_header, mut attached_document_with_header) =
+            key_id_header::decode_version_prefixed_value(value)?;
+        if attached_document_with_header.len() > 2 {
+            let edek_len = attached_document_with_header.get_u16();
+            if attached_document_with_header.len() > edek_len as usize {
+                let header_bytes = attached_document_with_header.split_to(edek_len as usize);
+                let edek = protobuf::Message::parse_from_bytes(&header_bytes[..])
+                    .map_err(|e| Error::HeaderParseErr(e.to_string()))?;
+                Ok(AttachedDocument {
+                    key_id_header,
+                    edek,
+                    edoc: IvAndCiphertext(attached_document_with_header),
+                })
+            } else {
+                Err(Error::HeaderParseErr(
+                    "The EDEK you passed in was too short based on the length bytes.".to_string(),
+                ))
+            }
+        } else {
+            Err(Error::HeaderParseErr("Header is too short.".to_string()))
+        }
+    }
+}
+
+impl AttachedDocument {
+    /// Write out the entire v5 attached documents to bytes.
+    pub fn write_to_bytes(&self) -> Result<Bytes> {
+        let AttachedDocument {
+            key_id_header,
+            edek,
+            edoc,
+        } = self;
+        let key_id_header_bytes = key_id_header.write_to_bytes();
+        let encoded_edek = edek.write_to_bytes().expect("Writing to bytes is safe");
+        if encoded_edek.len() > u16::MAX as usize {
+            Err(Error::HeaderLengthOverflow(encoded_edek.len() as u64))
+        } else {
+            let len = encoded_edek.len() as u16;
+
+            let result = [
+                key_id_header_bytes.as_ref(),
+                &len.to_be_bytes(),
+                &encoded_edek,
+                &edoc.0, // Note that the edoc is written without the leading OIRON since it's an attached document.
+            ]
+            .concat();
+            Ok(result.into())
+        }
+    }
+}
+
+#[cfg(test)]
+mod test {
+    use crate::{
+        aes::IvAndCiphertext,
+        icl_header_v4::{
+            v4document_header::{
+                edek_wrapper::{Aes256GcmEncryptedDek, Edek},
+                EdekWrapper, SignedPayload,
+            },
+            V4DocumentHeader,
+        },
+        v5::key_id_header::{EdekType, KeyId, KeyIdHeader, PayloadType},
+        Error,
+    };
+
+    use super::*;
+
+    #[test]
+    fn test_roundtrip() {
+        let key_id_header = KeyIdHeader::new(
+            EdekType::SaasShield,
+            PayloadType::StandardEdek,
+            KeyId(u32::MAX),
+        );
+
+        let edek_wrapper = EdekWrapper {
+            edek: Some(Edek::Aes256GcmEdek(Aes256GcmEncryptedDek {
+                ciphertext: [42u8; 255].as_ref().into(),
+                ..Default::default()
+            })),
+            ..Default::default()
+        };
+
+        let edek = V4DocumentHeader {
+            signed_payload: Some(SignedPayload {
+                edeks: vec![edek_wrapper],
+                ..Default::default()
+            })
+            .into(),
+            ..Default::default()
+        };
+
+        let edoc = IvAndCiphertext(vec![100, 200, 0].into());
+        let expected_result = AttachedDocument {
+            key_id_header,
+            edek,
+            edoc,
+        };
+        let encoded = expected_result.write_to_bytes().unwrap();
+        let result: AttachedDocument = encoded.try_into().unwrap();
+        assert_eq!(result, expected_result);
+    }
+
+    #[test]
+    fn test_len_too_long() {
+        // The `0,255` after the `0` is the u16 representing length. It is too large.
+        let encoded = vec![
+            255u8, 255, 255, 255, 2, 0, 0, 255, 18, 7, 18, 5, 26, 3, 18, 1, 42, 100, 200,
+        ];
+        let result = AttachedDocument::try_from(encoded).unwrap_err();
+        assert_eq!(
+            result,
+            Error::HeaderParseErr(
+                "The EDEK you passed in was too short based on the length bytes.".to_string()
+            )
+        );
+    }
+
+    #[test]
+    fn test_header_too_short() {
+        // This value is just a key_id header with 1 0 after.
+        let encoded = vec![255u8, 255, 255, 255, 2, 0, 0];
+        let result = AttachedDocument::try_from(encoded).unwrap_err();
+        assert_eq!(
+            result,
+            Error::HeaderParseErr("Header is too short.".to_string())
+        );
+    }
+}

src/v5/key_id_header.rs

@@ -109,6 +109,7 @@ impl EdekType {
 }
 
 /// The key id header parsed into its pieces.
+#[derive(Debug, PartialEq)]
 pub struct KeyIdHeader {
     pub key_id: KeyId,
     pub edek_type: EdekType,

src/v5/mod.rs

@@ -1,5 +1,6 @@
 // Reexport the v4 aes, because we also use it for v5.
 pub use crate::v4::aes;
+pub mod attached;
 pub mod key_id_header;
 use crate::{
     aes::{aes_encrypt, EncryptionKey, IvAndCiphertext, PlaintextDocument},
@@ -17,6 +18,8 @@ use rand::{CryptoRng, RngCore};
 type Result<T> = core::result::Result<T, Error>;
 const MAGIC: &[u8; 4] = crate::v4::MAGIC;
 pub(crate) const V0: u8 = 0u8;
+/// For external users to check the first bytes of an edoc.
+pub const VERSION_AND_MAGIC: [u8; 5] = [V0, MAGIC[0], MAGIC[1], MAGIC[2], MAGIC[3]];
 /// This is 0 + IRON
 pub(crate) const DETACHED_HEADER_LEN: usize = 5;