Prevent multiple Set-Cookie headers when calling RegenerateToken

Closes #61

Author: aeneasr

context.go

@@ -16,6 +16,12 @@ type csrfContext struct {
 	token string
 	// reason for the failure of CSRF check
 	reason error
+	// wasSent is true if `Set-Cookie` was called
+	// for the `name=csrf_token` already. This prevents
+	// duplicate `Set-Cookie: csrf_token` headers.
+	// For more information see:
+	// https://github.com/justinas/nosurf/pull/61
+	wasSent bool
 }
 
 // Token takes an HTTP request and returns
@@ -53,6 +59,17 @@ func ctxSetToken(req *http.Request, token []byte) {
 	ctx.token = b64encode(maskToken(token))
 }
 
+func ctxSetSent(req *http.Request) {
+	ctx := req.Context().Value(nosurfKey).(*csrfContext)
+	ctx.wasSent = true
+}
+
+func ctxWasSent(req *http.Request) bool {
+	ctx := req.Context().Value(nosurfKey).(*csrfContext)
+
+	return ctx.wasSent
+}
+
 func ctxSetReason(req *http.Request, reason error) {
 	ctx := req.Context().Value(nosurfKey).(*csrfContext)
 	if ctx.token == "" {

context_legacy.go

@@ -17,6 +17,12 @@ type csrfContext struct {
 	token string
 	// reason for the failure of CSRF check
 	reason error
+	// wasSent is true if `Set-Cookie` was called
+	// for the `name=csrf_token` already. This prevents
+	// duplicate `Set-Cookie: csrf_token` headers.
+	// For more information see:
+	// https://github.com/justinas/nosurf/pull/61
+	wasSent bool
 }
 
 var (
@@ -79,6 +85,32 @@ func ctxSetToken(req *http.Request, token []byte) *http.Request {
 	return req
 }
 
+func ctxSetSent(req *http.Request) {
+	cmMutex.Lock()
+	defer cmMutex.Unlock()
+
+	ctx, ok := contextMap[req]
+	if !ok {
+		ctx = new(csrfContext)
+		contextMap[req] = ctx
+	}
+
+	ctx.wasSent = true
+}
+
+func ctxWasSent(req *http.Request) bool {
+	cmMutex.RLock()
+	defer cmMutex.RUnlock()
+
+	ctx, ok := contextMap[req]
+
+	if !ok {
+		return false
+	}
+
+	return ctx.wasSent
+}
+
 func ctxSetReason(req *http.Request, reason error) *http.Request {
 	cmMutex.Lock()
 	defer cmMutex.Unlock()

handler.go

@@ -195,6 +195,16 @@ func (h *CSRFHandler) handleFailure(w http.ResponseWriter, r *http.Request) {
 
 // Generates a new token, sets it on the given request and returns it
 func (h *CSRFHandler) RegenerateToken(w http.ResponseWriter, r *http.Request) string {
+	if ctxWasSent(r) {
+		// The CSRF Cookie was set already by an earlier call to `RegenerateToken`
+		// in the same request context. It therefore does not make sense to regenerate
+		// it again as it will lead to two or more `Set-Cookie` instructions which will in turn
+		// cause CSRF to fail depending on the resulting order of the `Set-Cookie` instructions.
+		//
+		// No warning is necessary as the only caller to `setTokenCookie` is `RegenerateToken`.
+		return Token(r)
+	}
+
 	token := generateToken()
 	h.setTokenCookie(w, r, token)
 
@@ -210,6 +220,7 @@ func (h *CSRFHandler) setTokenCookie(w http.ResponseWriter, r *http.Request, tok
 	cookie.Value = b64encode(token)
 
 	http.SetCookie(w, &cookie)
+	ctxSetSent(r)
 
 }