From 67118d884b66f7d783fbefab58b20829e2662ea4 Mon Sep 17 00:00:00 2001
From: Michal Budzyn <michalbudzyn@gmail.com>
Date: Fri, 4 May 2018 00:05:51 +0200
Subject: [PATCH] HTTP Proxy dialer

---
 proxy/client.go | 11 ++++++--
 proxy/dial.go   | 74 ++++++++++++++++++++++++++++++++++++++++++++++++-
 2 files changed, 82 insertions(+), 3 deletions(-)

diff --git a/proxy/client.go b/proxy/client.go
index d2764d6b..4116374b 100644
--- a/proxy/client.go
+++ b/proxy/client.go
@@ -198,6 +198,10 @@ func (c *Client) DialAndAuth(brokerAddress string) (net.Conn, error) {
 	if err != nil {
 		return nil, err
 	}
+	if err := conn.SetDeadline(time.Time{}); err != nil {
+		conn.Close()
+		return nil, err
+	}
 	err = c.auth(conn)
 	if err != nil {
 		return nil, err
@@ -212,6 +216,7 @@ func (c *Client) auth(conn net.Conn) error {
 			return err
 		}
 		if err := conn.SetDeadline(time.Time{}); err != nil {
+			conn.Close()
 			return err
 		}
 	}
@@ -221,8 +226,10 @@ func (c *Client) auth(conn net.Conn) error {
 			conn.Close()
 			return err
 		}
-		// reset deadlines
-		return conn.SetDeadline(time.Time{})
+		if err := conn.SetDeadline(time.Time{}); err != nil {
+			conn.Close()
+			return err
+		}
 	}
 	return nil
 }
diff --git a/proxy/dial.go b/proxy/dial.go
index 8de4b7f7..f8d48eda 100644
--- a/proxy/dial.go
+++ b/proxy/dial.go
@@ -1,10 +1,15 @@
 package proxy
 
 import (
+	"bufio"
 	"crypto/tls"
+	"encoding/base64"
+	"fmt"
 	"github.com/pkg/errors"
 	"golang.org/x/net/proxy"
 	"net"
+	"net/http"
+	"net/url"
 	"strings"
 	"time"
 )
@@ -23,7 +28,16 @@ func (d directDialer) Dial(network, addr string) (net.Conn, error) {
 		Timeout:   d.dialTimeout,
 		KeepAlive: d.keepAlive,
 	}
-	return dialer.Dial(network, addr)
+	conn, err := dialer.Dial(network, addr)
+	if err != nil {
+		return nil, err
+	}
+	err = conn.SetDeadline(time.Now().Add(d.dialTimeout))
+	if err != nil {
+		conn.Close()
+		return nil, err
+	}
+	return conn, err
 }
 
 type socks5Dialer struct {
@@ -122,3 +136,61 @@ func (d tlsDialer) Dial(network, addr string) (net.Conn, error) {
 
 	return conn, nil
 }
+
+type httpProxy struct {
+	forwardDialer      Dialer
+	hostPort           string
+	username, password string
+}
+
+func (s *httpProxy) Dial(network, addr string) (net.Conn, error) {
+	reqURL, err := url.Parse("http://" + addr)
+	if err != nil {
+		return nil, err
+	}
+	req, err := http.NewRequest("CONNECT", reqURL.String(), nil)
+	if err != nil {
+		return nil, err
+	}
+	req.Close = false
+	if s.username != "" && s.password != "" {
+		req.Header.Set("Proxy-Authorization", base64.StdEncoding.EncodeToString([]byte(s.username+":"+s.password)))
+	}
+
+	c, err := s.forwardDialer.Dial("tcp", s.hostPort)
+	if err != nil {
+		return nil, err
+	}
+	err = req.Write(c)
+	if err != nil {
+		c.Close()
+		return nil, err
+	}
+
+	resp, err := http.ReadResponse(bufio.NewReader(c), req)
+	if err != nil {
+		c.Close()
+		return nil, err
+	}
+	resp.Body.Close()
+	if resp.StatusCode != 200 {
+		c.Close()
+		return nil, fmt.Errorf("connect server using proxy error, statuscode [%d]", resp.StatusCode)
+	}
+
+	return c, nil
+}
+
+func newHTTPProxy(uri *url.URL, forward Dialer) (Dialer, error) {
+	s := new(httpProxy)
+	s.hostPort = uri.Host
+	if uri.Port() == "" {
+		return nil, fmt.Errorf("http proxy url doesn't contain a port [%v]", uri)
+	}
+	s.forwardDialer = forward
+	if uri.User != nil {
+		s.username = uri.User.Username()
+		s.password, _ = uri.User.Password()
+	}
+	return s, nil
+}