Detect indirect superusers via role membership on PG 16+

collector/pg_unexpected_superusers.go

@@ -18,6 +18,7 @@ import (
 	"database/sql"
 	"log/slog"
 
+	"github.com/blang/semver/v4"
 	"github.com/prometheus/client_golang/prometheus"
 )
 
@@ -55,20 +56,54 @@ var (
 			"role",
 		),
 		"Unexpected superuser role (value is always 1)",
-		[]string{"rolname"}, nil,
+		[]string{"rolname", "access_type"}, nil,
 	)
 
 	// Roles that are expected to have superuser privileges.
 	expectedSuperusers = map[string]struct{}{
 		"pscale_admin": {},
 	}
 
-	pgUnexpectedSuperusersQuery = "SELECT rolname FROM pg_roles WHERE rolsuper"
+	// pgLargeRolesThreshold is the number of roles above which the expensive recursive
+	// CTE for indirect superuser detection is skipped in favour of the simple query.
+	pgLargeRolesThreshold = 1000
+
+	pgRoleCountQuery = "SELECT pg_catalog.count(*) FROM pg_catalog.pg_roles"
+
+	pgUnexpectedSuperusersQuery = "SELECT rolname, 'direct'::pg_catalog.text AS access_type FROM pg_catalog.pg_roles WHERE rolsuper"
+
+	pgUnexpectedSuperusersQueryPG16 = `WITH RECURSIVE superuser_oids AS (
+    SELECT oid FROM pg_catalog.pg_roles WHERE rolsuper
+    UNION
+    SELECT m.member
+    FROM pg_catalog.pg_auth_members m
+    JOIN superuser_oids s ON m.roleid OPERATOR(pg_catalog.=) s.oid
+    WHERE m.set_option OPERATOR(pg_catalog.=) true OR m.admin_option OPERATOR(pg_catalog.=) true
+)
+SELECT r.rolname,
+    CASE WHEN r.rolsuper
+        THEN 'direct'::pg_catalog.text
+        ELSE 'indirect'::pg_catalog.text
+    END AS access_type
+FROM superuser_oids so
+JOIN pg_catalog.pg_roles r ON r.oid OPERATOR(pg_catalog.=) so.oid`
 )
 
 func (c PGUnexpectedSuperusersCollector) Update(ctx context.Context, instance *Instance, ch chan<- prometheus.Metric) error {
 	db := instance.getDB()
-	rows, err := db.QueryContext(ctx, pgUnexpectedSuperusersQuery)
+
+	query := pgUnexpectedSuperusersQuery
+	if instance.version.GTE(semver.MustParse("16.0.0")) {
+		var roleCount int
+		if err := db.QueryRowContext(ctx, pgRoleCountQuery).Scan(&roleCount); err != nil {
+			return err
+		}
+		if roleCount < pgLargeRolesThreshold {
+			query = pgUnexpectedSuperusersQueryPG16
+		}
+	}
+
+	rows, err := db.QueryContext(ctx, query)
 	if err != nil {
 		return err
 	}
@@ -77,7 +112,8 @@ func (c PGUnexpectedSuperusersCollector) Update(ctx context.Context, instance *I
 	var count float64
 	for rows.Next() {
 		var rolname sql.NullString
-		if err := rows.Scan(&rolname); err != nil {
+		var accessType sql.NullString
+		if err := rows.Scan(&rolname, &accessType); err != nil {
 			return err
 		}
 
@@ -89,10 +125,15 @@ func (c PGUnexpectedSuperusersCollector) Update(ctx context.Context, instance *I
 			continue
 		}
 
+		accessTypeLabel := "direct"
+		if accessType.Valid {
+			accessTypeLabel = accessType.String
+		}
+
 		count++
 		ch <- prometheus.MustNewConstMetric(
 			pgUnexpectedSuperuserDesc,
-			prometheus.GaugeValue, 1, rolname.String,
+			prometheus.GaugeValue, 1, rolname.String, accessTypeLabel,
 		)
 	}
 

collector/pg_unexpected_superusers_test.go

@@ -17,6 +17,7 @@ import (
 	"testing"
 
 	"github.com/DATA-DOG/go-sqlmock"
+	"github.com/blang/semver/v4"
 	"github.com/prometheus/client_golang/prometheus"
 	dto "github.com/prometheus/client_model/go"
 	"github.com/smartystreets/goconvey/convey"
@@ -29,11 +30,11 @@ func TestPGUnexpectedSuperusersCollectorNoUnexpected(t *testing.T) {
 	}
 	defer db.Close()
 
-	inst := &Instance{db: db}
+	inst := &Instance{db: db, version: semver.MustParse("15.0.0")}
 
 	mock.ExpectQuery(sanitizeQuery(pgUnexpectedSuperusersQuery)).
-		WillReturnRows(sqlmock.NewRows([]string{"rolname"}).
-			AddRow("pscale_admin"))
+		WillReturnRows(sqlmock.NewRows([]string{"rolname", "access_type"}).
+			AddRow("pscale_admin", "direct"))
 
 	ch := make(chan prometheus.Metric)
 	go func() {
@@ -65,13 +66,13 @@ func TestPGUnexpectedSuperusersCollectorWithUnexpected(t *testing.T) {
 	}
 	defer db.Close()
 
-	inst := &Instance{db: db}
+	inst := &Instance{db: db, version: semver.MustParse("15.0.0")}
 
 	mock.ExpectQuery(sanitizeQuery(pgUnexpectedSuperusersQuery)).
-		WillReturnRows(sqlmock.NewRows([]string{"rolname"}).
-			AddRow("pscale_admin").
-			AddRow("rogue_admin").
-			AddRow("another_bad_user"))
+		WillReturnRows(sqlmock.NewRows([]string{"rolname", "access_type"}).
+			AddRow("pscale_admin", "direct").
+			AddRow("rogue_admin", "direct").
+			AddRow("another_bad_user", "direct"))
 
 	ch := make(chan prometheus.Metric)
 	go func() {
@@ -83,8 +84,8 @@ func TestPGUnexpectedSuperusersCollectorWithUnexpected(t *testing.T) {
 	}()
 
 	expected := []MetricResult{
-		{labels: labelMap{"rolname": "rogue_admin"}, value: 1, metricType: dto.MetricType_GAUGE},
-		{labels: labelMap{"rolname": "another_bad_user"}, value: 1, metricType: dto.MetricType_GAUGE},
+		{labels: labelMap{"rolname": "rogue_admin", "access_type": "direct"}, value: 1, metricType: dto.MetricType_GAUGE},
+		{labels: labelMap{"rolname": "another_bad_user", "access_type": "direct"}, value: 1, metricType: dto.MetricType_GAUGE},
 		{labels: labelMap{}, value: 2, metricType: dto.MetricType_GAUGE},
 	}
 	convey.Convey("Unexpected superusers detected", t, func() {
@@ -105,10 +106,10 @@ func TestPGUnexpectedSuperusersCollectorNoSuperusers(t *testing.T) {
 	}
 	defer db.Close()
 
-	inst := &Instance{db: db}
+	inst := &Instance{db: db, version: semver.MustParse("15.0.0")}
 
 	mock.ExpectQuery(sanitizeQuery(pgUnexpectedSuperusersQuery)).
-		WillReturnRows(sqlmock.NewRows([]string{"rolname"}))
+		WillReturnRows(sqlmock.NewRows([]string{"rolname", "access_type"}))
 
 	ch := make(chan prometheus.Metric)
 	go func() {
@@ -132,3 +133,167 @@ func TestPGUnexpectedSuperusersCollectorNoSuperusers(t *testing.T) {
 		t.Errorf("there were unfulfilled expectations: %s", err)
 	}
 }
+
+func TestPGUnexpectedSuperusersCollectorIndirectPG16(t *testing.T) {
+	db, mock, err := sqlmock.New()
+	if err != nil {
+		t.Fatalf("Error opening a stub db connection: %s", err)
+	}
+	defer db.Close()
+
+	inst := &Instance{db: db, version: semver.MustParse("16.0.0")}
+
+	mock.ExpectQuery(sanitizeQuery(pgRoleCountQuery)).
+		WillReturnRows(sqlmock.NewRows([]string{"count"}).AddRow(10))
+
+	mock.ExpectQuery(sanitizeQuery(pgUnexpectedSuperusersQueryPG16)).
+		WillReturnRows(sqlmock.NewRows([]string{"rolname", "access_type"}).
+			AddRow("pscale_admin", "direct").
+			AddRow("sneaky_user", "indirect"))
+
+	ch := make(chan prometheus.Metric)
+	go func() {
+		defer close(ch)
+		c := PGUnexpectedSuperusersCollector{}
+		if err := c.Update(context.Background(), inst, ch); err != nil {
+			t.Errorf("Error calling PGUnexpectedSuperusersCollector.Update: %s", err)
+		}
+	}()
+
+	expected := []MetricResult{
+		{labels: labelMap{"rolname": "sneaky_user", "access_type": "indirect"}, value: 1, metricType: dto.MetricType_GAUGE},
+		{labels: labelMap{}, value: 1, metricType: dto.MetricType_GAUGE},
+	}
+	convey.Convey("Indirect superuser detected on PG 16", t, func() {
+		for _, expect := range expected {
+			m := readMetric(<-ch)
+			convey.So(expect, convey.ShouldResemble, m)
+		}
+	})
+	if err := mock.ExpectationsWereMet(); err != nil {
+		t.Errorf("there were unfulfilled expectations: %s", err)
+	}
+}
+
+func TestPGUnexpectedSuperusersCollectorMixedPG16(t *testing.T) {
+	db, mock, err := sqlmock.New()
+	if err != nil {
+		t.Fatalf("Error opening a stub db connection: %s", err)
+	}
+	defer db.Close()
+
+	inst := &Instance{db: db, version: semver.MustParse("16.0.0")}
+
+	mock.ExpectQuery(sanitizeQuery(pgRoleCountQuery)).
+		WillReturnRows(sqlmock.NewRows([]string{"count"}).AddRow(10))
+
+	mock.ExpectQuery(sanitizeQuery(pgUnexpectedSuperusersQueryPG16)).
+		WillReturnRows(sqlmock.NewRows([]string{"rolname", "access_type"}).
+			AddRow("pscale_admin", "direct").
+			AddRow("rogue_admin", "direct").
+			AddRow("sneaky_user", "indirect"))
+
+	ch := make(chan prometheus.Metric)
+	go func() {
+		defer close(ch)
+		c := PGUnexpectedSuperusersCollector{}
+		if err := c.Update(context.Background(), inst, ch); err != nil {
+			t.Errorf("Error calling PGUnexpectedSuperusersCollector.Update: %s", err)
+		}
+	}()
+
+	expected := []MetricResult{
+		{labels: labelMap{"rolname": "rogue_admin", "access_type": "direct"}, value: 1, metricType: dto.MetricType_GAUGE},
+		{labels: labelMap{"rolname": "sneaky_user", "access_type": "indirect"}, value: 1, metricType: dto.MetricType_GAUGE},
+		{labels: labelMap{}, value: 2, metricType: dto.MetricType_GAUGE},
+	}
+	convey.Convey("Mix of direct and indirect unexpected superusers on PG 16", t, func() {
+		for _, expect := range expected {
+			m := readMetric(<-ch)
+			convey.So(expect, convey.ShouldResemble, m)
+		}
+	})
+	if err := mock.ExpectationsWereMet(); err != nil {
+		t.Errorf("there were unfulfilled expectations: %s", err)
+	}
+}
+
+func TestPGUnexpectedSuperusersCollectorExpectedIndirectFilteredPG16(t *testing.T) {
+	db, mock, err := sqlmock.New()
+	if err != nil {
+		t.Fatalf("Error opening a stub db connection: %s", err)
+	}
+	defer db.Close()
+
+	inst := &Instance{db: db, version: semver.MustParse("16.0.0")}
+
+	mock.ExpectQuery(sanitizeQuery(pgRoleCountQuery)).
+		WillReturnRows(sqlmock.NewRows([]string{"count"}).AddRow(10))
+
+	mock.ExpectQuery(sanitizeQuery(pgUnexpectedSuperusersQueryPG16)).
+		WillReturnRows(sqlmock.NewRows([]string{"rolname", "access_type"}).
+			AddRow("pscale_admin", "indirect"))
+
+	ch := make(chan prometheus.Metric)
+	go func() {
+		defer close(ch)
+		c := PGUnexpectedSuperusersCollector{}
+		if err := c.Update(context.Background(), inst, ch); err != nil {
+			t.Errorf("Error calling PGUnexpectedSuperusersCollector.Update: %s", err)
+		}
+	}()
+
+	expected := []MetricResult{
+		{labels: labelMap{}, value: 0, metricType: dto.MetricType_GAUGE},
+	}
+	convey.Convey("Expected superuser filtered even when indirect", t, func() {
+		for _, expect := range expected {
+			m := readMetric(<-ch)
+			convey.So(expect, convey.ShouldResemble, m)
+		}
+	})
+	if err := mock.ExpectationsWereMet(); err != nil {
+		t.Errorf("there were unfulfilled expectations: %s", err)
+	}
+}
+
+func TestPGUnexpectedSuperusersCollectorLargeRolesFallbackPG16(t *testing.T) {
+	db, mock, err := sqlmock.New()
+	if err != nil {
+		t.Fatalf("Error opening a stub db connection: %s", err)
+	}
+	defer db.Close()
+
+	inst := &Instance{db: db, version: semver.MustParse("16.0.0")}
+
+	mock.ExpectQuery(sanitizeQuery(pgRoleCountQuery)).
+		WillReturnRows(sqlmock.NewRows([]string{"count"}).AddRow(1000))
+
+	mock.ExpectQuery(sanitizeQuery(pgUnexpectedSuperusersQuery)).
+		WillReturnRows(sqlmock.NewRows([]string{"rolname", "access_type"}).
+			AddRow("pscale_admin", "direct").
+			AddRow("rogue_admin", "direct"))
+
+	ch := make(chan prometheus.Metric)
+	go func() {
+		defer close(ch)
+		c := PGUnexpectedSuperusersCollector{}
+		if err := c.Update(context.Background(), inst, ch); err != nil {
+			t.Errorf("Error calling PGUnexpectedSuperusersCollector.Update: %s", err)
+		}
+	}()
+
+	expected := []MetricResult{
+		{labels: labelMap{"rolname": "rogue_admin", "access_type": "direct"}, value: 1, metricType: dto.MetricType_GAUGE},
+		{labels: labelMap{}, value: 1, metricType: dto.MetricType_GAUGE},
+	}
+	convey.Convey("Falls back to simple query when role count >= 1000", t, func() {
+		for _, expect := range expected {
+			m := readMetric(<-ch)
+			convey.So(expect, convey.ShouldResemble, m)
+		}
+	})
+	if err := mock.ExpectationsWereMet(); err != nil {
+		t.Errorf("there were unfulfilled expectations: %s", err)
+	}
+}