1 files changed, 27 insertions, 2 deletions
diff --git a/modules/context/response.go b/modules/context/response.go
index 549bd30ee0..1881ec7b33 100644
--- a/modules/context/response.go
+++ b/modules/context/response.go
@@ -11,6 +11,7 @@ type ResponseWriter interface {
 	http.ResponseWriter
 	Flush()
 	Status() int
+	Before(func(ResponseWriter))
 }
 
 var (
@@ -20,11 +21,19 @@ var (
 // Response represents a response
 type Response struct {
 	http.ResponseWriter
-	status int
+	status         int
+	befores        []func(ResponseWriter)
+	beforeExecuted bool
 }
 
 // Write writes bytes to HTTP endpoint
 func (r *Response) Write(bs []byte) (int, error) {
+	if !r.beforeExecuted {
+		for _, before := range r.befores {
+			before(r)
+		}
+		r.beforeExecuted = true
+	}
 	size, err := r.ResponseWriter.Write(bs)
 	if err != nil {
 		return 0, err
@@ -37,6 +46,12 @@ func (r *Response) Write(bs []byte) (int, error) {
 
 // WriteHeader write status code
 func (r *Response) WriteHeader(statusCode int) {
+	if !r.beforeExecuted {
+		for _, before := range r.befores {
+			before(r)
+		}
+		r.beforeExecuted = true
+	}
 	r.status = statusCode
 	r.ResponseWriter.WriteHeader(statusCode)
 }
@@ -53,10 +68,20 @@ func (r *Response) Status() int {
 	return r.status
 }
 
+// Before allows for a function to be called before the ResponseWriter has been written to. This is
+// useful for setting headers or any other operations that must happen before a response has been written.
+func (r *Response) Before(f func(ResponseWriter)) {
+	r.befores = append(r.befores, f)
+}
+
 // NewResponse creates a response
 func NewResponse(resp http.ResponseWriter) *Response {
 	if v, ok := resp.(*Response); ok {
 		return v
 	}
-	return &Response{resp, 0}
+	return &Response{
+		ResponseWriter: resp,
+		status:         0,
+		befores:        make([]func(ResponseWriter), 0),
+	}
 }