diff --git a/authorization/authz.go b/authorization/authz.go index c0ef387..8a15b2b 100644 --- a/authorization/authz.go +++ b/authorization/authz.go @@ -116,7 +116,7 @@ func (ctx *Ctx) AuthZResponse(rm ResponseModifier, r *http.Request) error { } } - rm.Flush() + rm.FlushAll() return nil } diff --git a/authorization/authz_test.go b/authorization/authz_test.go index 75c549b..3a6a991 100644 --- a/authorization/authz_test.go +++ b/authorization/authz_test.go @@ -118,7 +118,7 @@ func TestResponseModifier(t *testing.T) { m.Write([]byte("body")) m.WriteHeader(500) - m.Flush() + m.FlushAll() if r.Header().Get("h1") != "v1" { t.Fatalf("Header value must exists %s", r.Header().Get("h1")) } @@ -147,7 +147,7 @@ func TestResponseModifierOverride(t *testing.T) { m.OverrideHeader(overrideHeaderBytes) m.OverrideBody([]byte("override body")) m.OverrideStatusCode(404) - m.Flush() + m.FlushAll() if r.Header().Get("h1") != "v2" { t.Fatalf("Header value must exists %s", r.Header().Get("h1")) } diff --git a/authorization/response.go b/authorization/response.go index 8acbcd1..abe3c9f 100644 --- a/authorization/response.go +++ b/authorization/response.go @@ -5,6 +5,7 @@ import ( "bytes" "encoding/json" "fmt" + "github.com/Sirupsen/logrus" "net" "net/http" ) @@ -12,6 +13,8 @@ import ( // ResponseModifier allows authorization plugins to read and modify the content of the http.response type ResponseModifier interface { http.ResponseWriter + http.Flusher + http.CloseNotifier // RawBody returns the current http content RawBody() []byte @@ -32,7 +35,10 @@ type ResponseModifier interface { OverrideStatusCode(statusCode int) // Flush flushes all data to the HTTP response - Flush() error + FlushAll() error + + // Hijacked indicates the response has been hijacked by the Docker daemon + Hijacked() bool } // NewResponseModifier creates a wrapper to an http.ResponseWriter to allow inspecting and modifying the content @@ -44,7 +50,10 @@ func NewResponseModifier(rw http.ResponseWriter) ResponseModifier { // the http request/response from docker daemon type responseModifier struct { // The original response writer - rw http.ResponseWriter + rw http.ResponseWriter + + r *http.Request + status int // body holds the response body body []byte @@ -52,15 +61,34 @@ type responseModifier struct { header http.Header // statusCode holds the response status code statusCode int + // hijacked indicates the request has been hijacked + hijacked bool +} + +func (rm *responseModifier) Hijacked() bool { + return rm.hijacked } // WriterHeader stores the http status code func (rm *responseModifier) WriteHeader(s int) { + + // Use original request if hijacked + if rm.hijacked { + rm.rw.WriteHeader(s) + return + } + rm.statusCode = s } // Header returns the internal http header func (rm *responseModifier) Header() http.Header { + + // Use original header if hijacked + if rm.hijacked { + return rm.rw.Header() + } + return rm.header } @@ -90,6 +118,11 @@ func (rm *responseModifier) OverrideHeader(b []byte) error { // Write stores the byte array inside content func (rm *responseModifier) Write(b []byte) (int, error) { + + if rm.hijacked { + return rm.rw.Write(b) + } + rm.body = append(rm.body, b...) return len(b), nil } @@ -109,6 +142,10 @@ func (rm *responseModifier) RawHeaders() ([]byte, error) { // Hijack returns the internal connection of the wrapped http.ResponseWriter func (rm *responseModifier) Hijack() (net.Conn, *bufio.ReadWriter, error) { + + rm.hijacked = true + rm.FlushAll() + hijacker, ok := rm.rw.(http.Hijacker) if !ok { return nil, nil, fmt.Errorf("Internal reponse writer doesn't support the Hijacker interface") @@ -116,8 +153,30 @@ func (rm *responseModifier) Hijack() (net.Conn, *bufio.ReadWriter, error) { return hijacker.Hijack() } -// Flush flushes all data to the HTTP response -func (rm *responseModifier) Flush() error { +// CloseNotify uses the internal close notify API of the wrapped http.ResponseWriter +func (rm *responseModifier) CloseNotify() <-chan bool { + closeNotifier, ok := rm.rw.(http.CloseNotifier) + if !ok { + logrus.Errorf("Internal reponse writer doesn't support the CloseNotifier interface") + return nil + } + return closeNotifier.CloseNotify() +} + +// Flush uses the internal flush API of the wrapped http.ResponseWriter +func (rm *responseModifier) Flush() { + flusher, ok := rm.rw.(http.Flusher) + if !ok { + logrus.Errorf("Internal reponse writer doesn't support the Flusher interface") + return + } + + rm.FlushAll() + flusher.Flush() +} + +// FlushAll flushes all data to the HTTP response +func (rm *responseModifier) FlushAll() error { // Copy the status code if rm.statusCode > 0 { rm.rw.WriteHeader(rm.statusCode) @@ -130,7 +189,15 @@ func (rm *responseModifier) Flush() error { } } - // Write body - _, err := rm.rw.Write(rm.body) + var err error + if len(rm.body) > 0 { + // Write body + _, err = rm.rw.Write(rm.body) + } + + // Clean previous data + rm.body = nil + rm.statusCode = 0 + rm.header = http.Header{} return err }