From c06c6ba3bf2117c5756fb1336e17c1306027ce11 Mon Sep 17 00:00:00 2001
From: Stephen J Day <stephen.day@docker.com>
Date: Mon, 20 Apr 2015 11:15:01 -0700
Subject: [PATCH] Return instrumented response writer from context

This is ensures that users of the ResponseWriter from the context correctly
track usage. Otherwise, context reporting is incorrect.

Signed-off-by: Stephen J Day <stephen.day@docker.com>
---
 context/http.go      |  2 +-
 context/http_test.go | 13 +++++++++++--
 2 files changed, 12 insertions(+), 3 deletions(-)

diff --git a/context/http.go b/context/http.go
index 98ab436d..2b17074b 100644
--- a/context/http.go
+++ b/context/http.go
@@ -302,7 +302,7 @@ func (irw *instrumentedResponseWriter) Flush() {
 func (irw *instrumentedResponseWriter) Value(key interface{}) interface{} {
 	if keyStr, ok := key.(string); ok {
 		if keyStr == "http.response" {
-			return irw.ResponseWriter
+			return irw
 		}
 
 		if !strings.HasPrefix(keyStr, "http.response.") {
diff --git a/context/http_test.go b/context/http_test.go
index 42c78b75..f133574f 100644
--- a/context/http_test.go
+++ b/context/http_test.go
@@ -132,8 +132,17 @@ func TestWithResponseWriter(t *testing.T) {
 	trw := testResponseWriter{}
 	ctx, rw := WithResponseWriter(Background(), &trw)
 
-	if ctx.Value("http.response") != &trw {
-		t.Fatalf("response not available in context: %v != %v", ctx.Value("http.response"), &trw)
+	if ctx.Value("http.response") != rw {
+		t.Fatalf("response not available in context: %v != %v", ctx.Value("http.response"), rw)
+	}
+
+	grw, err := GetResponseWriter(ctx)
+	if err != nil {
+		t.Fatalf("error getting response writer: %v", err)
+	}
+
+	if grw != rw {
+		t.Fatalf("unexpected response writer returned: %#v != %#v", grw, rw)
 	}
 
 	if n, err := rw.Write(make([]byte, 1024)); err != nil {