From c133604f9d6ffc802428f4b85ee66b035b603be8 Mon Sep 17 00:00:00 2001
From: Brian Dwyer <Brian.Dwyer@broadridge.com>
Date: Sat, 18 Apr 2020 21:13:50 -0400
Subject: [PATCH] Clean up NTLM process

Signed-off-by: Brian Dwyer <Brian.Dwyer@broadridge.com>
---
 ntlm_windows.go | 23 ++++++++++++-----------
 1 file changed, 12 insertions(+), 11 deletions(-)

diff --git a/ntlm_windows.go b/ntlm_windows.go
index 7d1436b..7ce9ba4 100644
--- a/ntlm_windows.go
+++ b/ntlm_windows.go
@@ -55,7 +55,7 @@ func dialNTLM(p Proxy, addr string, baseDial func() (net.Conn, error)) (net.Conn
 		Host:   addr,
 		Header: h,
 	}
-	if err := connect.Write(conn); err != nil {
+	if err := connect.WriteProxy(conn); err != nil {
 		debugf("ntlm> Could not write negotiate message to proxy: %s", err)
 		return conn, err
 	}
@@ -65,6 +65,9 @@ func dialNTLM(p Proxy, addr string, baseDial func() (net.Conn, error)) (net.Conn
 		debugf("ntlm> Could not read negotiate response from proxy: %s", err)
 		return conn, err
 	}
+	if err := resp.Body.Close(); err != nil {
+		return conn, err
+	}
 
 	if resp.StatusCode != http.StatusProxyAuthRequired {
 		debugf("ntlm> Expected %d as return status, got: %d", http.StatusProxyAuthRequired, resp.StatusCode)
@@ -94,17 +97,15 @@ func dialNTLM(p Proxy, addr string, baseDial func() (net.Conn, error)) (net.Conn
 		return conn, err
 	}
 
-	resp.Body.Close()
-	h = p.Headers.Clone()
-	h.Set("Proxy-Authorization", fmt.Sprintf("NTLM %s", base64.StdEncoding.EncodeToString(authenticate)))
-	h.Set("Proxy-Connection", "Keep-Alive")
-	connect = &http.Request{
-		Method: "CONNECT",
-		URL:    &url.URL{Opaque: addr},
-		Host:   addr,
-		Header: h,
+	// Rewind the request body, the handshake needs it
+	if connect.GetBody != nil {
+		if connect.Body, err = connect.GetBody(); err != nil {
+			return conn, err
+		}
 	}
-	if err := connect.Write(conn); err != nil {
+
+	connect.Header.Set("Proxy-Authorization", fmt.Sprintf("NTLM %s", base64.StdEncoding.EncodeToString(authenticate)))
+	if err := connect.WriteProxy(conn); err != nil {
 		debugf("ntlm> Could not write authenticate message to proxy: %s", err)
 		return conn, err
 	}