db: refactor postgres test helpers (#27811)

* db: refactor postgres test helpers

* fix references to refactored test helper

* fix references to refactored test helper

* fix failing test

builtin/logical/database/backend_test.go

@@ -359,7 +359,7 @@ func TestBackend_BadConnectionString(t *testing.T) {
 	}
 	defer b.Cleanup(context.Background())
 
-	cleanup, _ := postgreshelper.PrepareTestContainer(t, "13.4-buster")
+	cleanup, _ := postgreshelper.PrepareTestContainer(t)
 	defer cleanup()
 
 	respCheck := func(req *logical.Request) {
@@ -410,7 +410,7 @@ func TestBackend_basic(t *testing.T) {
 	}
 	defer b.Cleanup(context.Background())
 
-	cleanup, connURL := postgreshelper.PrepareTestContainer(t, "13.4-buster")
+	cleanup, connURL := postgreshelper.PrepareTestContainer(t)
 	defer cleanup()
 
 	// Configure a connection
@@ -665,7 +665,7 @@ func TestBackend_connectionCrud(t *testing.T) {
 	dbFactory.sys = sys
 	client := cluster.Cores[0].Client.Logical()
 
-	cleanup, connURL := postgreshelper.PrepareTestContainer(t, "13.4-buster")
+	cleanup, connURL := postgreshelper.PrepareTestContainer(t)
 	defer cleanup()
 
 	// Mount the database plugin.
@@ -872,7 +872,7 @@ func TestBackend_roleCrud(t *testing.T) {
 	}
 	defer b.Cleanup(context.Background())
 
-	cleanup, connURL := postgreshelper.PrepareTestContainer(t, "13.4-buster")
+	cleanup, connURL := postgreshelper.PrepareTestContainer(t)
 	defer cleanup()
 
 	// Configure a connection
@@ -1121,7 +1121,7 @@ func TestBackend_allowedRoles(t *testing.T) {
 	}
 	defer b.Cleanup(context.Background())
 
-	cleanup, connURL := postgreshelper.PrepareTestContainer(t, "13.4-buster")
+	cleanup, connURL := postgreshelper.PrepareTestContainer(t)
 	defer cleanup()
 
 	// Configure a connection
@@ -1318,7 +1318,7 @@ func TestBackend_RotateRootCredentials(t *testing.T) {
 	}
 	defer b.Cleanup(context.Background())
 
-	cleanup, connURL := postgreshelper.PrepareTestContainer(t, "13.4-buster")
+	cleanup, connURL := postgreshelper.PrepareTestContainer(t)
 	defer cleanup()
 
 	connURL = strings.ReplaceAll(connURL, "postgres:secret", "{{username}}:{{password}}")

builtin/logical/database/path_roles_test.go

@@ -222,7 +222,7 @@ func TestBackend_StaticRole_Config(t *testing.T) {
 	}
 	defer b.Cleanup(context.Background())
 
-	cleanup, connURL := postgreshelper.PrepareTestContainer(t, "")
+	cleanup, connURL := postgreshelper.PrepareTestContainer(t)
 	defer cleanup()
 
 	// create the database user
@@ -487,7 +487,7 @@ func TestBackend_StaticRole_ReadCreds(t *testing.T) {
 	}
 	defer b.Cleanup(context.Background())
 
-	cleanup, connURL := postgreshelper.PrepareTestContainer(t, "")
+	cleanup, connURL := postgreshelper.PrepareTestContainer(t)
 	defer cleanup()
 
 	// create the database user
@@ -667,7 +667,7 @@ func TestBackend_StaticRole_Updates(t *testing.T) {
 	}
 	defer b.Cleanup(context.Background())
 
-	cleanup, connURL := postgreshelper.PrepareTestContainer(t, "")
+	cleanup, connURL := postgreshelper.PrepareTestContainer(t)
 	defer cleanup()
 
 	// create the database user
@@ -966,7 +966,7 @@ func TestBackend_StaticRole_Role_name_check(t *testing.T) {
 	}
 	defer b.Cleanup(context.Background())
 
-	cleanup, connURL := postgreshelper.PrepareTestContainer(t, "")
+	cleanup, connURL := postgreshelper.PrepareTestContainer(t)
 	defer cleanup()
 
 	// create the database user

builtin/logical/database/rollback_test.go

@@ -44,7 +44,7 @@ func TestBackend_RotateRootCredentials_WAL_rollback(t *testing.T) {
 	}
 	defer lb.Cleanup(context.Background())
 
-	cleanup, connURL := postgreshelper.PrepareTestContainer(t, "")
+	cleanup, connURL := postgreshelper.PrepareTestContainer(t)
 	defer cleanup()
 
 	connURL = strings.ReplaceAll(connURL, "postgres:secret", "{{username}}:{{password}}")
@@ -183,7 +183,7 @@ func TestBackend_RotateRootCredentials_WAL_no_rollback_1(t *testing.T) {
 	}
 	defer lb.Cleanup(context.Background())
 
-	cleanup, connURL := postgreshelper.PrepareTestContainer(t, "")
+	cleanup, connURL := postgreshelper.PrepareTestContainer(t)
 	defer cleanup()
 
 	connURL = strings.ReplaceAll(connURL, "postgres:secret", "{{username}}:{{password}}")
@@ -291,7 +291,7 @@ func TestBackend_RotateRootCredentials_WAL_no_rollback_2(t *testing.T) {
 	}
 	defer lb.Cleanup(context.Background())
 
-	cleanup, connURL := postgreshelper.PrepareTestContainer(t, "")
+	cleanup, connURL := postgreshelper.PrepareTestContainer(t)
 	defer cleanup()
 
 	connURL = strings.ReplaceAll(connURL, "postgres:secret", "{{username}}:{{password}}")

builtin/logical/database/rotation_test.go

@@ -63,7 +63,7 @@ func TestBackend_StaticRole_Rotation_basic(t *testing.T) {
 
 	b.schedule = &TestSchedule{}
 
-	cleanup, connURL := postgreshelper.PrepareTestContainer(t, "")
+	cleanup, connURL := postgreshelper.PrepareTestContainer(t)
 	defer cleanup()
 
 	// create the database user
@@ -274,7 +274,7 @@ func TestBackend_StaticRole_Rotation_Schedule_ErrorRecover(t *testing.T) {
 
 	b.schedule = &TestSchedule{}
 
-	cleanup, connURL := postgreshelper.PrepareTestContainer(t, "")
+	cleanup, connURL := postgreshelper.PrepareTestContainer(t)
 	t.Cleanup(cleanup)
 
 	// create the database user
@@ -458,7 +458,7 @@ func TestBackend_StaticRole_Rotation_NonStaticError(t *testing.T) {
 	}
 	defer b.Cleanup(context.Background())
 
-	cleanup, connURL := postgreshelper.PrepareTestContainer(t, "")
+	cleanup, connURL := postgreshelper.PrepareTestContainer(t)
 	defer cleanup()
 
 	// create the database user
@@ -562,7 +562,7 @@ func TestBackend_StaticRole_Rotation_Revoke_user(t *testing.T) {
 	}
 	defer b.Cleanup(context.Background())
 
-	cleanup, connURL := postgreshelper.PrepareTestContainer(t, "")
+	cleanup, connURL := postgreshelper.PrepareTestContainer(t)
 	defer cleanup()
 
 	// create the database user
@@ -783,7 +783,7 @@ func TestBackend_StaticRole_Rotation_QueueWAL_discard_role_newer_rotation_date(t
 		t.Fatal("could not convert to db backend")
 	}
 
-	cleanup, connURL := postgreshelper.PrepareTestContainer(t, "")
+	cleanup, connURL := postgreshelper.PrepareTestContainer(t)
 	defer cleanup()
 
 	// create the database user
@@ -952,7 +952,7 @@ func assertWALCount(t *testing.T, s logical.Storage, expected int, key string) {
 type userCreator func(t *testing.T, username, password string)
 
 func TestBackend_StaticRole_Rotation_PostgreSQL(t *testing.T) {
-	cleanup, connURL := postgreshelper.PrepareTestContainer(t, "13.4-buster")
+	cleanup, connURL := postgreshelper.PrepareTestContainer(t)
 	defer cleanup()
 	uc := userCreator(func(t *testing.T, username, password string) {
 		createTestPGUser(t, connURL, username, password, testRoleStaticCreate)
@@ -1246,7 +1246,7 @@ func TestBackend_StaticRole_Rotation_LockRegression(t *testing.T) {
 	}
 	defer b.Cleanup(context.Background())
 
-	cleanup, connURL := postgreshelper.PrepareTestContainer(t, "")
+	cleanup, connURL := postgreshelper.PrepareTestContainer(t)
 	defer cleanup()
 
 	// Configure a connection
@@ -1325,7 +1325,7 @@ func TestBackend_StaticRole_Rotation_Invalid_Role(t *testing.T) {
 	}
 	defer b.Cleanup(context.Background())
 
-	cleanup, connURL := postgreshelper.PrepareTestContainer(t, "")
+	cleanup, connURL := postgreshelper.PrepareTestContainer(t)
 	defer cleanup()
 
 	// create the database user

helper/testhelpers/postgresql/postgresqlhelper.go

@@ -14,13 +14,29 @@ import (
 	"github.com/hashicorp/vault/sdk/helper/docker"
 )
 
-func PrepareTestContainer(t *testing.T, version string) (func(), string) {
-	env := []string{
-		"POSTGRES_PASSWORD=secret",
-		"POSTGRES_DB=database",
-	}
+const postgresVersion = "13.4-buster"
 
-	_, cleanup, url, _ := prepareTestContainer(t, "postgres", "docker.mirror.hashicorp.services/postgres", version, "secret", true, false, false, env)
+func defaultRunOpts(t *testing.T) docker.RunOptions {
+	return docker.RunOptions{
+		ContainerName: "postgres",
+		ImageRepo:     "docker.mirror.hashicorp.services/postgres",
+		ImageTag:      postgresVersion,
+		Env: []string{
+			"POSTGRES_PASSWORD=secret",
+			"POSTGRES_DB=database",
+		},
+		Ports:           []string{"5432/tcp"},
+		DoNotAutoRemove: false,
+		LogConsumer: func(s string) {
+			if t.Failed() {
+				t.Logf("container logs: %s", s)
+			}
+		},
+	}
+}
+
+func PrepareTestContainer(t *testing.T) (func(), string) {
+	_, cleanup, url, _ := prepareTestContainer(t, defaultRunOpts(t), "secret", true, false)
 
 	return cleanup, url
 }
@@ -28,13 +44,8 @@ func PrepareTestContainer(t *testing.T, version string) (func(), string) {
 // PrepareTestContainerWithVaultUser will setup a test container with a Vault
 // admin user configured so that we can safely call rotate-root without
 // rotating the root DB credentials
-func PrepareTestContainerWithVaultUser(t *testing.T, ctx context.Context, version string) (func(), string) {
-	env := []string{
-		"POSTGRES_PASSWORD=secret",
-		"POSTGRES_DB=database",
-	}
-
-	runner, cleanup, url, id := prepareTestContainer(t, "postgres", "docker.mirror.hashicorp.services/postgres", version, "secret", true, false, false, env)
+func PrepareTestContainerWithVaultUser(t *testing.T, ctx context.Context) (func(), string) {
+	runner, cleanup, url, id := prepareTestContainer(t, defaultRunOpts(t), "secret", true, false)
 
 	cmd := []string{"psql", "-U", "postgres", "-c", "CREATE USER vaultadmin WITH LOGIN PASSWORD 'vaultpass' SUPERUSER"}
 	_, err := runner.RunCmdInBackground(ctx, id, cmd)
@@ -45,47 +56,70 @@ func PrepareTestContainerWithVaultUser(t *testing.T, ctx context.Context, versio
 	return cleanup, url
 }
 
+func PrepareTestContainerWithSSL(t *testing.T, ctx context.Context, version string) (func(), string) {
+	runOpts := defaultRunOpts(t)
+	runOpts.Cmd = []string{"-c", "log_statement=all"}
+	runner, cleanup, url, id := prepareTestContainer(t, runOpts, "secret", true, false)
+
+	content := "echo 'hostssl all all all cert clientcert=verify-ca' > /var/lib/postgresql/data/pg_hba.conf"
+	// Copy the ssl init script into the newly running container.
+	buildCtx := docker.NewBuildContext()
+	buildCtx["ssl-conf.sh"] = docker.PathContentsFromBytes([]byte(content))
+	if err := runner.CopyTo(id, "/usr/local/bin", buildCtx); err != nil {
+		t.Fatalf("Could not copy ssl init script into container: %v", err)
+	}
+
+	// run the ssl init script to overwrite the pg_hba.conf file and set it to
+	// require SSL for each connection
+	cmd := []string{"bash", "/usr/local/bin/ssl-conf.sh"}
+	_, err := runner.RunCmdInBackground(ctx, id, cmd)
+	if err != nil {
+		t.Fatalf("Could not run command (%v) in container: %v", cmd, err)
+	}
+
+	// reload so the config changes take effect
+	cmd = []string{"psql", "-U", "postgres", "-c", "SELECT pg_reload_conf()"}
+	_, err = runner.RunCmdInBackground(ctx, id, cmd)
+	if err != nil {
+		t.Fatalf("Could not run command (%v) in container: %v", cmd, err)
+	}
+
+	return cleanup, url
+}
+
 func PrepareTestContainerWithPassword(t *testing.T, version, password string) (func(), string) {
-	env := []string{
+	runOpts := defaultRunOpts(t)
+	runOpts.Env = []string{
 		"POSTGRES_PASSWORD=" + password,
 		"POSTGRES_DB=database",
 	}
 
-	_, cleanup, url, _ := prepareTestContainer(t, "postgres", "docker.mirror.hashicorp.services/postgres", version, password, true, false, false, env)
+	_, cleanup, url, _ := prepareTestContainer(t, runOpts, password, true, false)
 
 	return cleanup, url
 }
 
 func PrepareTestContainerRepmgr(t *testing.T, name, version string, envVars []string) (*docker.Runner, func(), string, string) {
-	env := append(envVars,
+	runOpts := defaultRunOpts(t)
+	runOpts.ImageRepo = "docker.mirror.hashicorp.services/bitnami/postgresql-repmgr"
+	runOpts.ImageTag = version
+	runOpts.Env = append(envVars,
 		"REPMGR_PARTNER_NODES=psql-repl-node-0,psql-repl-node-1",
 		"REPMGR_PRIMARY_HOST=psql-repl-node-0",
 		"REPMGR_PASSWORD=repmgrpass",
 		"POSTGRESQL_PASSWORD=secret")
+	runOpts.DoNotAutoRemove = true
 
-	return prepareTestContainer(t, name, "docker.mirror.hashicorp.services/bitnami/postgresql-repmgr", version, "secret", false, true, true, env)
+	return prepareTestContainer(t, runOpts, "secret", false, true)
 }
 
-func prepareTestContainer(t *testing.T, name, repo, version, password string,
-	addSuffix, forceLocalAddr, doNotAutoRemove bool, envVars []string,
+func prepareTestContainer(t *testing.T, runOpts docker.RunOptions, password string, addSuffix, forceLocalAddr bool,
 ) (*docker.Runner, func(), string, string) {
 	if os.Getenv("PG_URL") != "" {
 		return nil, func() {}, "", os.Getenv("PG_URL")
 	}
 
-	if version == "" {
-		version = "11"
-	}
-
-	runOpts := docker.RunOptions{
-		ContainerName:   name,
-		ImageRepo:       repo,
-		ImageTag:        version,
-		Env:             envVars,
-		Ports:           []string{"5432/tcp"},
-		DoNotAutoRemove: doNotAutoRemove,
-	}
-	if repo == "bitnami/postgresql-repmgr" {
+	if runOpts.ImageRepo == "bitnami/postgresql-repmgr" {
 		runOpts.NetworkID = os.Getenv("POSTGRES_MULTIHOST_NET")
 	}
 
@@ -94,7 +128,7 @@ func prepareTestContainer(t *testing.T, name, repo, version, password string,
 		t.Fatalf("Could not start docker Postgres: %s", err)
 	}
 
-	svc, containerID, err := runner.StartNewService(context.Background(), addSuffix, forceLocalAddr, connectPostgres(password, repo))
+	svc, containerID, err := runner.StartNewService(context.Background(), addSuffix, forceLocalAddr, connectPostgres(password, runOpts.ImageRepo))
 	if err != nil {
 		t.Fatalf("Could not start docker Postgres: %s", err)
 	}

physical/postgresql/postgresql_test.go

@@ -22,7 +22,7 @@ func TestPostgreSQLBackend(t *testing.T) {
 	// Use docker as pg backend if no url is provided via environment variables
 	connURL := os.Getenv("PGURL")
 	if connURL == "" {
-		cleanup, u := postgresql.PrepareTestContainer(t, "11.1")
+		cleanup, u := postgresql.PrepareTestContainer(t)
 		defer cleanup()
 		connURL = u
 	}

plugins/database/postgresql/postgresql_test.go

@@ -24,7 +24,7 @@ import (
 )
 
 func getPostgreSQL(t *testing.T, options map[string]interface{}) (*PostgreSQL, func()) {
-	cleanup, connURL := postgresql.PrepareTestContainer(t, "13.4-buster")
+	cleanup, connURL := postgresql.PrepareTestContainer(t)
 
 	connectionDetails := map[string]interface{}{
 		"connection_url": connURL,
@@ -70,7 +70,7 @@ func TestPostgreSQL_InitializeWithStringVals(t *testing.T) {
 }
 
 func TestPostgreSQL_Initialize_ConnURLWithDSNFormat(t *testing.T) {
-	cleanup, connURL := postgresql.PrepareTestContainer(t, "13.4-buster")
+	cleanup, connURL := postgresql.PrepareTestContainer(t)
 	defer cleanup()
 
 	dsnConnURL, err := dbutil.ParseURL(connURL)
@@ -185,7 +185,7 @@ func TestPostgreSQL_Initialize_CloudGCP(t *testing.T) {
 // TestPostgreSQL_PasswordAuthentication tests that the default "password_authentication" is "none", and that
 // an error is returned if an invalid "password_authentication" is provided.
 func TestPostgreSQL_PasswordAuthentication(t *testing.T) {
-	cleanup, connURL := postgresql.PrepareTestContainer(t, "13.4-buster")
+	cleanup, connURL := postgresql.PrepareTestContainer(t)
 	defer cleanup()
 
 	dsnConnURL, err := dbutil.ParseURL(connURL)
@@ -227,7 +227,7 @@ func TestPostgreSQL_PasswordAuthentication(t *testing.T) {
 // TestPostgreSQL_PasswordAuthentication_SCRAMSHA256 tests that password_authentication works when set to scram-sha-256.
 // When sending an encrypted password, the raw password should still successfully authenticate the user.
 func TestPostgreSQL_PasswordAuthentication_SCRAMSHA256(t *testing.T) {
-	cleanup, connURL := postgresql.PrepareTestContainer(t, "13.4-buster")
+	cleanup, connURL := postgresql.PrepareTestContainer(t)
 	defer cleanup()
 
 	dsnConnURL, err := dbutil.ParseURL(connURL)
@@ -1092,7 +1092,7 @@ func TestUsernameGeneration(t *testing.T) {
 }
 
 func TestNewUser_CustomUsername(t *testing.T) {
-	cleanup, connURL := postgresql.PrepareTestContainer(t, "13.4-buster")
+	cleanup, connURL := postgresql.PrepareTestContainer(t)
 	defer cleanup()
 
 	type testCase struct {

vault/external_tests/plugin/external_plugin_test.go

@@ -696,7 +696,7 @@ func TestExternalPlugin_Database(t *testing.T) {
 			t.Run(dbName, func(t *testing.T) {
 				roleName := "test-role-" + dbName
 
-				cleanupContainer, connURL := postgreshelper.PrepareTestContainerWithVaultUser(t, context.Background(), "13.4-buster")
+				cleanupContainer, connURL := postgreshelper.PrepareTestContainerWithVaultUser(t, context.Background())
 				t.Cleanup(cleanupContainer)
 
 				_, err := client.Logical().Write("database/config/"+dbName, map[string]interface{}{
@@ -819,7 +819,7 @@ func TestExternalPlugin_DatabaseReload(t *testing.T) {
 	dbName := fmt.Sprintf("%s-%d", plugin.Name, 0)
 	roleName := "test-role-" + dbName
 
-	cleanupContainer, connURL := postgreshelper.PrepareTestContainerWithVaultUser(t, context.Background(), "13.4-buster")
+	cleanupContainer, connURL := postgreshelper.PrepareTestContainerWithVaultUser(t, context.Background())
 	t.Cleanup(cleanupContainer)
 
 	_, err := client.Logical().Write("database/config/"+dbName, map[string]interface{}{
@@ -1204,7 +1204,7 @@ func TestCore_UpgradePluginUsingPinnedVersion_Database(t *testing.T) {
 		t.Fatal(err)
 	}
 
-	cleanupPG, connURL := postgreshelper.PrepareTestContainerWithVaultUser(t, context.Background(), "13.4-buster")
+	cleanupPG, connURL := postgreshelper.PrepareTestContainerWithVaultUser(t, context.Background())
 	t.Cleanup(cleanupPG)
 
 	// Mount 1.0.0 then pin to 1.0.1