Merge pull request #89 from SenseUnit/handler_refactoring

Handler refactoring

Author: Snawoot

README.md

@@ -219,8 +219,6 @@ Usage of /home/user/go/bin/dumbproxy:
     	bcrypt password cost (for -passwd mode) (default 4)
   -proxy value
     	upstream proxy URL. Can be repeated multiple times to chain proxies. Examples: socks5h://127.0.0.1:9050; https://user:password@example.com:443
-  -timeout duration
-    	timeout for network operations (default 10s)
   -user-ip-hints
     	allow IP hints to be specified by user in X-Src-IP-Hints header
   -verbosity int

handler/adapter.go

@@ -0,0 +1,136 @@
+package handler
+
+import (
+	"io"
+	"net/http"
+)
+
+type wrappedH2 struct {
+	r io.ReadCloser
+	w io.Writer
+}
+
+func wrapH2(r io.ReadCloser, w io.Writer) wrappedH2 {
+	return wrappedH2{
+		r: r,
+		w: w,
+	}
+}
+
+func (w wrappedH2) Read(p []byte) (n int, err error) {
+	return w.r.Read(p)
+}
+
+func (w wrappedH2) Write(p []byte) (n int, err error) {
+	n, err = w.w.Write(p)
+	if err != nil {
+		return
+	}
+	if f, ok := w.w.(http.Flusher); ok {
+		f.Flush()
+	}
+	return
+}
+
+func (w wrappedH2) Close() error {
+	// can't really close response writer, but at least we can disrupt copy
+	// closing Reader
+	return w.r.Close()
+}
+
+var _ io.ReadWriteCloser = wrappedH2{}
+
+type wrappedH1ReqBody struct {
+	r io.ReadCloser
+}
+
+func wrapH1ReqBody(r io.ReadCloser) wrappedH1ReqBody {
+	return wrappedH1ReqBody{
+		r: r,
+	}
+}
+
+func (w wrappedH1ReqBody) Read(p []byte) (n int, err error) {
+	return w.r.Read(p)
+}
+
+func (w wrappedH1ReqBody) Write(p []byte) (n int, err error) {
+	return len(p), nil
+}
+
+func (w wrappedH1ReqBody) Close() error {
+	return w.r.Close()
+}
+
+func (w wrappedH1ReqBody) CloseWrite() error {
+	return nil
+}
+
+var _ io.ReadWriteCloser = wrappedH1ReqBody{}
+var _ interface{ CloseWrite() error } = wrappedH1ReqBody{}
+
+type h1ReqBodyPipe struct {
+	r *io.PipeReader
+	w *io.PipeWriter
+}
+
+func newH1ReqBodyPipe() h1ReqBodyPipe {
+	r, w := io.Pipe()
+	return h1ReqBodyPipe{
+		r: r,
+		w: w,
+	}
+}
+
+func (w h1ReqBodyPipe) Read(p []byte) (n int, err error) {
+	return 0, io.EOF
+}
+
+func (w h1ReqBodyPipe) Write(p []byte) (n int, err error) {
+	return w.w.Write(p)
+}
+
+func (w h1ReqBodyPipe) Close() error {
+	return w.CloseWrite()
+}
+
+func (w h1ReqBodyPipe) CloseWrite() error {
+	return w.w.Close()
+}
+
+func (w h1ReqBodyPipe) Body() io.ReadCloser {
+	return w.r
+}
+
+var _ io.ReadWriteCloser = h1ReqBodyPipe{}
+var _ interface{ CloseWrite() error } = h1ReqBodyPipe{}
+
+type wrappedH1RespWriter struct {
+	w io.Writer
+}
+
+func wrapH1RespWriter(w io.Writer) wrappedH1RespWriter {
+	return wrappedH1RespWriter{
+		w: w,
+	}
+}
+
+func (w wrappedH1RespWriter) Read(p []byte) (n int, err error) {
+	return 0, io.EOF
+}
+
+func (w wrappedH1RespWriter) Write(p []byte) (n int, err error) {
+	n, err = w.w.Write(p)
+	if f, ok := w.w.(http.Flusher); ok {
+		f.Flush()
+	}
+	return
+}
+
+func (w wrappedH1RespWriter) Close() error {
+	// can't really close response writer, just make copier return
+	// and finish request
+	return nil
+}
+
+var _ io.ReadWriteCloser = wrappedH1RespWriter{}

handler/config.go

@@ -0,0 +1,25 @@
+package handler
+
+import (
+	"context"
+	"io"
+
+	"github.com/SenseUnit/dumbproxy/auth"
+	clog "github.com/SenseUnit/dumbproxy/log"
+)
+
+type Config struct {
+	// Dialer optionally specifies dialer to use for creating
+	// connections originating from proxy.
+	Dialer HandlerDialer
+	// Auth optionally specifies request validator used to verify users
+	// and return their username.
+	Auth auth.Auth
+	// Logger specifies optional custom logger.
+	Logger *clog.CondLogger
+	// Forward optionally specifies custom connection pairing function
+	// which does actual data forwarding.
+	Forward func(ctx context.Context, username string, incoming, outgoing io.ReadWriteCloser) error
+	// UserIPHints specifies whether allow IP hints set by user or not
+	UserIPHints bool
+}

handler/handler.go

