From 10f602b158119eac73535ba2aa6510c1887857b4 Mon Sep 17 00:00:00 2001
From: Aaron Lehmann <aaron.lehmann@docker.com>
Date: Thu, 6 Aug 2015 15:14:44 -0700
Subject: [PATCH] Don't panic when a http.ResponseWriter does not implement
 CloseNotifier

Instead, provide a variant of instrumentedResponseWriter that does not
implement CloseNotifier, and use that when necessary. In
copyFullPayload, log instead of panicing when we encounter something
that doesn't implement CloseNotifier.

This is more complicated than I'd like, but it's necessary because
instrumentedResponseWriter must not embed CloseNotifier unless there's
really a CloseNotifier to embed.

Signed-off-by: Aaron Lehmann <aaron.lehmann@docker.com>
---
 context/http.go              | 30 +++++++++++++++++++++---------
 context/http_test.go         |  7 -------
 registry/handlers/helpers.go |  2 +-
 3 files changed, 22 insertions(+), 17 deletions(-)

diff --git a/context/http.go b/context/http.go
index f61e3bc24..97b1067fd 100644
--- a/context/http.go
+++ b/context/http.go
@@ -103,17 +103,21 @@ func GetRequestID(ctx Context) string {
 // WithResponseWriter returns a new context and response writer that makes
 // interesting response statistics available within the context.
 func WithResponseWriter(ctx Context, w http.ResponseWriter) (Context, http.ResponseWriter) {
-	closeNotifier, ok := w.(http.CloseNotifier)
-	if !ok {
-		panic("the ResponseWriter does not implement CloseNotifier")
-	}
-	irw := &instrumentedResponseWriter{
+	irw := instrumentedResponseWriter{
 		ResponseWriter: w,
-		CloseNotifier:  closeNotifier,
 		Context:        ctx,
 	}
 
-	return irw, irw
+	if closeNotifier, ok := w.(http.CloseNotifier); ok {
+		irwCN := &instrumentedResponseWriterCN{
+			instrumentedResponseWriter: irw,
+			CloseNotifier:              closeNotifier,
+		}
+
+		return irwCN, irwCN
+	}
+
+	return &irw, &irw
 }
 
 // GetResponseWriter returns the http.ResponseWriter from the provided
@@ -263,11 +267,19 @@ func (ctx *muxVarsContext) Value(key interface{}) interface{} {
 	return ctx.Context.Value(key)
 }
 
+// instrumentedResponseWriterCN provides response writer information in a
+// context. It implements http.CloseNotifier so that users can detect
+// early disconnects.
+type instrumentedResponseWriterCN struct {
+	instrumentedResponseWriter
+	http.CloseNotifier
+}
+
 // instrumentedResponseWriter provides response writer information in a
-// context.
+// context. This variant is only used in the case where CloseNotifier is not
+// implemented by the parent ResponseWriter.
 type instrumentedResponseWriter struct {
 	http.ResponseWriter
-	http.CloseNotifier
 	Context
 
 	mu      sync.Mutex
diff --git a/context/http_test.go b/context/http_test.go
index ae88a314e..3d4b3c8eb 100644
--- a/context/http_test.go
+++ b/context/http_test.go
@@ -110,13 +110,6 @@ func (trw *testResponseWriter) Header() http.Header {
 	return trw.header
 }
 
-// CloseNotify is only here to make the testResponseWriter implement the
-// http.CloseNotifier interface, which WithResponseWriter expects to be
-// implemented.
-func (trw *testResponseWriter) CloseNotify() <-chan bool {
-	return make(chan bool)
-}
-
 func (trw *testResponseWriter) Write(p []byte) (n int, err error) {
 	if trw.status == 0 {
 		trw.status = http.StatusOK
diff --git a/registry/handlers/helpers.go b/registry/handlers/helpers.go
index 1f9a8ee1a..a4f3abcc0 100644
--- a/registry/handlers/helpers.go
+++ b/registry/handlers/helpers.go
@@ -29,7 +29,7 @@ func copyFullPayload(responseWriter http.ResponseWriter, r *http.Request, destWr
 	if notifier, ok := responseWriter.(http.CloseNotifier); ok {
 		clientClosed = notifier.CloseNotify()
 	} else {
-		panic("the ResponseWriter does not implement CloseNotifier")
+		ctxu.GetLogger(context).Warn("the ResponseWriter does not implement CloseNotifier")
 	}
 
 	// Read in the data, if any.