websocket support & request handling improvements

.gitignore

@@ -1,2 +1,3 @@
+/.idea/
 bin
 *.swp

ctx.go

@@ -9,18 +9,19 @@ import (
 // every user function. Also used as a logger.
 type ProxyCtx struct {
 	// Will contain the client request from the proxy
-	Req *http.Request
+	Req          *http.Request
 	// Will contain the remote server's response (if available. nil if the request wasn't send yet)
 	Resp         *http.Response
 	RoundTripper RoundTripper
 	// will contain the recent error that occured while trying to send receive or parse traffic
-	Error error
+	Error        error
 	// A handle for the user to keep data in the context, from the call of ReqHandler to the
 	// call of RespHandler
-	UserData interface{}
+	UserData     interface{}
 	// Will connect a request to a response
-	Session int64
-	proxy   *ProxyHttpServer
+	Session      int64
+	Websocket	 bool
+	proxy        *ProxyHttpServer
 }
 
 type RoundTripper interface {
@@ -41,7 +42,7 @@ func (ctx *ProxyCtx) RoundTrip(req *http.Request) (*http.Response, error) {
 }
 
 func (ctx *ProxyCtx) printf(msg string, argv ...interface{}) {
-	ctx.proxy.Logger.Printf("[%03d] "+msg+"\n", append([]interface{}{ctx.Session & 0xFF}, argv...)...)
+	ctx.proxy.Logger.Printf("[%03d] " + msg + "\n", append([]interface{}{ctx.Session & 0xFF}, argv...)...)
 }
 
 // Logf prints a message to the proxy's log. Should be used in a ProxyHttpServer's filter
@@ -54,7 +55,7 @@ func (ctx *ProxyCtx) printf(msg string, argv ...interface{}) {
 //	})
 func (ctx *ProxyCtx) Logf(msg string, argv ...interface{}) {
 	if ctx.proxy.Verbose {
-		ctx.printf("INFO: "+msg, argv...)
+		ctx.printf("INFO: " + msg, argv...)
 	}
 }
 
@@ -70,7 +71,7 @@ func (ctx *ProxyCtx) Logf(msg string, argv ...interface{}) {
 //		return r, nil
 //	})
 func (ctx *ProxyCtx) Warnf(msg string, argv ...interface{}) {
-	ctx.printf("WARN: "+msg, argv...)
+	ctx.printf("WARN: " + msg, argv...)
 }
 
 var charsetFinder = regexp.MustCompile("charset=([^ ;]*)")

examples/goproxy-websocket/main.go

@@ -0,0 +1,27 @@
+package main
+
+import (
+	"github.com/elazarl/goproxy"
+	"log"
+	"net/http"
+	"regexp"
+)
+
+func main() {
+	// Init
+	https := regexp.MustCompile("^.*:(443|8443)$")
+
+	proxy := goproxy.NewProxyHttpServer()
+	proxy.Verbose = true
+
+	// MitM
+	proxy.OnRequest().HandleConnectFunc(func(host string, ctx *goproxy.ProxyCtx) (*goproxy.ConnectAction, string) {
+		if https.MatchString(host) {
+			return goproxy.MitmConnect, host
+		} else {
+			return goproxy.HTTPMitmConnect, host
+		}
+	})
+
+	log.Fatal(http.ListenAndServe(":8888", proxy))
+}

https.go

@@ -11,7 +11,6 @@ import (
 	"net/url"
 	"os"
 	"regexp"
-	"strconv"
 	"strings"
 	"sync"
 	"sync/atomic"
@@ -64,7 +63,7 @@ func (proxy *ProxyHttpServer) connectDial(network, addr string) (c net.Conn, err
 	return proxy.ConnectDial(network, addr)
 }
 
-func (proxy *ProxyHttpServer) handleHttps(w http.ResponseWriter, r *http.Request) {
+func (proxy *ProxyHttpServer) handleConnect(w http.ResponseWriter, r *http.Request) {
 	ctx := &ProxyCtx{Req: r, Session: atomic.AddInt64(&proxy.sess, 1), proxy: proxy}
 
 	hij, ok := w.(http.Hijacker)
@@ -89,6 +88,7 @@ func (proxy *ProxyHttpServer) handleHttps(w http.ResponseWriter, r *http.Request
 			break
 		}
 	}
+
 	switch todo.Action {
 	case ConnectAccept:
 		if !hasPort.MatchString(host) {
@@ -116,51 +116,44 @@ func (proxy *ProxyHttpServer) handleHttps(w http.ResponseWriter, r *http.Request
 				wg.Wait()
 				proxyClient.Close()
 				targetSiteCon.Close()
-
 			}()
 		}
 
 	case ConnectHijack:
 		ctx.Logf("Hijacking CONNECT to %s", host)
 		proxyClient.Write([]byte("HTTP/1.0 200 OK\r\n\r\n"))
 		todo.Hijack(r, proxyClient, ctx)
+
 	case ConnectHTTPMitm:
-		proxyClient.Write([]byte("HTTP/1.0 200 OK\r\n\r\n"))
 		ctx.Logf("Assuming CONNECT is plain HTTP tunneling, mitm proxying it")
-		targetSiteCon, err := proxy.connectDial("tcp", host)
-		if err != nil {
-			ctx.Warnf("Error dialing to %s: %s", host, err.Error())
-			return
-		}
+		proxyClient.Write([]byte("HTTP/1.0 200 OK\r\n\r\n"))
+
+		client := bufio.NewReader(proxyClient)
+
 		for {
-			client := bufio.NewReader(proxyClient)
-			remote := bufio.NewReader(targetSiteCon)
 			req, err := http.ReadRequest(client)
-			if err != nil && err != io.EOF {
-				ctx.Warnf("cannot read request of MITM HTTP client: %+#v", err)
-			}
+
 			if err != nil {
-				return
-			}
-			req, resp := proxy.filterRequest(req, ctx)
-			if resp == nil {
-				if err := req.Write(targetSiteCon); err != nil {
-					httpError(proxyClient, ctx, err)
-					return
+				if err != io.EOF {
+					ctx.Warnf("cannot read request of MITM HTTP client: %+#v", err)
 				}
-				resp, err = http.ReadResponse(remote, req)
+
+				break
+			}
+
+			req.URL, err = url.Parse("http://" + req.Host + req.URL.String())
+
+			if end, err := proxy.handleRequest(NewConnResponseWriter(proxyClient), req); end {
 				if err != nil {
-					httpError(proxyClient, ctx, err)
-					return
+					ctx.Warnf("Error during serving MITM HTTP request: %+#v", err)
 				}
-				defer resp.Body.Close()
-			}
-			resp = proxy.filterResponse(resp, ctx)
-			if err := resp.Write(proxyClient); err != nil {
-				httpError(proxyClient, ctx, err)
-				return
+
+				break
 			}
 		}
+
+		proxyClient.Close()
+
 	case ConnectMitm:
 		proxyClient.Write([]byte("HTTP/1.0 200 OK\r\n\r\n"))
 		ctx.Logf("Assuming CONNECT is TLS, mitm proxying it")
@@ -177,93 +170,45 @@ func (proxy *ProxyHttpServer) handleHttps(w http.ResponseWriter, r *http.Request
 				return
 			}
 		}
-		go func() {
-			//TODO: cache connections to the remote website
-			rawClientTls := tls.Server(proxyClient, tlsConfig)
-			if err := rawClientTls.Handshake(); err != nil {
-				ctx.Warnf("Cannot handshake client %v %v", r.Host, err)
-				return
-			}
-			defer rawClientTls.Close()
-			clientTlsReader := bufio.NewReader(rawClientTls)
-			for !isEof(clientTlsReader) {
-				req, err := http.ReadRequest(clientTlsReader)
-				if err != nil && err != io.EOF {
-					return
-				}
-				if err != nil {
-					ctx.Warnf("Cannot read TLS request from mitm'd client %v %v", r.Host, err)
-					return
-				}
-				req.RemoteAddr = r.RemoteAddr // since we're converting the request, need to carry over the original connecting IP as well
-				ctx.Logf("req %v", r.Host)
 
-				if !httpsRegexp.MatchString(req.URL.String()) {
-					req.URL, err = url.Parse("https://" + r.Host + req.URL.String())
-				}
+		//TODO: cache connections to the remote website
+		rawClientTls := tls.Server(proxyClient, tlsConfig)
+		if err := rawClientTls.Handshake(); err != nil {
+			ctx.Warnf("Cannot handshake client %v %v", r.Host, err)
+			return
+		}
+		defer rawClientTls.Close()
+		clientTlsReader := bufio.NewReader(rawClientTls)
 
-				// Bug fix which goproxy fails to provide request
-				// information URL in the context when does HTTPS MITM
-				ctx.Req = req
-
-				req, resp := proxy.filterRequest(req, ctx)
-				if resp == nil {
-					if err != nil {
-						ctx.Warnf("Illegal URL %s", "https://"+r.Host+req.URL.Path)
-						return
-					}
-					removeProxyHeaders(ctx, req)
-					resp, err = ctx.RoundTrip(req)
-					if err != nil {
-						ctx.Warnf("Cannot read TLS response from mitm'd server %v", err)
-						return
-					}
-					ctx.Logf("resp %v", resp.Status)
-				}
-				resp = proxy.filterResponse(resp, ctx)
-				defer resp.Body.Close()
+		for {
+			req, err := http.ReadRequest(clientTlsReader)
 
-				text := resp.Status
-				statusCode := strconv.Itoa(resp.StatusCode) + " "
-				if strings.HasPrefix(text, statusCode) {
-					text = text[len(statusCode):]
-				}
-				// always use 1.1 to support chunked encoding
-				if _, err := io.WriteString(rawClientTls, "HTTP/1.1"+" "+statusCode+text+"\r\n"); err != nil {
-					ctx.Warnf("Cannot write TLS response HTTP status from mitm'd client: %v", err)
-					return
-				}
-				// Since we don't know the length of resp, return chunked encoded response
-				// TODO: use a more reasonable scheme
-				resp.Header.Del("Content-Length")
-				resp.Header.Set("Transfer-Encoding", "chunked")
-				if err := resp.Header.Write(rawClientTls); err != nil {
-					ctx.Warnf("Cannot write TLS response header from mitm'd client: %v", err)
-					return
-				}
-				if _, err = io.WriteString(rawClientTls, "\r\n"); err != nil {
-					ctx.Warnf("Cannot write TLS response header end from mitm'd client: %v", err)
-					return
-				}
-				chunked := newChunkedWriter(rawClientTls)
-				if _, err := io.Copy(chunked, resp.Body); err != nil {
-					ctx.Warnf("Cannot write TLS response body from mitm'd client: %v", err)
-					return
-				}
-				if err := chunked.Close(); err != nil {
-					ctx.Warnf("Cannot write TLS chunked EOF from mitm'd client: %v", err)
-					return
+			if err != nil {
+				if err != io.EOF {
+					ctx.Warnf("cannot read request of MITM HTTPS client: %+#v", err)
 				}
-				if _, err = io.WriteString(rawClientTls, "\r\n"); err != nil {
-					ctx.Warnf("Cannot write TLS response chunked trailer from mitm'd client: %v", err)
-					return
+
+				break
+			}
+
+			req.URL, err = url.Parse("https://" + req.Host + req.URL.String())
+
+			if end, err := proxy.handleRequest(NewConnResponseWriter(rawClientTls), req); end {
+				if err != nil {
+					ctx.Warnf("Error during serving MITM HTTPS request: %+#v", err)
 				}
+
+				break
 			}
-			ctx.Logf("Exiting on EOF")
-		}()
+		}
+
+		ctx.Logf("Exiting on EOF")
+		proxyClient.Close()
+
 	case ConnectProxyAuthHijack:
 		proxyClient.Write([]byte("HTTP/1.1 407 Proxy Authentication Required\r\n"))
 		todo.Hijack(r, proxyClient, ctx)
+
 	case ConnectReject:
 		if ctx.Resp != nil {
 			if err := ctx.Resp.Write(proxyClient); err != nil {