proxy header-only responses properly

fixes redirects

proxy/proxy.go

@@ -79,7 +79,7 @@ func (p *Proxy) proxyResponse(w *ResponseWriter, r *ResponseReader) error {
 	w.takeHeaders(r)
 	transcoder, found := p.transcoders[r.ContentType()]
 	if !found {
-		return w.writeFrom(r)
+		return w.ReadFrom(r)
 	}
 	w.setChunked()
 	if err := transcoder.Transcode(w, r); err != nil {

proxy/response.go

@@ -55,28 +55,33 @@ func (w *ResponseWriter) takeHeaders(r *ResponseReader) {
 	w.WriteHeader(r.r.StatusCode)
 }
 
-func (w *ResponseWriter) WriteHeader(s int) {
-	w.statusCode = s
-}
-
 func (w *ResponseWriter) Header() http.Header {
 	return w.rw.Header()
 }
 
-func (w *ResponseWriter) writeFrom(r *ResponseReader) error {
-	w.rw.WriteHeader(r.r.StatusCode)
-	_, err := io.Copy(w.rw, r)
+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")
 }
-
-func (w *ResponseWriter) Write(b []byte) (int, error) {
-	if !w.headersDone {
-		w.rw.WriteHeader(w.statusCode)
-		w.headersDone = true
-	}
-	return w.Writer.Write(b)
-}

transcoder/identity.go

@@ -2,12 +2,10 @@ package transcoder
 
 import (
 	"github.com/barnacs/compy/proxy"
-	"io"
 )
 
 type Identity struct{}
 
 func (i *Identity) Transcode(w *proxy.ResponseWriter, r *proxy.ResponseReader) error {
-	_, err := io.Copy(w, r)
-	return err
+	return w.ReadFrom(r)
 }