From af8c16dcc400e10cc87359f4432de300b6619c6e Mon Sep 17 00:00:00 2001
From: Eryn Wells <eryn@erynwells.me>
Date: Sun, 25 Mar 2018 18:50:45 -0400
Subject: [PATCH] Use HexBytes traits for base64

---
 src/b64.rs | 49 +++++++++++++++++--------------
 src/hex.rs | 86 +++++++++++++++++++++++++++++++++++++++++++++++++-----
 2 files changed, 106 insertions(+), 29 deletions(-)

diff --git a/src/b64.rs b/src/b64.rs
index 9c3985c..0d46290 100644
--- a/src/b64.rs
+++ b/src/b64.rs
@@ -1,32 +1,37 @@
+use hex::{AsHexBytes, HexResult};
+
 static B64: &'static str = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+\\";
 
 pub fn base64(hex: &str) -> Result<String, String> {
     let mut out = String::from("");
     let mut num_bits = 0;
     let mut acc: u32 = 0;
-    for (idx, c) in hex.char_indices() {
-        if let Some(c_digit) = c.to_digit(16) {
-            // Accumulate bytes until we have 6 chunks of 4.
-            acc = (acc << 4) + c_digit;
-            num_bits += 4;
-            if idx % 6 != 5 {
-                continue;
-            }
-        
-            // Read out 4 chunks of 6.
-            for i in (0..4).rev() {
-                let out_char_idx = ((acc >> (6 * i)) & 0x3F) as usize;
-                // TODO: I don't like this nth() call.
-                if let Some(out_char) = B64.chars().nth(out_char_idx) {
-                    out.push(out_char);
-                } else {
-                    return Err(format!("Couldn't make output char from {}", out_char_idx));
+    for (idx, c) in hex.hex_bytes().enumerate() {
+        match c {
+            HexResult::Byte(c) => {
+                // Accumulate bytes until we have 6 chunks of 4.
+                acc = (acc << 4) + (c as u32);
+                num_bits += 4;
+                if idx % 6 != 5 {
+                    continue;
                 }
-            }
-            acc = 0;
-            num_bits = 0;
-        } else {
-            return Err(format!("Invalid input: {}", c));
+        
+                // Read out 4 chunks of 6.
+                for i in (0..4).rev() {
+                    let out_char_idx = ((acc >> (6 * i)) & 0x3F) as usize;
+                    // TODO: I don't like this nth() call.
+                    if let Some(out_char) = B64.chars().nth(out_char_idx) {
+                        out.push(out_char);
+                    } else {
+                        return Err(format!("Couldn't make output char from {}", out_char_idx));
+                    }
+                }
+                acc = 0;
+                num_bits = 0;
+            },
+            HexResult::Invalid(c) => {
+                return Err(format!("Invalid input: {}", c));
+            },
         }
     }
 
diff --git a/src/hex.rs b/src/hex.rs
index badcdd1..0a0078c 100644
--- a/src/hex.rs
+++ b/src/hex.rs
@@ -1,3 +1,4 @@
+use std::char;
 use std::iter;
 use std::str::Chars;
 
@@ -6,17 +7,21 @@ pub enum HexResult {
     Invalid(char),
 }
 
-pub struct HexDecoder<'a> {
+pub struct HexBytes<'a> {
     input: Chars<'a>,
 }
 
-impl<'a> HexDecoder<'a> {
-    fn new(input: &'a str) -> HexDecoder {
-        HexDecoder { input: input.chars() }
+impl<'a> HexBytes<'a> {
+    fn new(input: &'a str) -> HexBytes {
+        HexBytes { input: input.chars() }
+    }
+
+    pub fn valid(self) -> ValidHexBytes<'a> {
+        ValidHexBytes::new(self)
     }
 }
 
-impl<'a> iter::Iterator for HexDecoder<'a> {
+impl<'a> iter::Iterator for HexBytes<'a> {
     type Item = HexResult;
 
     fn next(&mut self) -> Option<Self::Item> {
@@ -32,14 +37,81 @@ impl<'a> iter::Iterator for HexDecoder<'a> {
     }
 }
 
+/*
+ * ValidHexBytes --
+ */
+
+pub struct ValidHexBytes<'a> {
+    input: HexBytes<'a>,
+}
+
+impl<'a> ValidHexBytes<'a> {
+    fn new(input: HexBytes) -> ValidHexBytes {
+        ValidHexBytes { input: input }
+    }
+}
+
+impl<'a> iter::Iterator for ValidHexBytes<'a> {
+    type Item = u8;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        loop {
+            match self.input.next() {
+                Some(hex_c) => match hex_c {
+                    HexResult::Byte(c) => return Some(c),
+                    _ => continue,
+                },
+                None => return None,
+            }
+        }
+    }
+}
+
+/*
+ * HexDigest --
+ */
+
+pub trait HexDigest {
+    fn hex_digest(self) -> String;
+}
+
+impl<'a> HexDigest for HexBytes<'a> {
+    fn hex_digest(self) -> String {
+        self.valid().hex_digest()
+    }
+}
+
+impl<'a> HexDigest for ValidHexBytes<'a> {
+    fn hex_digest(self) -> String {
+        self.map(|x| char::from_digit(x as u32, 16).unwrap()).collect()
+    }
+}
+
+impl HexDigest for Vec<u8> {
+    fn hex_digest(self) -> String {
+        self.into_iter().map(|x| char::from_digit(x as u32, 16).unwrap()).collect()
+    }
+}
+
+/*
+ * AsHexbytes --
+ */
+
+pub trait AsHexBytes<'a> {
+    fn hex_bytes(&self) -> HexBytes<'a>;
+}
+
+impl<'a> AsHexBytes<'a> for &'a str {
+    fn hex_bytes(&self) -> HexBytes<'a> { HexBytes::new(self) }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
 
     #[test]
     fn simple() {
-        let decoder = HexDecoder::new("123");
-        let collected: Vec<u8> = decoder.map(|c| match c {
+        let collected: Vec<u8> = "123".hex_bytes().map(|c| match c {
             HexResult::Byte(c) => c,
             _ => panic!(),
         }).collect();