Postgres revocation sql, beta mode (#1972)

builtin/logical/postgresql/backend_test.go

@@ -233,6 +233,39 @@ func TestBackend_roleReadOnly(t *testing.T) {
 	})
 }
 
+func TestBackend_roleReadOnly_revocationSQL(t *testing.T) {
+	config := logical.TestBackendConfig()
+	config.StorageView = &logical.InmemStorage{}
+	b, err := Factory(config)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	cid, connURL := prepareTestContainer(t, config.StorageView, b)
+	if cid != "" {
+		defer cleanupTestContainer(t, cid)
+	}
+	connData := map[string]interface{}{
+		"connection_url": connURL,
+	}
+
+	logicaltest.Test(t, logicaltest.TestCase{
+		Backend: b,
+		Steps: []logicaltest.TestStep{
+			testAccStepConfig(t, connData, false),
+			testAccStepCreateRoleWithRevocationSQL(t, "web", testRole, defaultRevocationSQL, false),
+			testAccStepCreateRoleWithRevocationSQL(t, "web-readonly", testReadOnlyRole, defaultRevocationSQL, false),
+			testAccStepReadRole(t, "web-readonly", testReadOnlyRole),
+			testAccStepCreateTable(t, b, config.StorageView, "web", connURL),
+			testAccStepReadCreds(t, b, config.StorageView, "web-readonly", connURL),
+			testAccStepDropTable(t, b, config.StorageView, "web", connURL),
+			testAccStepDeleteRole(t, "web-readonly"),
+			testAccStepDeleteRole(t, "web"),
+			testAccStepReadRole(t, "web-readonly", ""),
+		},
+	})
+}
+
 func testAccStepConfig(t *testing.T, d map[string]interface{}, expectError bool) logicaltest.TestStep {
 	return logicaltest.TestStep{
 		Operation: logical.UpdateOperation,
@@ -273,6 +306,18 @@ func testAccStepCreateRole(t *testing.T, name string, sql string, expectFail boo
 	}
 }
 
+func testAccStepCreateRoleWithRevocationSQL(t *testing.T, name, sql, revocationSQL string, expectFail bool) logicaltest.TestStep {
+	return logicaltest.TestStep{
+		Operation: logical.UpdateOperation,
+		Path:      path.Join("roles", name),
+		Data: map[string]interface{}{
+			"sql":            sql,
+			"revocation_sql": revocationSQL,
+		},
+		ErrorOk: expectFail,
+	}
+}
+
 func testAccStepDeleteRole(t *testing.T, name string) logicaltest.TestStep {
 	return logicaltest.TestStep{
 		Operation: logical.DeleteOperation,
@@ -341,6 +386,7 @@ func testAccStepReadCreds(t *testing.T, b logical.Backend, s logical.Storage, na
 					InternalData: map[string]interface{}{
 						"secret_type": "creds",
 						"username":    d.Username,
+						"role":        name,
 					},
 				},
 			})
@@ -543,7 +589,8 @@ GRANT CONNECT ON DATABASE "postgres" TO "{{name}}";
 `
 
 var testBlockStatementRoleSlice = []string{
-	`DO $$
+	`
+DO $$
 BEGIN
    IF NOT EXISTS (SELECT * FROM pg_catalog.pg_roles WHERE rolname='foo-role') THEN
       CREATE ROLE "foo-role";
@@ -556,9 +603,18 @@ BEGIN
       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 defaultRevocationSQL = `
+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}};
+`

builtin/logical/postgresql/path_role_create.go

@@ -140,6 +140,7 @@ func (b *backend) pathRoleCreateRead(
 		"password": password,
 	}, map[string]interface{}{
 		"username": username,
+		"role":     name,
 	})
 	resp.Secret.TTL = lease.Lease
 	return resp, nil

builtin/logical/postgresql/path_roles.go

@@ -26,15 +26,20 @@ func pathRoles(b *backend) *framework.Path {
 	return &framework.Path{
 		Pattern: "roles/" + framework.GenericNameRegex("name"),
 		Fields: map[string]*framework.FieldSchema{
-			"name": &framework.FieldSchema{
+			"name": {
 				Type:        framework.TypeString,
 				Description: "Name of the role.",
 			},
 
-			"sql": &framework.FieldSchema{
+			"sql": {
 				Type:        framework.TypeString,
 				Description: "SQL string to create a user. See help for more info.",
 			},
+
+			"revocation_sql": {
+				Type:        framework.TypeString,
+				Description: "SQL string to revoke a user. This is in beta; use with caution.",
+			},
 		},
 
 		Callbacks: map[logical.Operation]framework.OperationFunc{
@@ -85,11 +90,19 @@ func (b *backend) pathRoleRead(
 		return nil, nil
 	}
 
-	return &logical.Response{
+	resp := &logical.Response{
 		Data: map[string]interface{}{
 			"sql": role.SQL,
 		},
-	}, nil
+	}
+
+	// TODO: This is separate because this is in beta in 0.6.2, so we don't
+	// want it to show up in the general case.
+	if role.RevocationSQL != "" {
+		resp.Data["revocation_sql"] = role.RevocationSQL
+	}
+
+	return resp, nil
 }
 
 func (b *backend) pathRoleList(
@@ -134,7 +147,8 @@ func (b *backend) pathRoleCreate(
 
 	// Store it
 	entry, err := logical.StorageEntryJSON("role/"+name, &roleEntry{
-		SQL: sql,
+		SQL:           sql,
+		RevocationSQL: data.Get("revocation_sql").(string),
 	})
 	if err != nil {
 		return nil, err
@@ -147,7 +161,8 @@ func (b *backend) pathRoleCreate(
 }
 
 type roleEntry struct {
-	SQL string `json:"sql"`
+	SQL           string `json:"sql" mapstructure:"sql" structs:"sql"`
+	RevocationSQL string `json:"revocation_sql" mapstructure:"revocation_sql" structs:"revocation_sql"`
 }
 
 const pathRoleHelpSyn = `

