fix(audit): flush pending audit logs before exit

audit/multi_auditor.go

@@ -5,10 +5,16 @@ import (
 	"fmt"
 	"log/slog"
 	"os"
+	"time"
 
 	"github.com/google/uuid"
 )
 
+// closeTimeout bounds how long the close function returned by SetupAuditor
+// waits for pending audit logs to reach the workspace agent, so an
+// unresponsive agent cannot block boundary from exiting.
+const closeTimeout = 5 * time.Second
+
 // MultiAuditor wraps multiple auditors and sends audit events to all of them.
 type MultiAuditor struct {
 	auditors []Auditor
@@ -30,13 +36,19 @@ func (m *MultiAuditor) AuditRequest(req Request) {
 // provided configuration. It always includes a LogAuditor for stderr logging,
 // and conditionally adds a SocketAuditor if audit logs are enabled and the
 // workspace agent's log proxy socket exists.
-func SetupAuditor(ctx context.Context, logger *slog.Logger, disableAuditLogs bool, logProxySocketPath string, sessionID uuid.UUID, confinedProcessName string) (Auditor, error) {
+//
+// The returned close function stops the SocketAuditor and waits up to
+// closeTimeout for it to send logs that are still queued. Call it once the
+// proxy has stopped producing audit events and before the process exits;
+// otherwise logs from the last batch window never reach coderd.
+func SetupAuditor(ctx context.Context, logger *slog.Logger, disableAuditLogs bool, logProxySocketPath string, sessionID uuid.UUID, confinedProcessName string) (Auditor, func(), error) {
 	stderrAuditor := NewLogAuditor(logger)
 	auditors := []Auditor{stderrAuditor}
+	closeAuditor := func() {}
 
 	if !disableAuditLogs {
 		if logProxySocketPath == "" {
-			return nil, fmt.Errorf("log proxy socket path is undefined")
+			return nil, nil, fmt.Errorf("log proxy socket path is undefined")
 		}
 		// Since boundary is separately versioned from a Coder deployment, it's possible
 		// Coder is on an older version that will not create the socket and listen for
@@ -46,12 +58,21 @@ func SetupAuditor(ctx context.Context, logger *slog.Logger, disableAuditLogs boo
 		// log proxy server before boundary runs.
 		_, err := os.Stat(logProxySocketPath)
 		if err != nil && !os.IsNotExist(err) {
-			return nil, fmt.Errorf("failed to stat log proxy socket: %v", err)
+			return nil, nil, fmt.Errorf("failed to stat log proxy socket: %v", err)
 		}
 		agentWillProxy := !os.IsNotExist(err)
 		if agentWillProxy {
 			socketAuditor := NewSocketAuditor(logger, logProxySocketPath, sessionID, confinedProcessName)
-			go socketAuditor.Loop(ctx)
+			stop := runSocketAuditor(ctx, socketAuditor)
+			closeAuditor = func() {
+				// The caller's context may already be canceled at exit, so
+				// the wait gets its own deadline.
+				waitCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), closeTimeout)
+				defer cancel()
+				if err := stop(waitCtx); err != nil {
+					logger.Warn("Audit logs may not have been sent to the workspace agent before exit", "error", err)
+				}
+			}
 			auditors = append(auditors, socketAuditor)
 		} else {
 			logger.Warn("Audit logs are disabled; workspace agent has not created log proxy socket",
@@ -61,5 +82,27 @@ func SetupAuditor(ctx context.Context, logger *slog.Logger, disableAuditLogs boo
 		logger.Warn("Audit logs are disabled by configuration")
 	}
 
-	return NewMultiAuditor(auditors...), nil
+	return NewMultiAuditor(auditors...), closeAuditor, nil
+}
+
+// runSocketAuditor runs s.Loop in the background and returns a function that
+// stops the loop and waits until it has flushed pending logs, or until waitCtx
+// is done.
+func runSocketAuditor(ctx context.Context, s *SocketAuditor) func(waitCtx context.Context) error {
+	loopCtx, cancel := context.WithCancel(ctx)
+	done := make(chan struct{})
+	go func() {
+		defer close(done)
+		s.Loop(loopCtx)
+	}()
+
+	return func(waitCtx context.Context) error {
+		cancel()
+		select {
+		case <-done:
+			return nil
+		case <-waitCtx.Done():
+			return waitCtx.Err()
+		}
+	}
 }

audit/multi_auditor_test.go

@@ -2,13 +2,18 @@ package audit
 
 import (
 	"context"
+	"errors"
 	"io"
 	"log/slog"
+	"net"
 	"os"
 	"path/filepath"
 	"testing"
+	"time"
 
 	"github.com/google/uuid"
+
+	"github.com/coder/coder/v2/agent/boundarylogproxy/codec"
 )
 
 type mockAuditor struct {
@@ -27,10 +32,12 @@ func TestSetupAuditor_DisabledAuditLogs(t *testing.T) {
 	logger := slog.New(slog.NewTextHandler(io.Discard, nil))
 	ctx := context.Background()
 
-	auditor, err := SetupAuditor(ctx, logger, true, "", uuid.New(), "")
+	auditor, closeAuditor, err := SetupAuditor(ctx, logger, true, "", uuid.New(), "")
 	if err != nil {
 		t.Fatalf("unexpected error: %v", err)
 	}
+	// Without a socket auditor there is nothing to flush.
+	closeAuditor()
 
 	multi, ok := auditor.(*MultiAuditor)
 	if !ok {
@@ -52,7 +59,7 @@ func TestSetupAuditor_EmptySocketPath(t *testing.T) {
 	logger := slog.New(slog.NewTextHandler(io.Discard, nil))
 	ctx := context.Background()
 
-	_, err := SetupAuditor(ctx, logger, false, "", uuid.New(), "")
+	_, _, err := SetupAuditor(ctx, logger, false, "", uuid.New(), "")
 	if err == nil {
 		t.Fatal("expected error for empty socket path, got nil")
 	}
@@ -64,7 +71,7 @@ func TestSetupAuditor_SocketDoesNotExist(t *testing.T) {
 	logger := slog.New(slog.NewTextHandler(io.Discard, nil))
 	ctx := context.Background()
 
-	auditor, err := SetupAuditor(ctx, logger, false, "/nonexistent/socket/path", uuid.New(), "")
+	auditor, _, err := SetupAuditor(ctx, logger, false, "/nonexistent/socket/path", uuid.New(), "")
 	if err != nil {
 		t.Fatalf("unexpected error: %v", err)
 	}
@@ -102,7 +109,7 @@ func TestSetupAuditor_SocketExists(t *testing.T) {
 		t.Fatalf("failed to close temp file: %v", err)
 	}
 
-	auditor, err := SetupAuditor(ctx, logger, false, socketPath, uuid.New(), "claude")
+	auditor, _, err := SetupAuditor(ctx, logger, false, socketPath, uuid.New(), "claude")
 	if err != nil {
 		t.Fatalf("unexpected error: %v", err)
 	}
@@ -142,3 +149,94 @@ func TestMultiAuditor_AuditRequest(t *testing.T) {
 		t.Error("expected second auditor to be called")
 	}
 }
+
+// TestSetupAuditor_CloseSendsQueuedLogs covers a wrapped command that exits
+// right after a request: the log is still queued for the next batch, and
+// close must send it to the workspace agent before returning.
+func TestSetupAuditor_CloseSendsQueuedLogs(t *testing.T) {
+	t.Parallel()
+
+	socketPath := filepath.Join(t.TempDir(), "agent.sock")
+	listener, err := net.Listen("unix", socketPath)
+	if err != nil {
+		t.Fatalf("failed to listen on unix socket: %v", err)
+	}
+	t.Cleanup(func() { _ = listener.Close() })
+
+	accepted := make(chan net.Conn, 1)
+	go func() {
+		conn, err := listener.Accept()
+		if err == nil {
+			accepted <- conn
+		}
+	}()
+
+	logger := slog.New(slog.NewTextHandler(io.Discard, nil))
+	auditor, closeAuditor, err := SetupAuditor(t.Context(), logger, false, socketPath, uuid.New(), "claude")
+	if err != nil {
+		t.Fatalf("unexpected error: %v", err)
+	}
+
+	// One request does not fill a batch, so without close it would wait for
+	// the batch timer.
+	auditor.AuditRequest(Request{Method: "GET", URL: "https://example.com/last", Allowed: true})
+	closeAuditor()
+
+	var conn net.Conn
+	select {
+	case conn = <-accepted:
+	case <-time.After(5 * time.Second):
+		t.Fatal("timeout waiting for the auditor to connect")
+	}
+	t.Cleanup(func() { _ = conn.Close() })
+
+	// Close returns only after the batch is written, so it is already in the
+	// socket buffer. The deadline only bounds the failure case.
+	if err := conn.SetReadDeadline(time.Now().Add(time.Second)); err != nil {
+		t.Fatalf("failed to set read deadline: %v", err)
+	}
+	msg, _, err := codec.ReadMessage(conn, make([]byte, 1<<10))
+	if err != nil {
+		t.Fatalf("expected the queued log to be sent before close returned: %v", err)
+	}
+	boundaryMsg, ok := msg.(*codec.BoundaryMessage)
+	if !ok {
+		t.Fatalf("expected *codec.BoundaryMessage, got %T", msg)
+	}
+	logs := boundaryMsg.GetLogs().GetLogs()
+	if len(logs) != 1 {
+		t.Fatalf("expected 1 log, got %d", len(logs))
+	}
+	if got := logs[0].GetHttpRequest().GetUrl(); got != "https://example.com/last" {
+		t.Errorf("expected URL https://example.com/last, got %s", got)
+	}
+}
+
+// TestRunSocketAuditor_StopIsBoundedByContext verifies that an unresponsive
+// workspace agent cannot block stop past its context deadline.
+func TestRunSocketAuditor_StopIsBoundedByContext(t *testing.T) {
+	t.Parallel()
+
+	unblockDial := make(chan struct{})
+	auditor := setupSocketAuditor(t)
+	auditor.dial = func() (net.Conn, error) {
+		<-unblockDial
+		return nil, errors.New("agent not responding")
+	}
+
+	stop := runSocketAuditor(t.Context(), auditor)
+
+	ctx, cancel := context.WithTimeout(t.Context(), 50*time.Millisecond)
+	defer cancel()
+	if err := stop(ctx); !errors.Is(err, context.DeadlineExceeded) {
+		t.Fatalf("expected context.DeadlineExceeded, got %v", err)
+	}
+
+	// Once the agent responds, the loop finishes and stop returns nil.
+	close(unblockDial)
+	waitCtx, waitCancel := context.WithTimeout(t.Context(), 5*time.Second)
+	defer waitCancel()
+	if err := stop(waitCtx); err != nil {
+		t.Fatalf("expected the loop to finish, got %v", err)
+	}
+}

landjail/parent.go

@@ -27,10 +27,13 @@ func RunParent(ctx context.Context, logger *slog.Logger, config config.AppConfig
 	ruleEngine := rulesengine.NewRuleEngine(allowRules, logger)
 
 	// Create auditor
-	auditor, err := audit.SetupAuditor(ctx, logger, config.DisableAuditLogs, config.LogProxySocketPath, config.SessionID, config.ConfinedProcessName)
+	auditor, closeAuditor, err := audit.SetupAuditor(ctx, logger, config.DisableAuditLogs, config.LogProxySocketPath, config.SessionID, config.ConfinedProcessName)
 	if err != nil {
 		return fmt.Errorf("failed to setup auditor: %v", err)
 	}
+	// Deferred calls run after Run has stopped the proxy, so no audit events
+	// arrive after the final flush.
+	defer closeAuditor()
 
 	// Create TLS certificate manager
 	certManager, err := tls.NewCertificateManager(tls.Config{

nsjail_manager/parent.go

@@ -28,10 +28,13 @@ func RunParent(ctx context.Context, logger *slog.Logger, config config.AppConfig
 	ruleEngine := rulesengine.NewRuleEngine(allowRules, logger)
 
 	// Create auditor
-	auditor, err := audit.SetupAuditor(ctx, logger, config.DisableAuditLogs, config.LogProxySocketPath, config.SessionID, config.ConfinedProcessName)
+	auditor, closeAuditor, err := audit.SetupAuditor(ctx, logger, config.DisableAuditLogs, config.LogProxySocketPath, config.SessionID, config.ConfinedProcessName)
 	if err != nil {
 		return fmt.Errorf("failed to setup auditor: %v", err)
 	}
+	// Deferred calls run after Run has stopped the proxy, so no audit events
+	// arrive after the final flush.
+	defer closeAuditor()
 
 	// Create TLS certificate manager
 	certManager, err := tls.NewCertificateManager(tls.Config{