compy/proxy/response.go
2017-01-14 03:44:07 +01:00

94 lines
1.7 KiB
Go

package proxy
import (
"io"
"mime"
"net/http"
"github.com/miolini/datacounter"
)
type ResponseReader struct {
io.Reader
counter *datacounter.ReaderCounter
r *http.Response
}
func newResponseReader(r *http.Response) *ResponseReader {
counter := datacounter.NewReaderCounter(r.Body)
return &ResponseReader{
Reader: counter,
counter: counter,
r: r,
}
}
func (r *ResponseReader) ContentType() string {
cth := r.Header().Get("Content-Type")
ct, _, _ := mime.ParseMediaType(cth)
return ct
}
func (r *ResponseReader) Header() http.Header {
return r.r.Header
}
func (r *ResponseReader) Request() *http.Request {
return r.r.Request
}
type ResponseWriter struct {
io.Writer
rw *datacounter.ResponseWriterCounter
statusCode int
headersDone bool
}
func newResponseWriter(w http.ResponseWriter) *ResponseWriter {
rw := datacounter.NewResponseWriterCounter(w)
return &ResponseWriter{
Writer: rw,
rw: rw,
}
}
func (w *ResponseWriter) takeHeaders(r *ResponseReader) {
for k, v := range r.Header() {
for _, v := range v {
w.Header().Add(k, v)
}
}
w.WriteHeader(r.r.StatusCode)
}
func (w *ResponseWriter) Header() http.Header {
return w.rw.Header()
}
func (w *ResponseWriter) WriteHeader(s int) {
w.statusCode = s
}
func (w *ResponseWriter) ReadFrom(r io.Reader) error {
w.flushHeaders()
_, err := io.Copy(w.Writer, r)
return err
}
func (w *ResponseWriter) Write(b []byte) (int, error) {
w.flushHeaders()
return w.Writer.Write(b)
}
func (w *ResponseWriter) flushHeaders() {
if w.headersDone {
return
}
w.rw.WriteHeader(w.statusCode)
w.headersDone = true
}
func (w *ResponseWriter) setChunked() {
w.Header().Del("Content-Length")
}