HTTP error handling customization

http.go

@@ -0,0 +1,86 @@
+package goproxy
+
+import (
+	"io"
+	"net/http"
+	"sync/atomic"
+)
+
+func (proxy *ProxyHttpServer) handleHttp(w http.ResponseWriter, r *http.Request) {
+	ctx := &ProxyCtx{Req: r, Session: atomic.AddInt64(&proxy.sess, 1), Proxy: proxy}
+
+	var err error
+	ctx.Logf("Got request %v %v %v %v", r.URL.Path, r.Host, r.Method, r.URL.String())
+	if !r.URL.IsAbs() {
+		proxy.NonproxyHandler.ServeHTTP(w, r)
+		return
+	}
+	r, resp := proxy.filterRequest(r, ctx)
+
+	if resp == nil {
+		if isWebSocketRequest(r) {
+			ctx.Logf("Request looks like websocket upgrade.")
+			proxy.serveWebsocket(ctx, w, r)
+		}
+
+		if !proxy.KeepHeader {
+			removeProxyHeaders(ctx, r)
+		}
+		resp, err = ctx.RoundTrip(r)
+		if err != nil {
+			ctx.Error = err
+			resp = proxy.filterResponse(nil, ctx)
+
+		}
+		if resp != nil {
+			ctx.Logf("Received response %v", resp.Status)
+		}
+	}
+
+	var origBody io.ReadCloser
+
+	if resp != nil {
+		origBody = resp.Body
+		defer origBody.Close()
+	}
+
+	resp = proxy.filterResponse(resp, ctx)
+
+	if resp == nil {
+		hij, ok := w.(http.Hijacker)
+		if !ok {
+			panic("httpserver does not support hijacking")
+		}
+
+		proxyClient, _, e := hij.Hijack()
+		if e != nil {
+			panic("Cannot hijack connection " + e.Error())
+		}
+		httpError(proxyClient, ctx, ctx.Error)
+		return
+	}
+
+	ctx.Logf("Copying response to client %v [%d]", resp.Status, resp.StatusCode)
+	// http.ResponseWriter will take care of filling the correct response length
+	// Setting it now, might impose wrong value, contradicting the actual new
+	// body the user returned.
+	// We keep the original body to remove the header only if things changed.
+	// This will prevent problems with HEAD requests where there's no body, yet,
+	// the Content-Length header should be set.
+	if origBody != resp.Body {
+		resp.Header.Del("Content-Length")
+	}
+	copyHeaders(w.Header(), resp.Header, proxy.KeepDestinationHeaders)
+	w.WriteHeader(resp.StatusCode)
+	var copyWriter io.Writer = w
+	if w.Header().Get("content-type") == "text/event-stream" {
+		// server-side events, flush the buffered data to the client.
+		copyWriter = &flushWriter{w: w}
+	}
+
+	nr, err := io.Copy(copyWriter, resp.Body)
+	if err := resp.Body.Close(); err != nil {
+		ctx.Warnf("Can't close response body %v", err)
+	}
+	ctx.Logf("Copied %v bytes to client error=%v", nr, err)
+}

https.go

@@ -324,6 +324,11 @@ func (proxy *ProxyHttpServer) handleHttps(w http.ResponseWriter, r *http.Request
 }
 
 func httpError(w io.WriteCloser, ctx *ProxyCtx, err error) {
+	if ctx.Proxy.HTTPErrorHandler != nil {
+		ctx.Proxy.HTTPErrorHandler(w, ctx, err)
+		return
+	}
+
 	if _, err := io.WriteString(w, "HTTP/1.1 502 Bad Gateway\r\n\r\n"); err != nil {
 		ctx.Warnf("Error responding to client: %s", err)
 	}

proxy.go

@@ -8,7 +8,6 @@ import (
 	"net/http"
 	"os"
 	"regexp"
-	"sync/atomic"
 )
 
 // The basic proxy type. Implements http.Handler.
@@ -26,6 +25,8 @@ type ProxyHttpServer struct {
 	respHandlers    []RespHandler
 	httpsHandlers   []HttpsHandler
 	Tr              *http.Transport
+	// Will be invoked to return a custom response to clients when goproxy fails to connect to a proxy target
+	HTTPErrorHandler func(io.WriteCloser, *ProxyCtx, error)
 	// ConnectDial will be used to create TCP connections for CONNECT requests
 	// if nil Tr.Dial will be used
 	ConnectDial        func(network string, addr string) (net.Conn, error)
@@ -128,81 +129,7 @@ func (proxy *ProxyHttpServer) ServeHTTP(w http.ResponseWriter, r *http.Request)
 	if r.Method == "CONNECT" {
 		proxy.handleHttps(w, r)
 	} else {
-		ctx := &ProxyCtx{Req: r, Session: atomic.AddInt64(&proxy.sess, 1), Proxy: proxy}
-
-		var err error
-		ctx.Logf("Got request %v %v %v %v", r.URL.Path, r.Host, r.Method, r.URL.String())
-		if !r.URL.IsAbs() {
-			proxy.NonproxyHandler.ServeHTTP(w, r)
-			return
-		}
-		r, resp := proxy.filterRequest(r, ctx)
-
-		if resp == nil {
-			if isWebSocketRequest(r) {
-				ctx.Logf("Request looks like websocket upgrade.")
-				proxy.serveWebsocket(ctx, w, r)
-			}
-
-			if !proxy.KeepHeader {
-				removeProxyHeaders(ctx, r)
-			}
-			resp, err = ctx.RoundTrip(r)
-			if err != nil {
-				ctx.Error = err
-				resp = proxy.filterResponse(nil, ctx)
-
-			}
-			if resp != nil {
-				ctx.Logf("Received response %v", resp.Status)
-			}
-		}
-
-		var origBody io.ReadCloser
-
-		if resp != nil {
-			origBody = resp.Body
-			defer origBody.Close()
-		}
-
-		resp = proxy.filterResponse(resp, ctx)
-
-		if resp == nil {
-			var errorString string
-			if ctx.Error != nil {
-				errorString = "error read response " + r.URL.Host + " : " + ctx.Error.Error()
-				ctx.Logf(errorString)
-				http.Error(w, ctx.Error.Error(), 500)
-			} else {
-				errorString = "error read response " + r.URL.Host
-				ctx.Logf(errorString)
-				http.Error(w, errorString, 500)
-			}
-			return
-		}
-		ctx.Logf("Copying response to client %v [%d]", resp.Status, resp.StatusCode)
-		// http.ResponseWriter will take care of filling the correct response length
-		// Setting it now, might impose wrong value, contradicting the actual new
-		// body the user returned.
-		// We keep the original body to remove the header only if things changed.
-		// This will prevent problems with HEAD requests where there's no body, yet,
-		// the Content-Length header should be set.
-		if origBody != resp.Body {
-			resp.Header.Del("Content-Length")
-		}
-		copyHeaders(w.Header(), resp.Header, proxy.KeepDestinationHeaders)
-		w.WriteHeader(resp.StatusCode)
-		var copyWriter io.Writer = w
-		if w.Header().Get("content-type") == "text/event-stream" {
-			// server-side events, flush the buffered data to the client.
-			copyWriter = &flushWriter{w: w}
-		}
-
-		nr, err := io.Copy(copyWriter, resp.Body)
-		if err := resp.Body.Close(); err != nil {
-			ctx.Warnf("Can't close response body %v", err)
-		}
-		ctx.Logf("Copied %v bytes to client error=%v", nr, err)
+		proxy.handleHttp(w, r)
 	}
 }