add proxy (http connect) hook

1 files changed, 125 insertions, 0 deletions

diff --git a/pkg/server/hook/proxy.go b/pkg/server/hook/proxy.go
new file mode 100644
index 0000000..8c3de7b
--- /dev/null
+++ b/pkg/server/hook/proxy.go
@@ -0,0 +1,125 @@
+package hook
+
+import (
+	"bufio"
+	"errors"
+	"fmt"
+	"io"
+	"regexp"
+	"tunnel/pkg/server/env"
+	"tunnel/pkg/server/opts"
+	"tunnel/pkg/server/queue"
+)
+
+var addrRe = regexp.MustCompile("^[0-9a-zA-Z-.]+:[0-9]+$")
+
+var errBadHttpResponse = errors.New("bad HTTP response")
+
+type proxyHook struct {
+	addr string
+}
+
+type proxy struct {
+	addr 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)
+
+	select {
+	case <-p.fail:
+		return nil
+	case <-p.ok:
+	}
+
+	return queue.Copy(rq, wq)
+}
+
+func parseProxyResponse(s string) error {
+	var version string
+	var code int
+	var desc string
+
+	if _, err := fmt.Sscanf(s, "%s %d %s", &version, &code, &desc); err != nil {
+		return errBadHttpResponse
+	}
+
+	if version != "HTTP/1.0" && version != "HTTP/1.1" {
+		return errBadHttpResponse
+	}
+
+	if code != 200 {
+		return fmt.Errorf("connect failed: %d %s", code, desc)
+	}
+
+	return nil
+}
+
+func (p *proxy) Recv(rq, wq queue.Q) (err error) {
+	defer func() {
+		if err != nil {
+			close(p.fail)
+		}
+	}()
+
+	s := bufio.NewScanner(rq.Reader())
+
+	var resp bool
+
+	for s.Scan() {
+		line := s.Text()
+
+		if !resp {
+			if err := parseProxyResponse(line); err != nil {
+				return err
+			}
+			resp = true
+			continue
+		}
+
+		if line == "" {
+			break
+		}
+	}
+
+	if err := s.Err(); err != nil {
+		return err
+	} else if !resp {
+		return io.ErrUnexpectedEOF
+	}
+
+	close(p.ok)
+
+	return queue.Copy(rq, wq)
+}
+
+func (h *proxyHook) Open(env.Env) (interface{}, error) {
+	p := &proxy{
+		addr: h.addr,
+		ok:   make(chan struct{}),
+		fail: make(chan struct{}),
+	}
+
+	return p, nil
+}
+
+func newProxyHook(opts opts.Opts, env env.Env) (hook, error) {
+	h := &proxyHook{}
+
+	if addr, ok := opts["addr"]; !ok {
+		return nil, errors.New("proxy: missing addr")
+	} else if !addrRe.MatchString(addr) {
+		return nil, errors.New("proxy: invalid addr")
+	} else {
+		h.addr = addr
+	}
+
+	return h, nil
+}
+
+func init() {
+	register("proxy", newProxyHook)
+}