Fixed #176, closed #831

Signed-off-by: Vishal Rana <vr@labstack.com>

README.md

@@ -1,4 +1,4 @@
-<a href="https://echo.labstack.com"><img class="logo" height="40" src="https://cdn.labstack.com/images/echo-logo.svg"></a>
+<a href="https://echo.labstack.com"><img height="80" src="https://cdn.labstack.com/images/echo-logo.svg"></a>
 
 [![Sourcegraph](https://sourcegraph.com/github.com/labstack/echo/-/badge.svg?style=flat-square)](https://sourcegraph.com/github.com/labstack/echo?badge)
 [![GoDoc](http://img.shields.io/badge/go-documentation-blue.svg?style=flat-square)](http://godoc.org/github.com/labstack/echo)

echo.go

@@ -72,14 +72,14 @@ type (
 		TLSServer        *http.Server
 		Listener         net.Listener
 		TLSListener      net.Listener
+		AutoTLSManager   autocert.Manager
 		DisableHTTP2     bool
 		Debug            bool
 		HideBanner       bool
 		HTTPErrorHandler HTTPErrorHandler
 		Binder           Binder
 		Validator        Validator
 		Renderer         Renderer
-		AutoTLSManager   autocert.Manager
 		// Mutex            sync.RWMutex
 		Logger Logger
 	}
@@ -293,14 +293,16 @@ func New() (e *Echo) {
 
 // NewContext returns a Context instance.
 func (e *Echo) NewContext(r *http.Request, w http.ResponseWriter) Context {
-	return &context{
+	c := &context{
 		request:  r,
-		response: NewResponse(w, e),
+		response: &Response{Writer: w},
 		store:    make(Map),
 		echo:     e,
 		pvalues:  make([]string, *e.maxParam),
 		handler:  NotFoundHandler,
 	}
+	c.response.context = c
+	return c
 }
 
 // Router returns router.

response.go

@@ -11,18 +11,18 @@ type (
 	// by an HTTP handler to construct an HTTP response.
 	// See: https://golang.org/pkg/net/http/#ResponseWriter
 	Response struct {
-		Writer    http.ResponseWriter
-		Status    int
-		Size      int64
-		Committed bool
-		echo      *Echo
+		context     Context
+		beforeFuncs []BeforeResponseFunc
+		Writer      http.ResponseWriter
+		Status      int
+		Size        int64
+		Committed   bool
 	}
-)
 
-// NewResponse creates a new instance of Response.
-func NewResponse(w http.ResponseWriter, e *Echo) (r *Response) {
-	return &Response{Writer: w, echo: e}
-}
+	// BeforeResponseFunc defines a function which is called just before writing the
+	// response.
+	BeforeResponseFunc func(Context)
+)
 
 // Header returns the header map for the writer that will be sent by
 // WriteHeader. Changing the header after a call to WriteHeader (or Write) has
@@ -34,15 +34,23 @@ func (r *Response) Header() http.Header {
 	return r.Writer.Header()
 }
 
+// Before registers a function which is called just before the response is written.
+func (r *Response) Before(fn BeforeResponseFunc) {
+	r.beforeFuncs = append(r.beforeFuncs, fn)
+}
+
 // WriteHeader sends an HTTP response header with status code. If WriteHeader is
 // not called explicitly, the first call to Write will trigger an implicit
 // WriteHeader(http.StatusOK). Thus explicit calls to WriteHeader are mainly
 // used to send error codes.
 func (r *Response) WriteHeader(code int) {
 	if r.Committed {
-		r.echo.Logger.Warn("response already committed")
+		r.context.Logger().Warn("response already committed")
 		return
 	}
+	// for _, fn := range r.beforeFuncs {
+	// 	fn(r.context)
+	// }
 	r.Status = code
 	r.Writer.WriteHeader(code)
 	r.Committed = true

response_test.go

@@ -0,0 +1,23 @@
+package echo
+
+import (
+	"net/http/httptest"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestResponse(t *testing.T) {
+	e := New()
+	req := httptest.NewRequest(GET, "/", nil)
+	rec := httptest.NewRecorder()
+	c := e.NewContext(req, rec)
+	res := &Response{context: c, Writer: rec}
+
+	// Before
+	res.Before(func(c Context) {
+		c.Response().Header().Set(HeaderServer, "echo")
+	})
+	res.Write([]byte("test"))
+	assert.Equal(t, "echo", rec.Header().Get(HeaderServer))
+}