[hook/proxy] add auth support

1 files changed, 33 insertions, 4 deletions

diff --git a/pkg/server/hook/proxy.go b/pkg/server/hook/proxy.go
index 8c3de7b..64db784 100644
--- a/pkg/server/hook/proxy.go
+++ b/pkg/server/hook/proxy.go
@@ -2,32 +2,46 @@ package hook
 
 import (
 	"bufio"
+	"bytes"
+	"encoding/base64"
 	"errors"
 	"fmt"
 	"io"
 	"regexp"
+	"strconv"
 	"tunnel/pkg/server/env"
 	"tunnel/pkg/server/opts"
 	"tunnel/pkg/server/queue"
 )
 
 var addrRe = regexp.MustCompile("^[0-9a-zA-Z-.]+:[0-9]+$")
+var respRe = regexp.MustCompile("^([^ ]+) +([0-9]+) +(.*)$")
 
 var errBadHttpResponse = errors.New("bad HTTP response")
 
 type proxyHook struct {
 	addr string
+	auth string
 }
 
 type proxy struct {
 	addr string
+	auth string
 	ok   chan struct{}
 	fail chan struct{}
 }
 
 func (p *proxy) Send(rq, wq queue.Q) error {
-	request := fmt.Sprintf("CONNECT %s HTTP/1.0\r\n\r\n", p.addr)
-	wq <- []byte(request)
+	var out bytes.Buffer
+
+	fmt.Fprintf(&out, "CONNECT %s HTTP/1.0\r\n", p.addr)
+	if p.auth != "" {
+		encoded := base64.StdEncoding.EncodeToString([]byte(p.auth))
+		fmt.Fprintf(&out, "Proxy-Authorization: Basic %s\r\n", encoded)
+	}
+	fmt.Fprintf(&out, "\r\n")
+
+	wq <- out.Bytes()
 
 	select {
 	case <-p.fail:
@@ -43,8 +57,16 @@ func parseProxyResponse(s string) error {
 	var code int
 	var desc string
 
-	if _, err := fmt.Sscanf(s, "%s %d %s", &version, &code, &desc); err != nil {
+	if m := respRe.FindStringSubmatch(s); m == nil {
 		return errBadHttpResponse
+	} else {
+		version = m[1]
+		if c, err := strconv.Atoi(m[2]); err != nil {
+			return errBadHttpResponse
+		} else {
+			code = c
+		}
+		desc = m[3]
 	}
 
 	if version != "HTTP/1.0" && version != "HTTP/1.1" {
@@ -96,13 +118,18 @@ func (p *proxy) Recv(rq, wq queue.Q) (err error) {
 	return queue.Copy(rq, wq)
 }
 
-func (h *proxyHook) Open(env.Env) (interface{}, error) {
+func (h *proxyHook) Open(env env.Env) (interface{}, error) {
 	p := &proxy{
 		addr: h.addr,
+		auth: h.auth,
 		ok:   make(chan struct{}),
 		fail: make(chan struct{}),
 	}
 
+	if p.auth == "" {
+		p.auth = getHookVar(env, "proxy.auth")
+	}
+
 	return p, nil
 }
 
@@ -117,6 +144,8 @@ func newProxyHook(opts opts.Opts, env env.Env) (hook, error) {
 		h.addr = addr
 	}
 
+	h.auth = opts["auth"]
+
 	return h, nil
 }