Merge branch 'pull_129'

elazarl#129

Author: jdamick

ctx.go

@@ -19,8 +19,10 @@ type ProxyCtx struct {
 	// call of RespHandler
 	UserData interface{}
 	// Will connect a request to a response
-	Session int64
-	proxy   *ProxyHttpServer
+	Session         int64
+	proxy           *ProxyHttpServer
+	BytesUpstream   int64
+	BytesDownstream int64
 }
 
 type RoundTripper interface {

https.go

@@ -14,6 +14,7 @@ import (
 	"strconv"
 	"strings"
 	"sync/atomic"
+	"time"
 )
 
 type ConnectActionLiteral int
@@ -98,8 +99,9 @@ func (proxy *ProxyHttpServer) handleHttps(w http.ResponseWriter, r *http.Request
 		}
 		ctx.Logf("Accepting CONNECT to %s", host)
 		proxyClient.Write([]byte("HTTP/1.0 200 OK\r\n\r\n"))
-		go copyAndClose(ctx, targetSiteCon, proxyClient)
-		go copyAndClose(ctx, proxyClient, targetSiteCon)
+
+		pipeAndClose(ctx, targetSiteCon, proxyClient)
+
 	case ConnectHijack:
 		ctx.Logf("Hijacking CONNECT to %s", host)
 		proxyClient.Write([]byte("HTTP/1.0 200 OK\r\n\r\n"))
@@ -254,15 +256,46 @@ func httpError(w io.WriteCloser, ctx *ProxyCtx, err error) {
 	}
 }
 
-func copyAndClose(ctx *ProxyCtx, w, r net.Conn) {
+func pipeAndClose(ctx *ProxyCtx, targetSiteCon, proxyClient net.Conn) {
+	upCh := make(chan int64)
+	downCh := make(chan int64)
+	closedConns := 0
+
+	go func() {
+		upCh <- copyAndClose(ctx, targetSiteCon, proxyClient)
+	}()
+
+	go func() {
+		downCh <- copyAndClose(ctx, proxyClient, targetSiteCon)
+	}()
+
+	for closedConns < 2 {
+		select {
+		case bytes := <-upCh:
+			ctx.BytesUpstream += bytes
+			closedConns++
+			targetSiteCon.SetReadDeadline(time.Now())
+
+		case bytes := <-downCh:
+			ctx.BytesDownstream += bytes
+			closedConns++
+			proxyClient.SetReadDeadline(time.Now())
+		}
+	}
+}
+
+func copyAndClose(ctx *ProxyCtx, w, r net.Conn) (bytes int64) {
 	connOk := true
-	if _, err := io.Copy(w, r); err != nil {
+	var err error
+
+	if bytes, err = io.Copy(w, r); err != nil && bytes <= 0 {
 		connOk = false
 		ctx.Warnf("Error copying to client: %s", err)
 	}
 	if err := r.Close(); err != nil && connOk {
 		ctx.Warnf("Error closing: %s", err)
 	}
+	return bytes
 }
 
 func dialerFromEnv(proxy *ProxyHttpServer) func(network, addr string) (net.Conn, error) {

proxy.go

@@ -110,6 +110,7 @@ func (proxy *ProxyHttpServer) ServeHTTP(w http.ResponseWriter, r *http.Request)
 		if resp == nil {
 			removeProxyHeaders(ctx, r)
 			resp, err = ctx.RoundTrip(r)
+			ctx.BytesUpstream += r.ContentLength
 			if err != nil {
 				ctx.Error = err
 				resp = proxy.filterResponse(nil, ctx)
@@ -120,6 +121,7 @@ func (proxy *ProxyHttpServer) ServeHTTP(w http.ResponseWriter, r *http.Request)
 				}
 			}
 			ctx.Logf("Received response %v", resp.Status)
+
 		}
 		origBody := resp.Body
 		resp = proxy.filterResponse(resp, ctx)
@@ -140,6 +142,7 @@ func (proxy *ProxyHttpServer) ServeHTTP(w http.ResponseWriter, r *http.Request)
 		if err := resp.Body.Close(); err != nil {
 			ctx.Warnf("Can't close response body %v", err)
 		}
+		ctx.BytesDownstream += nr
 		ctx.Logf("Copied %v bytes to client error=%v", nr, err)
 	}
 }