commit audit in Issue, revert tracing status writer changes

mafredri · mafredri · commit 8a61541cc306 · 2025-03-10T10:14:12.000Z
diff --git a/coderd/tracing/status_writer.go b/coderd/tracing/status_writer.go
@@ -27,7 +27,6 @@ type StatusWriter struct {
 	http.ResponseWriter
 	Status       int
 	Hijacked     bool
-	doneFuncs    []func() // If non-nil, this function will be called when the handler is done.
 	responseBody []byte
 
 	wroteHeader      bool
@@ -38,17 +37,9 @@ func StatusWriterMiddleware(next http.Handler) http.Handler {
 	return http.HandlerFunc(func(rw http.ResponseWriter, r *http.Request) {
 		sw := &StatusWriter{ResponseWriter: rw}
 		next.ServeHTTP(sw, r)
-		for _, done := range sw.doneFuncs {
-			done()
-		}
 	})
 }
 
-func (w *StatusWriter) AddDoneFunc(f func()) {
-	// Prepend, as if deferred.
-	w.doneFuncs = append([]func(){f}, w.doneFuncs...)
-}
-
 func (w *StatusWriter) WriteHeader(status int) {
 	if buildinfo.IsDev() || flag.Lookup("test.v") != nil {
 		if w.wroteHeader {
diff --git a/coderd/tracing/status_writer_test.go b/coderd/tracing/status_writer_test.go
@@ -121,21 +121,16 @@ func TestStatusWriter(t *testing.T) {
 		t.Parallel()
 
 		var (
-			sw   *tracing.StatusWriter
-			done = false
-			rr   = httptest.NewRecorder()
+			sw *tracing.StatusWriter
+			rr = httptest.NewRecorder()
 		)
 		tracing.StatusWriterMiddleware(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 			sw = w.(*tracing.StatusWriter)
-			sw.AddDoneFunc(func() {
-				done = true
-			})
 			w.WriteHeader(http.StatusNoContent)
 		})).ServeHTTP(rr, httptest.NewRequest("GET", "/", nil))
 
 		require.Equal(t, http.StatusNoContent, rr.Code, "rr status code not set")
 		require.Equal(t, http.StatusNoContent, sw.Status, "sw status code not set")
-		require.True(t, done, "done func not called")
 	})
 }
 
diff --git a/coderd/workspaceapps/db.go b/coderd/workspaceapps/db.go
@@ -96,7 +96,8 @@ func (p *DBTokenProvider) Issue(ctx context.Context, rw http.ResponseWriter, r *
 	//                 // permissions.
 	dangerousSystemCtx := dbauthz.AsSystemRestricted(ctx)
 
-	aReq := p.auditInitAutocommitRequest(ctx, rw, r)
+	aReq, commitAudit := p.auditInitRequest(ctx, rw, r)
+	defer commitAudit()
 
 	appReq := issueReq.AppRequest.Normalize()
 	err := appReq.Check()
@@ -371,14 +372,14 @@ type auditRequest struct {
 	dbReq  *databaseRequest
 }
 
-// auditInitAutocommitRequest creates a new audit session and audit log for the
-// given request, if one does not already exist. If an audit session already
-// exists, it will be updated with the current timestamp. A session is used to
-// reduce the number of audit logs created.
+// auditInitRequest creates a new audit session and audit log for the given
+// request, if one does not already exist. If an audit session already exists,
+// it will be updated with the current timestamp. A session is used to reduce
+// the number of audit logs created.
 //
 // A session is unique to the agent, app, user and users IP. If any of these
 // values change, a new session and audit log is created.
-func (p *DBTokenProvider) auditInitAutocommitRequest(ctx context.Context, w http.ResponseWriter, r *http.Request) (aReq *auditRequest) {
+func (p *DBTokenProvider) auditInitRequest(ctx context.Context, w http.ResponseWriter, r *http.Request) (aReq *auditRequest, commit func()) {
 	// Get the status writer from the request context so we can figure
 	// out the HTTP status and autocommit the audit log.
 	sw, ok := w.(*tracing.StatusWriter)
@@ -393,7 +394,13 @@ func (p *DBTokenProvider) auditInitAutocommitRequest(ctx context.Context, w http
 
 	// Set the commit function on the status writer to create an audit
 	// log, this ensures that the status and response body are available.
-	sw.AddDoneFunc(func() {
+	var committed bool
+	return aReq, func() {
+		if committed {
+			return
+		}
+		committed = true
+
 		if sw.Status == http.StatusSeeOther {
 			// Redirects aren't interesting as we will capture the audit
 			// log after the redirect.
@@ -548,7 +555,5 @@ func (p *DBTokenProvider) auditInitAutocommitRequest(ctx context.Context, w http
 				AdditionalFields: appInfoBytes,
 			})
 		}
-	})
-
-	return aReq
+	}
 }
