unify code-paths

mafredri · mafredri · commit 8411d318c4c1 · 2025-02-12T18:31:06.000Z
diff --git a/coderd/agentapi/audit.go b/coderd/agentapi/audit.go
@@ -34,28 +34,13 @@ func (a *AuditAPI) ReportConnection(ctx context.Context, req *agentproto.ReportC
 		return nil, xerrors.Errorf("connection id from bytes: %w", err)
 	}
 
-	var action database.AuditAction
-	switch req.GetConnection().GetAction() {
-	case agentproto.Connection_CONNECT:
-		action = database.AuditActionConnect
-	case agentproto.Connection_DISCONNECT:
-		action = database.AuditActionDisconnect
-	default:
-		return nil, xerrors.Errorf("unknown agent connection action %q", req.GetConnection().GetAction())
+	action, err := AgentProtoConnectionActionToAuditAction(req.GetConnection().GetAction())
+	if err != nil {
+		return nil, err
 	}
-
-	var connectionType agentsdk.ConnectionType
-	switch req.GetConnection().GetType() {
-	case agentproto.Connection_SSH:
-		connectionType = agentsdk.ConnectionTypeSSH
-	case agentproto.Connection_VSCODE:
-		connectionType = agentsdk.ConnectionTypeVSCode
-	case agentproto.Connection_JETBRAINS:
-		connectionType = agentsdk.ConnectionTypeJetBrains
-	case agentproto.Connection_RECONNECTING_PTY:
-		connectionType = agentsdk.ConnectionTypeReconnectingPTY
-	default:
-		return nil, xerrors.Errorf("unknown agent connection type %q", req.GetConnection().GetType())
+	connectionType, err := AgentProtoConnectionTypeToAgentConnectionType(req.GetConnection().GetType())
+	if err != nil {
+		return nil, err
 	}
 
 	// Fetch contextual data for this audit event.
@@ -115,3 +100,29 @@ func (a *AuditAPI) ReportConnection(ctx context.Context, req *agentproto.ReportC
 
 	return &emptypb.Empty{}, nil
 }
+
+func AgentProtoConnectionActionToAuditAction(action agentproto.Connection_Action) (database.AuditAction, error) {
+	switch action {
+	case agentproto.Connection_CONNECT:
+		return database.AuditActionConnect, nil
+	case agentproto.Connection_DISCONNECT:
+		return database.AuditActionDisconnect, nil
+	default:
+		return "", xerrors.Errorf("unknown agent connection action %q", action)
+	}
+}
+
+func AgentProtoConnectionTypeToAgentConnectionType(typ agentproto.Connection_Type) (agentsdk.ConnectionType, error) {
+	switch typ {
+	case agentproto.Connection_SSH:
+		return agentsdk.ConnectionTypeSSH, nil
+	case agentproto.Connection_VSCODE:
+		return agentsdk.ConnectionTypeVSCode, nil
+	case agentproto.Connection_JETBRAINS:
+		return agentsdk.ConnectionTypeJetBrains, nil
+	case agentproto.Connection_RECONNECTING_PTY:
+		return agentsdk.ConnectionTypeReconnectingPTY, nil
+	default:
+		return "", xerrors.Errorf("unknown agent connection type %q", typ)
+	}
+}
diff --git a/coderd/agentapi/audit_test.go b/coderd/agentapi/audit_test.go
@@ -147,32 +147,16 @@ func TestAuditReport(t *testing.T) {
 	}
 }
 
-func agentProtoConnectionActionToAudit(t *testing.T, typ agentproto.Connection_Action) database.AuditAction {
-	switch typ {
-	case agentproto.Connection_CONNECT:
-		return database.AuditActionConnect
-	case agentproto.Connection_DISCONNECT:
-		return database.AuditActionDisconnect
-	default:
-		t.Fatalf("unknown agent connection action %q", typ)
-		return ""
-	}
+func agentProtoConnectionActionToAudit(t *testing.T, action agentproto.Connection_Action) database.AuditAction {
+	a, err := agentapi.AgentProtoConnectionActionToAuditAction(action)
+	require.NoError(t, err)
+	return a
 }
 
 func agentProtoConnectionTypeToSDK(t *testing.T, typ agentproto.Connection_Type) agentsdk.ConnectionType {
-	switch typ {
-	case agentproto.Connection_SSH:
-		return agentsdk.ConnectionTypeSSH
-	case agentproto.Connection_VSCODE:
-		return agentsdk.ConnectionTypeVSCode
-	case agentproto.Connection_JETBRAINS:
-		return agentsdk.ConnectionTypeJetBrains
-	case agentproto.Connection_RECONNECTING_PTY:
-		return agentsdk.ConnectionTypeReconnectingPTY
-	default:
-		t.Fatalf("unknown agent connection type %q", typ)
-		return ""
-	}
+	action, err := agentapi.AgentProtoConnectionTypeToAgentConnectionType(typ)
+	require.NoError(t, err)
+	return action
 }
 
 func asAtomicPointer[T any](v T) *atomic.Pointer[T] {
