// Standard net/http function. Shouldn't be used directly, http.Serve will use it. func (proxy *ProxyHttpServer) ServeHTTP(w http.ResponseWriter, r *http.Request) { //r.Header["X-Forwarded-For"] = w.RemoteAddr() 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 { removeProxyHeaders(ctx, r) resp, err = ctx.RoundTrip(r) if err != nil { ctx.Error = err resp = proxy.filterResponse(nil, ctx) if resp == nil { ctx.Logf("error read response %v %v:", r.URL.Host, err.Error()) http.Error(w, err.Error(), 500) return } } ctx.Logf("Received response %v", resp.Status) } origBody := resp.Body resp = proxy.filterResponse(resp, ctx) 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") } // Throttle response reader var reader io.Reader if proxy.TokenBucket != nil { reader = ratelimit.Reader(resp.Body, proxy.TokenBucket) } else { reader = resp.Body } copyHeaders(w.Header(), resp.Header) w.WriteHeader(resp.StatusCode) nr, err := io.Copy(w, reader) 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) } }
func NewRateLimitReader(rc io.ReadCloser, rate float64, capacity int64) io.ReadCloser { var rlr rateLimitReader rlr.rc = rc rlr.rlr = ratelimit.Reader(rc, ratelimit.NewBucketWithRate(rate, capacity)) return &rlr }