Fixed Response#Before()

Implemented Response#After()

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

response.go

@@ -4,6 +4,7 @@ import (
 	"bufio"
 	"net"
 	"net/http"
+	"strconv"
 )
 
 type (
@@ -11,12 +12,14 @@ type (
 	// by an HTTP handler to construct an HTTP response.
 	// See: https://golang.org/pkg/net/http/#ResponseWriter
 	Response struct {
-		echo        *Echo
-		beforeFuncs []func()
-		Writer      http.ResponseWriter
-		Status      int
-		Size        int64
-		Committed   bool
+		echo          *Echo
+		contentLength int64
+		beforeFuncs   []func()
+		afterFuncs    []func()
+		Writer        http.ResponseWriter
+		Status        int
+		Size          int64
+		Committed     bool
 	}
 )
 
@@ -40,6 +43,12 @@ func (r *Response) Before(fn func()) {
 	r.beforeFuncs = append(r.beforeFuncs, fn)
 }
 
+// After registers a function which is called just after the response is written.
+// If the `Content-Length` is unknown, none of the after function is executed.
+func (r *Response) After(fn func()) {
+	r.afterFuncs = append(r.afterFuncs, 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
@@ -55,6 +64,7 @@ func (r *Response) WriteHeader(code int) {
 	r.Status = code
 	r.Writer.WriteHeader(code)
 	r.Committed = true
+	r.contentLength, _ = strconv.ParseInt(r.Header().Get(HeaderContentLength), 10, 0)
 }
 
 // Write writes the data to the connection as part of an HTTP reply.
@@ -64,6 +74,11 @@ func (r *Response) Write(b []byte) (n int, err error) {
 	}
 	n, err = r.Writer.Write(b)
 	r.Size += int64(n)
+	if r.Size == r.contentLength {
+		for _, fn := range r.afterFuncs {
+			fn()
+		}
+	}
 	return
 }
 
@@ -91,6 +106,9 @@ func (r *Response) CloseNotify() <-chan bool {
 }
 
 func (r *Response) reset(w http.ResponseWriter) {
+	r.contentLength = 0
+	r.beforeFuncs = nil
+	r.afterFuncs = nil
 	r.Writer = w
 	r.Size = 0
 	r.Status = http.StatusOK