mirror of
				https://github.com/optim-enterprises-bv/vault.git
				synced 2025-11-01 02:57:59 +00:00 
			
		
		
		
	Move plugins into main vault repo
This commit is contained in:
		
							
								
								
									
										308
									
								
								plugins/database/postgresql/postgresql_test.go
									
									
									
									
									
										Normal file
									
								
							
							
						
						
									
										308
									
								
								plugins/database/postgresql/postgresql_test.go
									
									
									
									
									
										Normal file
									
								
							| @@ -0,0 +1,308 @@ | ||||
| package postgresql | ||||
|  | ||||
| import ( | ||||
| 	"database/sql" | ||||
| 	"fmt" | ||||
| 	"os" | ||||
| 	"strings" | ||||
| 	"sync" | ||||
| 	"testing" | ||||
| 	"time" | ||||
|  | ||||
| 	"github.com/hashicorp/vault/builtin/logical/database/dbplugin" | ||||
| 	"github.com/hashicorp/vault/plugins/helper/database/connutil" | ||||
| 	dockertest "gopkg.in/ory-am/dockertest.v3" | ||||
| ) | ||||
|  | ||||
| var ( | ||||
| 	testPostgresImagePull sync.Once | ||||
| ) | ||||
|  | ||||
| func preparePostgresTestContainer(t *testing.T) (cleanup func(), retURL string) { | ||||
| 	if os.Getenv("PG_URL") != "" { | ||||
| 		return func() {}, os.Getenv("PG_URL") | ||||
| 	} | ||||
|  | ||||
| 	pool, err := dockertest.NewPool("") | ||||
| 	if err != nil { | ||||
| 		t.Fatalf("Failed to connect to docker: %s", err) | ||||
| 	} | ||||
|  | ||||
| 	resource, err := pool.Run("postgres", "latest", []string{"POSTGRES_PASSWORD=secret", "POSTGRES_DB=database"}) | ||||
| 	if err != nil { | ||||
| 		t.Fatalf("Could not start local PostgreSQL docker container: %s", err) | ||||
| 	} | ||||
|  | ||||
| 	cleanup = func() { | ||||
| 		err := pool.Purge(resource) | ||||
| 		if err != nil { | ||||
| 			t.Fatalf("Failed to cleanup local container: %s", err) | ||||
| 		} | ||||
| 	} | ||||
|  | ||||
| 	retURL = fmt.Sprintf("postgres://postgres:secret@localhost:%s/database?sslmode=disable", resource.GetPort("5432/tcp")) | ||||
|  | ||||
| 	// exponential backoff-retry | ||||
| 	if err = pool.Retry(func() error { | ||||
| 		var err error | ||||
| 		var db *sql.DB | ||||
| 		db, err = sql.Open("postgres", retURL) | ||||
| 		if err != nil { | ||||
| 			return err | ||||
| 		} | ||||
| 		return db.Ping() | ||||
| 	}); err != nil { | ||||
| 		t.Fatalf("Could not connect to PostgreSQL docker container: %s", err) | ||||
| 	} | ||||
|  | ||||
| 	return | ||||
| } | ||||
|  | ||||
| func TestPostgreSQL_Initialize(t *testing.T) { | ||||
| 	cleanup, connURL := preparePostgresTestContainer(t) | ||||
| 	defer cleanup() | ||||
|  | ||||
| 	connectionDetails := map[string]interface{}{ | ||||
| 		"connection_url": connURL, | ||||
| 	} | ||||
|  | ||||
| 	db := New() | ||||
| 	connProducer := db.ConnectionProducer.(*connutil.SQLConnectionProducer) | ||||
|  | ||||
| 	err := db.Initialize(connectionDetails, true) | ||||
| 	if err != nil { | ||||
| 		t.Fatalf("err: %s", err) | ||||
| 	} | ||||
|  | ||||
| 	if !connProducer.Initialized { | ||||
| 		t.Fatal("Database should be initalized") | ||||
| 	} | ||||
|  | ||||
| 	err = db.Close() | ||||
| 	if err != nil { | ||||
| 		t.Fatalf("err: %s", err) | ||||
| 	} | ||||
| } | ||||
|  | ||||
| func TestPostgreSQL_CreateUser(t *testing.T) { | ||||
| 	cleanup, connURL := preparePostgresTestContainer(t) | ||||
| 	defer cleanup() | ||||
|  | ||||
| 	connectionDetails := map[string]interface{}{ | ||||
| 		"connection_url": connURL, | ||||
| 	} | ||||
|  | ||||
| 	db := New() | ||||
| 	err := db.Initialize(connectionDetails, true) | ||||
| 	if err != nil { | ||||
| 		t.Fatalf("err: %s", err) | ||||
| 	} | ||||
|  | ||||
| 	// Test with no configured Creation Statememt | ||||
| 	_, _, err = db.CreateUser(dbplugin.Statements{}, "test", time.Now().Add(time.Minute)) | ||||
| 	if err == nil { | ||||
| 		t.Fatal("Expected error when no creation statement is provided") | ||||
| 	} | ||||
|  | ||||
| 	statements := dbplugin.Statements{ | ||||
| 		CreationStatements: testPostgresRole, | ||||
| 	} | ||||
|  | ||||
| 	username, password, err := db.CreateUser(statements, "test", time.Now().Add(time.Minute)) | ||||
| 	if err != nil { | ||||
| 		t.Fatalf("err: %s", err) | ||||
| 	} | ||||
|  | ||||
| 	if err = testCredsExist(t, connURL, username, password); err != nil { | ||||
| 		t.Fatalf("Could not connect with new credentials: %s", err) | ||||
| 	} | ||||
|  | ||||
| 	statements.CreationStatements = testPostgresReadOnlyRole | ||||
| 	username, password, err = db.CreateUser(statements, "test", time.Now().Add(time.Minute)) | ||||
| 	if err != nil { | ||||
| 		t.Fatalf("err: %s", err) | ||||
| 	} | ||||
|  | ||||
| 	if err = testCredsExist(t, connURL, username, password); err != nil { | ||||
| 		t.Fatalf("Could not connect with new credentials: %s", err) | ||||
| 	} | ||||
| } | ||||
|  | ||||
| func TestPostgreSQL_RenewUser(t *testing.T) { | ||||
| 	cleanup, connURL := preparePostgresTestContainer(t) | ||||
| 	defer cleanup() | ||||
|  | ||||
| 	connectionDetails := map[string]interface{}{ | ||||
| 		"connection_url": connURL, | ||||
| 	} | ||||
|  | ||||
| 	db := New() | ||||
| 	err := db.Initialize(connectionDetails, true) | ||||
| 	if err != nil { | ||||
| 		t.Fatalf("err: %s", err) | ||||
| 	} | ||||
|  | ||||
| 	statements := dbplugin.Statements{ | ||||
| 		CreationStatements: testPostgresRole, | ||||
| 	} | ||||
|  | ||||
| 	username, password, err := db.CreateUser(statements, "test", time.Now().Add(2*time.Second)) | ||||
| 	if err != nil { | ||||
| 		t.Fatalf("err: %s", err) | ||||
| 	} | ||||
|  | ||||
| 	if err = testCredsExist(t, connURL, username, password); err != nil { | ||||
| 		t.Fatalf("Could not connect with new credentials: %s", err) | ||||
| 	} | ||||
|  | ||||
| 	err = db.RenewUser(statements, username, time.Now().Add(time.Minute)) | ||||
| 	if err != nil { | ||||
| 		t.Fatalf("err: %s", err) | ||||
| 	} | ||||
|  | ||||
| 	// Sleep longer than the inital expiration time | ||||
| 	time.Sleep(2 * time.Second) | ||||
|  | ||||
| 	if err = testCredsExist(t, connURL, username, password); err != nil { | ||||
| 		t.Fatalf("Could not connect with new credentials: %s", err) | ||||
| 	} | ||||
| } | ||||
|  | ||||
| func TestPostgreSQL_RevokeUser(t *testing.T) { | ||||
| 	cleanup, connURL := preparePostgresTestContainer(t) | ||||
| 	defer cleanup() | ||||
|  | ||||
| 	connectionDetails := map[string]interface{}{ | ||||
| 		"connection_url": connURL, | ||||
| 	} | ||||
|  | ||||
| 	db := New() | ||||
| 	err := db.Initialize(connectionDetails, true) | ||||
| 	if err != nil { | ||||
| 		t.Fatalf("err: %s", err) | ||||
| 	} | ||||
|  | ||||
| 	statements := dbplugin.Statements{ | ||||
| 		CreationStatements: testPostgresRole, | ||||
| 	} | ||||
|  | ||||
| 	username, password, err := db.CreateUser(statements, "test", time.Now().Add(2*time.Second)) | ||||
| 	if err != nil { | ||||
| 		t.Fatalf("err: %s", err) | ||||
| 	} | ||||
|  | ||||
| 	if err = testCredsExist(t, connURL, username, password); err != nil { | ||||
| 		t.Fatalf("Could not connect with new credentials: %s", err) | ||||
| 	} | ||||
|  | ||||
| 	// Test default revoke statememts | ||||
| 	err = db.RevokeUser(statements, username) | ||||
| 	if err != nil { | ||||
| 		t.Fatalf("err: %s", err) | ||||
| 	} | ||||
|  | ||||
| 	if err := testCredsExist(t, connURL, username, password); err == nil { | ||||
| 		t.Fatal("Credentials were not revoked") | ||||
| 	} | ||||
|  | ||||
| 	username, password, err = db.CreateUser(statements, "test", time.Now().Add(2*time.Second)) | ||||
| 	if err != nil { | ||||
| 		t.Fatalf("err: %s", err) | ||||
| 	} | ||||
|  | ||||
| 	if err = testCredsExist(t, connURL, username, password); err != nil { | ||||
| 		t.Fatalf("Could not connect with new credentials: %s", err) | ||||
| 	} | ||||
|  | ||||
| 	// Test custom revoke statements | ||||
| 	statements.RevocationStatements = defaultPostgresRevocationSQL | ||||
| 	err = db.RevokeUser(statements, username) | ||||
| 	if err != nil { | ||||
| 		t.Fatalf("err: %s", err) | ||||
| 	} | ||||
|  | ||||
| 	if err := testCredsExist(t, connURL, username, password); err == nil { | ||||
| 		t.Fatal("Credentials were not revoked") | ||||
| 	} | ||||
| } | ||||
|  | ||||
| func testCredsExist(t testing.TB, connURL, username, password string) error { | ||||
| 	// Log in with the new creds | ||||
| 	connURL = strings.Replace(connURL, "postgres:secret", fmt.Sprintf("%s:%s", username, password), 1) | ||||
| 	db, err := sql.Open("postgres", connURL) | ||||
| 	if err != nil { | ||||
| 		return err | ||||
| 	} | ||||
| 	defer db.Close() | ||||
| 	return db.Ping() | ||||
| } | ||||
|  | ||||
| const testPostgresRole = ` | ||||
| CREATE ROLE "{{name}}" WITH | ||||
|   LOGIN | ||||
|   PASSWORD '{{password}}' | ||||
|   VALID UNTIL '{{expiration}}'; | ||||
| GRANT ALL PRIVILEGES ON ALL TABLES IN SCHEMA public TO "{{name}}"; | ||||
| ` | ||||
|  | ||||
| const testPostgresReadOnlyRole = ` | ||||
| CREATE ROLE "{{name}}" WITH | ||||
|   LOGIN | ||||
|   PASSWORD '{{password}}' | ||||
|   VALID UNTIL '{{expiration}}'; | ||||
| GRANT SELECT ON ALL TABLES IN SCHEMA public TO "{{name}}"; | ||||
| GRANT SELECT ON ALL SEQUENCES IN SCHEMA public TO "{{name}}"; | ||||
| ` | ||||
|  | ||||
| const testPostgresBlockStatementRole = ` | ||||
| DO $$ | ||||
| BEGIN | ||||
|    IF NOT EXISTS (SELECT * FROM pg_catalog.pg_roles WHERE rolname='foo-role') THEN | ||||
|       CREATE ROLE "foo-role"; | ||||
|       CREATE SCHEMA IF NOT EXISTS foo AUTHORIZATION "foo-role"; | ||||
|       ALTER ROLE "foo-role" SET search_path = foo; | ||||
|       GRANT TEMPORARY ON DATABASE "postgres" TO "foo-role"; | ||||
|       GRANT ALL PRIVILEGES ON SCHEMA foo TO "foo-role"; | ||||
|       GRANT ALL PRIVILEGES ON ALL TABLES IN SCHEMA foo TO "foo-role"; | ||||
|       GRANT ALL PRIVILEGES ON ALL SEQUENCES IN SCHEMA foo TO "foo-role"; | ||||
|       GRANT ALL PRIVILEGES ON ALL FUNCTIONS IN SCHEMA foo TO "foo-role"; | ||||
|    END IF; | ||||
| END | ||||
| $$ | ||||
|  | ||||
| CREATE ROLE "{{name}}" WITH LOGIN PASSWORD '{{password}}' VALID UNTIL '{{expiration}}'; | ||||
| GRANT "foo-role" TO "{{name}}"; | ||||
| ALTER ROLE "{{name}}" SET search_path = foo; | ||||
| GRANT CONNECT ON DATABASE "postgres" TO "{{name}}"; | ||||
| ` | ||||
|  | ||||
| var testPostgresBlockStatementRoleSlice = []string{ | ||||
| 	` | ||||
| DO $$ | ||||
| BEGIN | ||||
|    IF NOT EXISTS (SELECT * FROM pg_catalog.pg_roles WHERE rolname='foo-role') THEN | ||||
|       CREATE ROLE "foo-role"; | ||||
|       CREATE SCHEMA IF NOT EXISTS foo AUTHORIZATION "foo-role"; | ||||
|       ALTER ROLE "foo-role" SET search_path = foo; | ||||
|       GRANT TEMPORARY ON DATABASE "postgres" TO "foo-role"; | ||||
|       GRANT ALL PRIVILEGES ON SCHEMA foo TO "foo-role"; | ||||
|       GRANT ALL PRIVILEGES ON ALL TABLES IN SCHEMA foo TO "foo-role"; | ||||
|       GRANT ALL PRIVILEGES ON ALL SEQUENCES IN SCHEMA foo TO "foo-role"; | ||||
|       GRANT ALL PRIVILEGES ON ALL FUNCTIONS IN SCHEMA foo TO "foo-role"; | ||||
|    END IF; | ||||
| END | ||||
| $$ | ||||
| `, | ||||
| 	`CREATE ROLE "{{name}}" WITH LOGIN PASSWORD '{{password}}' VALID UNTIL '{{expiration}}';`, | ||||
| 	`GRANT "foo-role" TO "{{name}}";`, | ||||
| 	`ALTER ROLE "{{name}}" SET search_path = foo;`, | ||||
| 	`GRANT CONNECT ON DATABASE "postgres" TO "{{name}}";`, | ||||
| } | ||||
|  | ||||
| const defaultPostgresRevocationSQL = ` | ||||
| REVOKE ALL PRIVILEGES ON ALL TABLES IN SCHEMA public FROM "{{name}}"; | ||||
| REVOKE ALL PRIVILEGES ON ALL SEQUENCES IN SCHEMA public FROM "{{name}}"; | ||||
| REVOKE USAGE ON SCHEMA public FROM "{{name}}"; | ||||
|  | ||||
| DROP ROLE IF EXISTS "{{name}}"; | ||||
| ` | ||||
		Reference in New Issue
	
	Block a user
	 Brian Kassouf
					Brian Kassouf