Add client metadata to PostgreSQL database session start event (#50711)

* feat: add client metadata to database session start event

* refactor(test): update option to WithUserAgent

Author: gabrielcorado

api/proto/teleport/legacy/types/events/events.proto

@@ -3150,6 +3150,12 @@ message DatabaseSessionStart {
   // connection. This can be useful for backend process cancellation or
   // termination and it is not a sensitive or secret value.
   uint32 PostgresPID = 8 [(gogoproto.jsontag) = "postgres_pid,omitempty"];
+  // Client is the common client event metadata.
+  ClientMetadata Client = 9 [
+    (gogoproto.nullable) = false,
+    (gogoproto.embed) = true,
+    (gogoproto.jsontag) = ""
+  ];
 }
 
 // DatabaseSessionQuery is emitted when a user executes a database query.

api/types/events/events.pb.go

@@ -1629,13 +1629,13 @@ func (c *testContext) startHandlingConnections() {
 
 // postgresClient connects to test Postgres through database access as a
 // specified Teleport user and database account.
-func (c *testContext) postgresClient(ctx context.Context, teleportUser, dbService, dbUser, dbName string) (*pgconn.PgConn, error) {
-	return c.postgresClientWithAddr(ctx, c.mux.DB().Addr().String(), teleportUser, dbService, dbUser, dbName)
+func (c *testContext) postgresClient(ctx context.Context, teleportUser, dbService, dbUser, dbName string, opts ...common.ClientOption) (*pgconn.PgConn, error) {
+	return c.postgresClientWithAddr(ctx, c.mux.DB().Addr().String(), teleportUser, dbService, dbUser, dbName, opts...)
 }
 
 // postgresClientWithAddr is like postgresClient but allows to override connection address.
-func (c *testContext) postgresClientWithAddr(ctx context.Context, address, teleportUser, dbService, dbUser, dbName string) (*pgconn.PgConn, error) {
-	return postgres.MakeTestClient(ctx, common.TestClientConfig{
+func (c *testContext) postgresClientWithAddr(ctx context.Context, address, teleportUser, dbService, dbUser, dbName string, opts ...common.ClientOption) (*pgconn.PgConn, error) {
+	cfg := common.TestClientConfig{
 		AuthClient: c.authClient,
 		AuthServer: c.authServer,
 		Address:    address,
@@ -1647,7 +1647,13 @@ func (c *testContext) postgresClientWithAddr(ctx context.Context, address, telep
 			Username:    dbUser,
 			Database:    dbName,
 		},
-	})
+	}
+
+	for _, opt := range opts {
+		opt(&cfg)
+	}
+
+	return postgres.MakeTestClient(ctx, cfg)
 }
 
 // postgresClientLocalProxy connects to test Postgres through local ALPN proxy.

lib/srv/db/access_test.go

@@ -35,6 +35,7 @@ import (
 	"github.com/gravitational/teleport/lib/defaults"
 	libevents "github.com/gravitational/teleport/lib/events"
 	"github.com/gravitational/teleport/lib/events/eventstest"
+	"github.com/gravitational/teleport/lib/srv/db/common"
 	"github.com/gravitational/teleport/lib/srv/db/postgres"
 	"github.com/gravitational/teleport/lib/srv/db/redis"
 )
@@ -59,12 +60,14 @@ func TestAuditPostgres(t *testing.T) {
 	requireEvent(t, testCtx, libevents.DatabaseSessionStartFailureCode)
 
 	// Connect should trigger successful session start event.
-	psql, err := testCtx.postgresClient(ctx, "alice", "postgres", "postgres", "postgres")
+	userAgent := "psql"
+	psql, err := testCtx.postgresClient(ctx, "alice", "postgres", "postgres", "postgres", common.WithUserAgent(userAgent))
 	require.NoError(t, err)
 	startEvt, ok := requireEvent(t, testCtx, libevents.DatabaseSessionStartCode).(*events.DatabaseSessionStart)
 	require.True(t, ok)
 	require.NotNil(t, startEvt)
 	require.NotZero(t, startEvt.PostgresPID)
+	require.Equal(t, userAgent, startEvt.ClientMetadata.UserAgent)
 
 	// Simple query should trigger the query event.
 	_, err = psql.Exec(ctx, "select 1").ReadAll()

lib/srv/db/audit_test.go

@@ -146,6 +146,9 @@ func (a *audit) OnSessionStart(ctx context.Context, session *Session, sessionErr
 			Success: true,
 		},
 		PostgresPID: session.PostgresPID,
+		ClientMetadata: events.ClientMetadata{
+			UserAgent: session.UserAgent,
+		},
 	}
 	event.SetTime(session.StartTime)
 

lib/srv/db/common/audit.go

@@ -68,6 +68,8 @@ type Session struct {
 	StartTime time.Time
 	// PostgresPID is the Postgres backend PID for the session.
 	PostgresPID uint32
+	// UserAgent identifies the type of client used on the session.
+	UserAgent string
 }
 
 // String returns string representation of the session parameters.

lib/srv/db/common/session.go

@@ -160,6 +160,16 @@ func MakeTestServerTLSConfig(config TestServerConfig) (*tls.Config, error) {
 	}, nil
 }
 
+// ClientOption represents a database client config option.
+type ClientOption func(config *TestClientConfig)
+
+// WithUserAgent set client user agent.
+func WithUserAgent(userAgent string) ClientOption {
+	return func(config *TestClientConfig) {
+		config.UserAgent = userAgent
+	}
+}
+
 // TestClientConfig combines parameters for a test Postgres/MySQL client.
 type TestClientConfig struct {
 	// AuthClient will be used to retrieve trusted CA.
@@ -176,6 +186,8 @@ type TestClientConfig struct {
 	PinnedIP string
 	// RouteToDatabase contains database routing information.
 	RouteToDatabase tlsca.RouteToDatabase
+	// UserAgent contains the client user agent.
+	UserAgent string
 }
 
 // MakeTestClientTLSCert returns TLS certificate suitable for configuring test

lib/srv/db/common/test.go

@@ -213,6 +213,9 @@ func (e *Engine) handleStartup(client *pgproto3.Backend, sessionCtx *common.Sess
 				sessionCtx.DatabaseName = value
 			case "user":
 				sessionCtx.DatabaseUser = value
+			// https://www.postgresql.org/docs/17/libpq-connect.html#LIBPQ-CONNECT-APPLICATION-NAME
+			case "application_name":
+				sessionCtx.UserAgent = value
 			default:
 				sessionCtx.StartupParameters[key] = value
 			}

lib/srv/db/postgres/engine.go

@@ -58,6 +58,9 @@ func MakeTestClient(ctx context.Context, config common.TestClientConfig) (*pgcon
 	if err != nil {
 		return nil, trace.Wrap(err)
 	}
+	if config.UserAgent != "" {
+		pgconnConfig.RuntimeParams["application_name"] = config.UserAgent
+	}
 	pgConn, err := pgconn.ConnectConfig(ctx, pgconnConfig)
 	if err != nil {
 		return nil, trace.Wrap(err)