diff --git a/compy.go b/compy.go index aaf87ea..38943ee 100644 --- a/compy.go +++ b/compy.go @@ -6,6 +6,9 @@ import ( "github.com/barnacs/compy/proxy" tc "github.com/barnacs/compy/transcoder" "log" + "os" + "os/signal" + "sync/atomic" ) var ( @@ -55,6 +58,20 @@ func main() { p.AddTranscoder("application/javascript", ttc) p.AddTranscoder("application/x-javascript", ttc) + c := make(chan os.Signal, 2) + signal.Notify(c, os.Interrupt) + go func() { + for _ = range c { + read := atomic.LoadUint64(&p.ReadCount) + written := atomic.LoadUint64(&p.WriteCount) + log.Printf("compy exiting, total transcoded: %d -> %d (%3.1f%%)", + read, written, float64(written)/float64(read)*100) + os.Exit(0) + } + }() + + log.Printf("compy listening on %s", *host) + var err error if *cert != "" { err = p.StartTLS(*host, *cert, *key) diff --git a/proxy/proxy.go b/proxy/proxy.go index 7dfc3ff..bef25a1 100644 --- a/proxy/proxy.go +++ b/proxy/proxy.go @@ -5,11 +5,14 @@ import ( "log" "net" "net/http" + "sync/atomic" ) type Proxy struct { transcoders map[string]Transcoder ml *mitmListener + ReadCount uint64 + WriteCount uint64 } type Transcoder interface { @@ -47,6 +50,7 @@ func (p *Proxy) StartTLS(host, cert, key string) error { } func (p *Proxy) ServeHTTP(w http.ResponseWriter, r *http.Request) { + log.Printf("serving request: %s", r.URL) if err := p.handle(w, r); err != nil { log.Printf("%s while serving request: %s", err, r.URL) } @@ -62,7 +66,15 @@ func (p *Proxy) handle(w http.ResponseWriter, r *http.Request) error { return fmt.Errorf("error forwarding request: %s", err) } defer resp.Body.Close() - return p.proxyResponse(newResponseWriter(w), newResponseReader(resp)) + rw := newResponseWriter(w) + rr := newResponseReader(resp) + err = p.proxyResponse(rw, rr) + read := rr.counter.Count() + written := rw.rw.Count() + log.Printf("transcoded: %d -> %d (%3.1f%%)", read, written, float64(written)/float64(read)*100) + atomic.AddUint64(&p.ReadCount, read) + atomic.AddUint64(&p.WriteCount, written) + return err } func forward(r *http.Request) (*http.Response, error) { diff --git a/proxy/response.go b/proxy/response.go index ed251fe..6ab67dc 100644 --- a/proxy/response.go +++ b/proxy/response.go @@ -4,17 +4,22 @@ import ( "io" "mime" "net/http" + + "github.com/miolini/datacounter" ) type ResponseReader struct { io.Reader - r *http.Response + counter *datacounter.ReaderCounter + r *http.Response } func newResponseReader(r *http.Response) *ResponseReader { + counter := datacounter.NewReaderCounter(r.Body) return &ResponseReader{ - Reader: r.Body, - r: r, + Reader: counter, + counter: counter, + r: r, } } @@ -34,15 +39,16 @@ func (r *ResponseReader) Request() *http.Request { type ResponseWriter struct { io.Writer - rw http.ResponseWriter + rw *datacounter.ResponseWriterCounter statusCode int headersDone bool } func newResponseWriter(w http.ResponseWriter) *ResponseWriter { + rw := datacounter.NewResponseWriterCounter(w) return &ResponseWriter{ - Writer: w, - rw: w, + Writer: rw, + rw: rw, } }