From bb498ce6993349b99395dc59e664d95a787d7360 Mon Sep 17 00:00:00 2001
From: rustysec <russ@infocyte.com>
Date: Fri, 13 Mar 2020 11:55:34 -0700
Subject: [PATCH] added basic authentication method

---
 src/error.rs |  5 +++++
 src/proxy.rs | 35 ++++++++++++++++++++++++++---------
 2 files changed, 31 insertions(+), 9 deletions(-)

diff --git a/src/error.rs b/src/error.rs
index d4d57c8..e7db332 100644
--- a/src/error.rs
+++ b/src/error.rs
@@ -29,6 +29,8 @@ pub enum Error {
     BadProxyCreds,
     /// Proxy could not connect
     ProxyConnect,
+    /// Incorrect credentials for proxy
+    InvalidProxyCreds,
 }
 
 impl Error {
@@ -56,6 +58,7 @@ impl Error {
             Error::BadProxy => 500,
             Error::BadProxyCreds => 500,
             Error::ProxyConnect => 500,
+            Error::InvalidProxyCreds => 500,
         }
     }
 
@@ -74,6 +77,7 @@ impl Error {
             Error::BadProxy => "Malformed proxy",
             Error::BadProxyCreds => "Failed to parse proxy credentials",
             Error::ProxyConnect => "Proxy failed to connect",
+            Error::InvalidProxyCreds => "Provided proxy credentials are incorrect",
         }
     }
 
@@ -92,6 +96,7 @@ impl Error {
             Error::BadProxy => "Malformed proxy".to_string(),
             Error::BadProxyCreds => "Failed to parse proxy credentials".to_string(),
             Error::ProxyConnect => "Proxy failed to connect".to_string(),
+            Error::InvalidProxyCreds => "Provided proxy credentials are incorrect".to_string(),
         }
     }
 }
diff --git a/src/proxy.rs b/src/proxy.rs
index 71c8fed..641f805 100644
--- a/src/proxy.rs
+++ b/src/proxy.rs
@@ -1,5 +1,4 @@
 use crate::error::Error;
-use std::fmt;
 
 /// Kind of proxy connection (Basic, Digest, etc)
 #[derive(Clone, Debug)]
@@ -86,29 +85,47 @@ impl Proxy {
         })
     }
 
-    pub(crate) fn kind(mut self, kind: ProxyKind) -> Self {
-        self.kind = kind;
-        self
-    }
-
     pub fn connect<S: AsRef<str>>(&self, host: S, port: u16) -> String {
+        let authorization = if self.use_authorization() {
+            match self.kind {
+                ProxyKind::Basic => {
+                    let creds = base64::encode(&format!(
+                        "{}:{}",
+                        self.user.clone().unwrap_or_default(),
+                        self.password.clone().unwrap_or_default()
+                    ));
+                    format!("Proxy-Authorization: basic {}\r\n", creds)
+                }
+            }
+        } else {
+            String::new()
+        };
+
         format!(
             "CONNECT {}:{} HTTP/1.1\r\n\
 Host: {}:{}\r\n\
 User-Agent: something/1.0.0\r\n\
-Proxy-Connection: Keep-Alive\r\n\r\n",
+Proxy-Connection: Keep-Alive\r\n\
+{}\
+\r\n",
             host.as_ref(),
             port,
             host.as_ref(),
-            port
+            port,
+            authorization
         )
     }
 
     pub(crate) fn verify_response(response: &[u8]) -> Result<(), Error> {
         let response_string = String::from_utf8_lossy(response);
         let top_line = response_string.lines().next().ok_or(Error::ProxyConnect)?;
+        let status_code = top_line.split_whitespace().nth(1).ok_or(Error::BadProxy)?;
 
-        Ok(())
+        match status_code {
+            "200" => Ok(()),
+            "401" | "407" => Err(Error::InvalidProxyCreds),
+            _ => Err(Error::BadProxy),
+        }
     }
 }