From a903b6a9c8b0ac1468b2d241857b519c5a9882ef Mon Sep 17 00:00:00 2001
From: Liron Levin <liron@twistlock.com>
Date: Thu, 4 Feb 2016 17:41:41 +0200
Subject: [PATCH] Fix 19575: Docker events doesn't work with authorization
 plugin

To support the requirement of blocking the request after the daemon
responded the authorization plugin use a `response recorder` that replay
the response after the flow ends.

This commit adds support for commands that hijack the connection and
flushes data via the http.Flusher interface. This resolves the error
with the event endpoint.

Signed-off-by: Liron Levin <liron@twistlock.com>
---
 authorization/authz.go      |  2 +-
 authorization/authz_test.go |  4 +-
 authorization/response.go   | 79 ++++++++++++++++++++++++++++++++++---
 3 files changed, 76 insertions(+), 9 deletions(-)

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
 }