From 8d39bfd343685ca8a3e7a3c2dd26f89b4f8675eb Mon Sep 17 00:00:00 2001
From: Matt McCoy <mattnenterprise@yahoo.com>
Date: Tue, 28 Jun 2016 22:29:55 -0400
Subject: [PATCH] Hardening the oauth2 example for gmail

---
 examples/gmail_oauth2.rs | 11 ++++----
 src/authenticator.rs     |  1 +
 src/client.rs            | 59 ++++++++++++++++++++++++----------------
 3 files changed, 42 insertions(+), 29 deletions(-)

diff --git a/examples/gmail_oauth2.rs b/examples/gmail_oauth2.rs
index 381b4f0..5814e58 100644
--- a/examples/gmail_oauth2.rs
+++ b/examples/gmail_oauth2.rs
@@ -3,7 +3,7 @@ extern crate openssl;
 extern crate base64;
 
 use openssl::ssl::{SslContext, SslMethod};
-use base64::{encode, decode};
+use base64::{encode};
 use imap::client::Client;
 use imap::authenticator::Authenticator;
 
@@ -13,15 +13,16 @@ struct GmailOAuth2 {
 }
 
 impl Authenticator for GmailOAuth2 {
+    #[allow(unused_variables)]
     fn process(&self, data: String) -> String {
-        String::from("dXNlcj1tYXR0bWNjb3kxMTBAZ21haWwuY29tAWF1dGg9QmVhcmVyIHlhMjkuQ2k4UUEzQ1Y5SW1OQ0Z1NDNpbkZRcngtSUR0cjVFSkZHNXdEM1IySzBXdTdiM1dzVG1Md")
+        encode(format!("user={}\x01auth=Bearer {}\x01\x01", self.user, self.access_token).as_bytes())
     }
 }
 
 fn main() {
-    let mut gmail_auth = GmailOAuth2{
-        user: String::from("email@gmail.com"),
-        access_token: String::from("")
+    let gmail_auth = GmailOAuth2{
+        user: String::from("sombody@gmail.com"),
+        access_token: String::from("<access_token>")
     };
     let mut imap_socket = Client::secure_connect(("imap.gmail.com", 993), SslContext::new(SslMethod::Sslv23).unwrap()).unwrap();
 
diff --git a/src/authenticator.rs b/src/authenticator.rs
index 16004a1..825ec9a 100644
--- a/src/authenticator.rs
+++ b/src/authenticator.rs
@@ -1,3 +1,4 @@
+/// This will allow plugable authentication mechanisms.
 pub trait Authenticator {
     fn process(&self, String) -> String;
 }
diff --git a/src/client.rs b/src/client.rs
index 4954172..9a8f8e0 100644
--- a/src/client.rs
+++ b/src/client.rs
@@ -63,29 +63,41 @@ impl<T: Read+Write> Client<T> {
 	pub fn authenticate<A: Authenticator>(&mut self, auth_type: &str, authenticator: A) -> Result<()> {
 		match self.run_command(&format!("AUTHENTICATE {}", auth_type).to_string()) {
 			Ok(_) => {
-				let line = match self.readline() {
-					Ok(l) => l,
-					Err(e) => return Err(e)
-				};
-				// TODO test this for the many authentication use cases
-				let data = match parse_authenticate_response(String::from_utf8(line).unwrap()) {
-					Ok(d) => d,
-					Err(e) => return Err(e)
-				};
-				println!("Done parsing authenticating response");
-				let auth_response = authenticator.process(data);
-				println!("Writing: {}", auth_response.clone());
-				match self.stream.write_all(auth_response.into_bytes().as_slice()) {
-					Err(e) => return Err(Error::Io(e)),
-					_ => {}
-				};
-				match self.stream.write(vec![0x0d, 0x0a].as_slice()) {
-					Err(e) => return Err(Error::Io(e)),
-					_ => {}
-				};
-				match self.read_response() {
-					Ok(_) => Ok(()),
-					Err(e) => Err(e)
+				loop {
+					let line = match self.readline() {
+						Ok(l) => l,
+						Err(e) => return Err(e)
+					};
+					if line.starts_with(b"+") {
+						let data = match parse_authenticate_response(String::from_utf8(line).unwrap()) {
+							Ok(d) => d,
+							Err(e) => return Err(e)
+						};
+						let auth_response = authenticator.process(data);
+						match self.stream.write_all(auth_response.into_bytes().as_slice()) {
+							Err(e) => return Err(Error::Io(e)),
+							_ => {}
+						};
+						match self.stream.write(vec![0x0d, 0x0a].as_slice()) {
+							Err(e) => return Err(Error::Io(e)),
+							_ => {}
+						};
+					} else if line.starts_with(format!("{}{} ", TAG_PREFIX, self.tag).as_bytes()) {
+						match parse_response(vec![String::from_utf8(line).unwrap()]) {
+							Ok(_) => return Ok(()),
+							Err(e) => return Err(e)
+						};
+					} else {
+						let mut lines = match self.read_response() {
+							Ok(l) => l,
+							Err(e) => return Err(e)
+						};
+						lines.insert(0, String::from_utf8(line).unwrap());
+						match parse_response(lines.clone()) {
+							Ok(_) => return Ok(()),
+							Err(e) => return Err(e)
+						};
+					}
 				}
 			},
 			Err(e) => Err(e)
@@ -279,7 +291,6 @@ impl<T: Read+Write> Client<T> {
 				print!("{}", String::from_utf8_lossy(byte_buffer));
 				line_buffer.push(byte_buffer[0]);
 		}
-		println!("{}", String::from_utf8(line_buffer.clone()).unwrap());
 		Ok(line_buffer)
 	}