@@ -3,11 +3,12 @@ package handler
 import (
 	"context"
 	"fmt"
+	"io"
+	"log"
 	"net"
 	"net/http"
 	"strings"
 	"sync"
-	"time"
 
 	"github.com/SenseUnit/dumbproxy/auth"
 	"github.com/SenseUnit/dumbproxy/dialer"
@@ -21,36 +22,50 @@ type HandlerDialer interface {
 }
 
 type ProxyHandler struct {
-	timeout       time.Duration
 	auth          auth.Auth
 	logger        *clog.CondLogger
 	dialer        HandlerDialer
+	forward       func(ctx context.Context, username string, incoming, outgoing io.ReadWriteCloser) error
 	httptransport http.RoundTripper
 	outbound      map[string]string
 	outboundMux   sync.RWMutex
 	userIPHints   bool
 }
 
-func NewProxyHandler(timeout time.Duration, auth auth.Auth, dialer HandlerDialer,
-	userIPHints bool, logger *clog.CondLogger) *ProxyHandler {
+func NewProxyHandler(config *Config) *ProxyHandler {
+	d := config.Dialer
+	if d == nil {
+		d = dialer.NewBoundDialer(nil, "")
+	}
 	httptransport := &http.Transport{
-		DialContext:       dialer.DialContext,
+		DialContext:       d.DialContext,
 		DisableKeepAlives: true,
 	}
+	a := config.Auth
+	if a == nil {
+		a = auth.NoAuth{}
+	}
+	l := config.Logger
+	if l == nil {
+		l = clog.NewCondLogger(log.New(io.Discard, "", 0), 0)
+	}
+	f := config.Forward
+	if f == nil {
+		f = PairConnections
+	}
 	return &ProxyHandler{
-		timeout:       timeout,
-		auth:          auth,
-		logger:        logger,
-		dialer:        dialer,
+		auth:          a,
+		logger:        l,
+		dialer:        d,
+		forward:       f,
 		httptransport: httptransport,
 		outbound:      make(map[string]string),
-		userIPHints:   userIPHints,
+		userIPHints:   config.UserIPHints,
 	}
 }
 
-func (s *ProxyHandler) HandleTunnel(wr http.ResponseWriter, req *http.Request) {
-	ctx, _ := context.WithTimeout(req.Context(), s.timeout)
-	conn, err := s.dialer.DialContext(ctx, "tcp", req.RequestURI)
+func (s *ProxyHandler) HandleTunnel(wr http.ResponseWriter, req *http.Request, username string) {
+	conn, err := s.dialer.DialContext(req.Context(), "tcp", req.RequestURI)
 	if err != nil {
 		s.logger.Error("Can't satisfy CONNECT request: %v", err)
 		http.Error(wr, "Can't satisfy CONNECT request", http.StatusBadGateway)
@@ -81,21 +96,27 @@ func (s *ProxyHandler) HandleTunnel(wr http.ResponseWriter, req *http.Request) {
 		// Inform client connection is built
 		fmt.Fprintf(localconn, "HTTP/%d.%d 200 OK\r\n\r\n", req.ProtoMajor, req.ProtoMinor)
 
-		proxy(req.Context(), localconn, conn)
+		s.forward(req.Context(), username, localconn, conn)
 	} else if req.ProtoMajor == 2 {
 		wr.Header()["Date"] = nil
 		wr.WriteHeader(http.StatusOK)
 		flush(wr)
-		proxyh2(req.Context(), req.Body, wr, conn)
+		s.forward(req.Context(), username, wrapH2(req.Body, wr), conn)
 	} else {
 		s.logger.Error("Unsupported protocol version: %s", req.Proto)
 		http.Error(wr, "Unsupported protocol version.", http.StatusBadRequest)
 		return
 	}
 }
 
-func (s *ProxyHandler) HandleRequest(wr http.ResponseWriter, req *http.Request) {
+func (s *ProxyHandler) HandleRequest(wr http.ResponseWriter, req *http.Request, username string) {
 	req.RequestURI = ""
+	forwardReqBody := newH1ReqBodyPipe()
+	origBody := req.Body
+	req.Body = forwardReqBody.Body()
+	go func() {
+		s.forward(req.Context(), username, wrapH1ReqBody(origBody), forwardReqBody)
+	}()
 	if req.ProtoMajor == 2 {
 		req.URL.Scheme = "http" // We can't access :scheme pseudo-header, so assume http
 		req.URL.Host = req.Host
@@ -112,7 +133,7 @@ func (s *ProxyHandler) HandleRequest(wr http.ResponseWriter, req *http.Request)
 	copyHeader(wr.Header(), resp.Header)
 	wr.WriteHeader(resp.StatusCode)
 	flush(wr)
-	copyBody(wr, resp.Body)
+	s.forward(req.Context(), username, wrapH1RespWriter(wr), wrapH1ReqBody(resp.Body))
 }
 
 func (s *ProxyHandler) isLoopback(req *http.Request) (string, bool) {
@@ -160,9 +181,9 @@ func (s *ProxyHandler) ServeHTTP(wr http.ResponseWriter, req *http.Request) {
 	req = req.WithContext(newCtx)
 	delHopHeaders(req.Header)
 	if isConnect {
-		s.HandleTunnel(wr, req)
+		s.HandleTunnel(wr, req, username)
 	} else {
-		s.HandleRequest(wr, req)
+		s.HandleRequest(wr, req, username)
 	}
 }