From cc90defb730a7f1632155509437b42d59bc8b552 Mon Sep 17 00:00:00 2001
From: Eryn Wells <eryn@erynwells.me>
Date: Sun, 25 Mar 2018 13:15:09 -0400
Subject: [PATCH] 1.1: Handle padding

---
 src/lib.rs | 40 +++++++++++++++++++++++++++++++++++-----
 1 file changed, 35 insertions(+), 5 deletions(-)

diff --git a/src/lib.rs b/src/lib.rs
index 8b74d77..b882f02 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -2,17 +2,19 @@ static B64: &'static str = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz
 
 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 24 of them. (We can do this with any number of bytes % 6.
+            // Accumulate bytes until we have 6 chunks of 4.
             acc = (acc << 4) + c_digit;
-            println!("idx={}, c='{}'({:2}), acc={:024b}", idx, c, c_digit, acc);
+            num_bits += 4;
+            println!("idx={}, c='{}'({:2}), acc={:024b}, n={}", idx, c, c_digit, acc, num_bits);
             if idx % 6 != 5 {
                 continue;
             }
         
-            // Read out bytes in chunks of 6.
+            // Read out 4 chunks of 6.
             for i in (0..4).rev() {
                 let out_char_idx = ((acc >> (6 * i)) & 0x3F) as usize;
                 println!("  i:{}, shift:{}, out_char_idx:{:08b}", i, 6*i, out_char_idx); 
@@ -24,11 +26,30 @@ pub fn base64(hex: &str) -> Result<String, String> {
                 }
             }
             acc = 0;
+            num_bits = 0;
         } else {
             return Err(format!("Invalid input: {}", c));
         }
     }
 
+    if acc != 0 {
+        acc = acc << (24 - num_bits);
+        let padding = (24 - num_bits) / 6;
+        println!("Making up for missing bits: n={}, acc={:024b}", num_bits, acc);
+        for i in (0..4).rev() {
+            let out_char_idx = ((acc >> (6 * i)) & 0x3F) as usize;
+            println!("  i:{}, shift:{}, out_char_idx:{:08b}", i, 6*i, out_char_idx); 
+            // TODO: I don't like this nth() call.
+            if i < padding {
+                out.push('=');
+            } else 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 {} (shifted: {})", out_char_idx, 6 * i));
+            }
+        }
+    }
+
     Ok(out)
 }
 
@@ -37,7 +58,7 @@ mod tests {
     use super::base64;
 
     #[test]
-    fn simple_single() {
+    fn small_wikipedia_example() {
         let input = "4d616e";
         let ex_output = "TWFu";
         println!("");
@@ -46,7 +67,7 @@ mod tests {
     }
 
     #[test]
-    fn padding() {
+    fn one_byte_padding() {
         let input = "6f6d";
         let ex_output = "b20=";
         println!("");
@@ -54,6 +75,15 @@ mod tests {
         assert_eq!(output, ex_output);
     }
 
+    #[test]
+    fn two_byte_padding() {
+        let input = "6f";
+        let ex_output = "bw==";
+        println!("");
+        let output = base64(input).expect("Error converting to base64");
+        assert_eq!(output, ex_output);
+    }
+
     #[test]
     fn cryptopals_input() {
         let input = "49276d206b696c6c696e6720796f757220627261696e206c696b65206120706f69736f6e6f7573206d757368726f6f6d";