builtin/logical/postgresql/secret_creds.go

@@ -3,7 +3,9 @@ package postgresql
 import (
 	"database/sql"
 	"fmt"
+	"strings"
 
+	"github.com/hashicorp/vault/helper/strutil"
 	"github.com/hashicorp/vault/logical"
 	"github.com/hashicorp/vault/logical/framework"
 	"github.com/lib/pq"
@@ -91,120 +93,178 @@ func (b *backend) secretCredsRevoke(
 	}
 	username, ok := usernameRaw.(string)
 
+	var revocationSQL string
+	var resp *logical.Response
+
+	roleNameRaw, ok := req.Secret.InternalData["role"]
+	if ok {
+		role, err := b.Role(req.Storage, roleNameRaw.(string))
+		if err != nil {
+			return nil, err
+		}
+		if role == nil {
+			if resp == nil {
+				resp = &logical.Response{}
+			}
+			resp.AddWarning(fmt.Sprintf("Role %q cannot be found. Using default revocation SQL.", roleNameRaw.(string)))
+		} else {
+			revocationSQL = role.RevocationSQL
+		}
+	}
+
 	// Get our connection
 	db, err := b.DB(req.Storage)
 	if err != nil {
 		return nil, err
 	}
 
-	// Check if the role exists
-	var exists bool
-	err = db.QueryRow("SELECT exists (SELECT rolname FROM pg_roles WHERE rolname=$1);", username).Scan(&exists)
-	if err != nil && err != sql.ErrNoRows {
-		return nil, err
-	}
+	switch revocationSQL {
 
-	if exists == false {
-		return nil, nil
-	}
-
-	// Query for permissions; we need to revoke permissions before we can drop
-	// the role
-	// This isn't done in a transaction because even if we fail along the way,
-	// we want to remove as much access as possible
-	stmt, err := db.Prepare("SELECT DISTINCT table_schema FROM information_schema.role_column_grants WHERE grantee=$1;")
-	if err != nil {
-		return nil, err
-	}
-	defer stmt.Close()
-
-	rows, err := stmt.Query(username)
-	if err != nil {
-		return nil, err
-	}
-	defer rows.Close()
-
-	const initialNumRevocations = 16
-	revocationStmts := make([]string, 0, initialNumRevocations)
-	for rows.Next() {
-		var schema string
-		err = rows.Scan(&schema)
-		if err != nil {
-			// keep going; remove as many permissions as possible right now
-			continue
+	// This is the default revocation logic. If revocation SQL is provided it
+	// is simply executed as-is.
+	case "":
+		// Check if the role exists
+		var exists bool
+		err = db.QueryRow("SELECT exists (SELECT rolname FROM pg_roles WHERE rolname=$1);", username).Scan(&exists)
+		if err != nil && err != sql.ErrNoRows {
+			return nil, err
 		}
-		revocationStmts = append(revocationStmts, fmt.Sprintf(
-			`REVOKE ALL PRIVILEGES ON ALL TABLES IN SCHEMA %s FROM %s;`,
-			pq.QuoteIdentifier(schema),
-			pq.QuoteIdentifier(username)))
 
-		revocationStmts = append(revocationStmts, fmt.Sprintf(
-			`REVOKE USAGE ON SCHEMA %s FROM %s;`,
-			pq.QuoteIdentifier(schema),
-			pq.QuoteIdentifier(username)))
-	}
+		if exists == false {
+			return resp, nil
+		}
 
-	// for good measure, revoke all privileges and usage on schema public
-	revocationStmts = append(revocationStmts, fmt.Sprintf(
-		`REVOKE ALL PRIVILEGES ON ALL TABLES IN SCHEMA public FROM %s;`,
-		pq.QuoteIdentifier(username)))
-
-	revocationStmts = append(revocationStmts, fmt.Sprintf(
-		"REVOKE ALL PRIVILEGES ON ALL SEQUENCES IN SCHEMA public FROM %s;",
-		pq.QuoteIdentifier(username)))
-
-	revocationStmts = append(revocationStmts, fmt.Sprintf(
-		"REVOKE USAGE ON SCHEMA public FROM %s;",
-		pq.QuoteIdentifier(username)))
-
-	// get the current database name so we can issue a REVOKE CONNECT for
-	// this username
-	var dbname sql.NullString
-	if err := db.QueryRow("SELECT current_database();").Scan(&dbname); err != nil {
-		return nil, err
-	}
-
-	if dbname.Valid {
-		revocationStmts = append(revocationStmts, fmt.Sprintf(
-			`REVOKE CONNECT ON DATABASE %s FROM %s;`,
-			pq.QuoteIdentifier(dbname.String),
-			pq.QuoteIdentifier(username)))
-	}
-
-	// again, here, we do not stop on error, as we want to remove as
-	// many permissions as possible right now
-	var lastStmtError error
-	for _, query := range revocationStmts {
-		stmt, err := db.Prepare(query)
+		// Query for permissions; we need to revoke permissions before we can drop
+		// the role
+		// This isn't done in a transaction because even if we fail along the way,
+		// we want to remove as much access as possible
+		stmt, err := db.Prepare("SELECT DISTINCT table_schema FROM information_schema.role_column_grants WHERE grantee=$1;")
 		if err != nil {
-			lastStmtError = err
-			continue
+			return nil, err
 		}
 		defer stmt.Close()
-		_, err = stmt.Exec()
+
+		rows, err := stmt.Query(username)
 		if err != nil {
-			lastStmtError = err
+			return nil, err
+		}
+		defer rows.Close()
+
+		const initialNumRevocations = 16
+		revocationStmts := make([]string, 0, initialNumRevocations)
+		for rows.Next() {
+			var schema string
+			err = rows.Scan(&schema)
+			if err != nil {
+				// keep going; remove as many permissions as possible right now
+				continue
+			}
+			revocationStmts = append(revocationStmts, fmt.Sprintf(
+				`REVOKE ALL PRIVILEGES ON ALL TABLES IN SCHEMA %s FROM %s;`,
+				pq.QuoteIdentifier(schema),
+				pq.QuoteIdentifier(username)))
+
+			revocationStmts = append(revocationStmts, fmt.Sprintf(
+				`REVOKE USAGE ON SCHEMA %s FROM %s;`,
+				pq.QuoteIdentifier(schema),
+				pq.QuoteIdentifier(username)))
+		}
+
+		// for good measure, revoke all privileges and usage on schema public
+		revocationStmts = append(revocationStmts, fmt.Sprintf(
+			`REVOKE ALL PRIVILEGES ON ALL TABLES IN SCHEMA public FROM %s;`,
+			pq.QuoteIdentifier(username)))
+
+		revocationStmts = append(revocationStmts, fmt.Sprintf(
+			"REVOKE ALL PRIVILEGES ON ALL SEQUENCES IN SCHEMA public FROM %s;",
+			pq.QuoteIdentifier(username)))
+
+		revocationStmts = append(revocationStmts, fmt.Sprintf(
+			"REVOKE USAGE ON SCHEMA public FROM %s;",
+			pq.QuoteIdentifier(username)))
+
+		// get the current database name so we can issue a REVOKE CONNECT for
+		// this username
+		var dbname sql.NullString
+		if err := db.QueryRow("SELECT current_database();").Scan(&dbname); err != nil {
+			return nil, err
+		}
+
+		if dbname.Valid {
+			revocationStmts = append(revocationStmts, fmt.Sprintf(
+				`REVOKE CONNECT ON DATABASE %s FROM %s;`,
+				pq.QuoteIdentifier(dbname.String),
+				pq.QuoteIdentifier(username)))
+		}
+
+		// again, here, we do not stop on error, as we want to remove as
+		// many permissions as possible right now
+		var lastStmtError error
+		for _, query := range revocationStmts {
+			stmt, err := db.Prepare(query)
+			if err != nil {
+				lastStmtError = err
+				continue
+			}
+			defer stmt.Close()
+			_, err = stmt.Exec()
+			if err != nil {
+				lastStmtError = err
+			}
+		}
+
+		// can't drop if not all privileges are revoked
+		if rows.Err() != nil {
+			return nil, fmt.Errorf("could not generate revocation statements for all rows: %s", rows.Err())
+		}
+		if lastStmtError != nil {
+			return nil, fmt.Errorf("could not perform all revocation statements: %s", lastStmtError)
+		}
+
+		// Drop this user
+		stmt, err = db.Prepare(fmt.Sprintf(
+			`DROP ROLE IF EXISTS %s;`, pq.QuoteIdentifier(username)))
+		if err != nil {
+			return nil, err
+		}
+		defer stmt.Close()
+		if _, err := stmt.Exec(); err != nil {
+			return nil, err
+		}
+
+	// We have revocation SQL, execute directly, within a transaction
+	default:
+		tx, err := db.Begin()
+		if err != nil {
+			return nil, err
+		}
+		defer func() {
+			tx.Rollback()
+		}()
+
+		for _, query := range strutil.ParseArbitraryStringSlice(revocationSQL, ";") {
+			query = strings.TrimSpace(query)
+			if len(query) == 0 {
+				continue
+			}
+
+			stmt, err := tx.Prepare(Query(query, map[string]string{
+				"name": pq.QuoteIdentifier(username),
+			}))
+			if err != nil {
+				return nil, err
+			}
+			defer stmt.Close()
+
+			if _, err := stmt.Exec(); err != nil {
+				return nil, err
+			}
+		}
+
+		if err := tx.Commit(); err != nil {
+			return nil, err
 		}
 	}
 
-	// can't drop if not all privileges are revoked
-	if rows.Err() != nil {
-		return nil, fmt.Errorf("could not generate revocation statements for all rows: %s", rows.Err())
-	}
-	if lastStmtError != nil {
-		return nil, fmt.Errorf("could not perform all revocation statements: %s", lastStmtError)
-	}
-
-	// Drop this user
-	stmt, err = db.Prepare(fmt.Sprintf(
-		`DROP ROLE IF EXISTS %s;`, pq.QuoteIdentifier(username)))
-	if err != nil {
-		return nil, err
-	}
-	defer stmt.Close()
-	if _, err := stmt.Exec(); err != nil {
-		return nil, err
-	}
-
-	return nil, nil
+	return resp, nil
 }