add the ability to set headers in the Pem file

jcreekmore · Feb 26, 2023 · 5d78060 · 5d78060
1 parent 62754ef
commit 5d78060
Show file tree

Hide file tree

Showing 2 changed files with 39 additions and 0 deletions.
diff --git a/src/errors.rs b/src/errors.rs
@@ -57,3 +57,13 @@ impl Error for PemError {
 
 /// The `pem` result type.
 pub type Result<T> = ::core::result::Result<T, PemError>;
+
+#[allow(missing_docs)]
+#[macro_export]
+macro_rules! ensure {
+    ($cond:expr, $err:expr) => {
+        if !$cond {
+            return Err($err);
+        }
+    };
+}
diff --git a/src/lib.rs b/src/lib.rs
@@ -218,6 +218,20 @@ impl HeaderMap {
     pub fn get(&self, key: &str) -> Option<&str> {
         self.iter().rev().find(|(k, _)| *k == key).map(|(_, v)| v)
     }
+
+    /// Get the last set value corresponding to the header key
+    pub fn add(&mut self, key: &str, value: &str) -> Result<()> {
+        ensure!(
+            !(key.contains(':') || key.contains('\n')),
+            PemError::InvalidHeader(key.to_string())
+        );
+        ensure!(
+            !(value.contains(':') || value.contains('\n')),
+            PemError::InvalidHeader(value.to_string())
+        );
+        self.0.push(format!("{}: {}", key.trim(), value.trim()));
+        Ok(())
+    }
 }
 
 impl Pem {
@@ -250,6 +264,11 @@ impl Pem {
         &self.headers
     }
 
+    /// Get the header map for modification
+    pub fn headers_mut(&mut self) -> &mut HeaderMap {
+        &mut self.headers
+    }
+
     fn new_from_captures(caps: Captures) -> Result<Pem> {
         fn as_utf8(bytes: &[u8]) -> Result<&str> {
             str::from_utf8(bytes).map_err(PemError::NotUtf8)
@@ -899,6 +918,16 @@ RzHX0lkJl9Stshd/7Gbt65/QYq+v+xvAeT0CoyIg",
             let pem = Pem::new(tag, contents);
             prop_assert_eq!(&pem, &pem.to_string().parse::<Pem>().unwrap());
         }
+
+        #[test]
+        fn test_str_parse_and_display_with_headers(tag in "[A-Z ]+",
+                                                   key in "[a-zA-Z]+",
+                                                   value in "[a-zA-A]+",
+                                                   contents in prop::collection::vec(0..255u8, 0..200)) {
+            let mut pem = Pem::new(tag, contents);
+            pem.headers_mut().add(&key, &value).unwrap();
+            prop_assert_eq!(&pem, &pem.to_string().parse::<Pem>().unwrap());
+        }
     }
 
     #[test]