From 33e9d1e61375db830d3ffc2465acf7033ba84847 Mon Sep 17 00:00:00 2001
From: Muyao CHEN <muyao.chen@intersec.com>
Date: Sun, 15 Sep 2024 22:38:41 +0200
Subject: [PATCH] Context: use timeout in JSON writer

---
 framework/context.go | 16 +++++++++++++++-
 1 file changed, 15 insertions(+), 1 deletion(-)

diff --git a/framework/context.go b/framework/context.go
index ccfec52..fc41c22 100644
--- a/framework/context.go
+++ b/framework/context.go
@@ -28,14 +28,21 @@ type Context struct {
 // NewContext create a new context
 func NewContext(w http.ResponseWriter, r *http.Request) *Context {
 	return &Context{
+		ctx:            r.Context(),
 		request:        r,
 		responseWriter: w,
+		writerMux:      &sync.Mutex{},
 	}
 }
 
 // {{{ Basic functions
 
-// WriterMux returns the writer mutex
+// WriterMux returns the writer mutex.
+//
+// This is useful when goroutines concurrently write into responseWriter,
+// while at the same time we are writing into the responseWriter for a
+// panic or timeout.
+// We can protect it at the context level.
 func (ctx *Context) WriterMux() *sync.Mutex {
 	return ctx.writerMux
 }
@@ -51,6 +58,8 @@ func (ctx *Context) GetResponseWriter() http.ResponseWriter {
 }
 
 // SetHasTimeout indicates that the context has timeout.
+//
+// So that other goroutines won't write into the responseWriter anymore
 func (ctx *Context) SetHasTimeout() {
 	ctx.hasTimeout = true
 }
@@ -237,6 +246,11 @@ func (ctx *Context) ReadJSON(obj any) error {
 
 // WriteJSON send back an object in JSON format with the status code
 func (ctx *Context) WriteJSON(status int, obj any) error {
+	// There is a timeout, some error message data must have already been
+	// written to the output. Stop writing anything into the responseWriter.
+	if ctx.HasTimeout() {
+		return nil
+	}
 	data, err := json.Marshal(obj)
 	if err != nil {
 		return err