Skip to content

Commit

Permalink
PostgreSQL - Add username customization (#10766)
Browse files Browse the repository at this point in the history
  • Loading branch information
@pcman312
pcman312 authored Feb 4, 2021
1 parent 2650dcd commit cf85a86
Show file tree
Hide file tree
Showing 6 changed files with 424 additions and 43 deletions.
78 changes: 58 additions & 20 deletions go.sum
View file

Large diffs are not rendered by default.

33 changes: 26 additions & 7 deletions plugins/database/postgresql/postgresql.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -11,10 +11,10 @@ import (
"github.com/hashicorp/go-multierror"
dbplugin "github.com/hashicorp/vault/sdk/database/dbplugin/v5"
"github.com/hashicorp/vault/sdk/database/helper/connutil"
"github.com/hashicorp/vault/sdk/database/helper/credsutil"
"github.com/hashicorp/vault/sdk/database/helper/dbutil"
"github.com/hashicorp/vault/sdk/helper/dbtxn"
"github.com/hashicorp/vault/sdk/helper/strutil"
"github.com/hashicorp/vault/sdk/helper/template"
"github.com/lib/pq"
)

Expand All @@ -28,6 +28,8 @@ ALTER ROLE "{{username}}" WITH PASSWORD '{{password}}';
`

expirationFormat = "2006-01-02 15:04:05-0700"

defaultUserNameTemplate = `{{ printf "v-%s-%s-%s-%s" (.DisplayName | truncate 8) (.RoleName | truncate 8) (random 20) (unix_time) | truncate 63 }}`
)

var (
Expand Down Expand Up @@ -68,13 +70,35 @@ func new() *PostgreSQL {

type PostgreSQL struct {
*connutil.SQLConnectionProducer

usernameProducer template.StringTemplate
}

func (p *PostgreSQL) Initialize(ctx context.Context, req dbplugin.InitializeRequest) (dbplugin.InitializeResponse, error) {
newConf, err := p.SQLConnectionProducer.Init(ctx, req.Config, req.VerifyConnection)
if err != nil {
return dbplugin.InitializeResponse{}, err
}

usernameTemplate, err := strutil.GetString(req.Config, "username_template")
if err != nil {
return dbplugin.InitializeResponse{}, fmt.Errorf("failed to retrieve username_template: %w", err)
}
if usernameTemplate == "" {
usernameTemplate = defaultUserNameTemplate
}

up, err := template.NewTemplate(template.Template(usernameTemplate))
if err != nil {
return dbplugin.InitializeResponse{}, fmt.Errorf("unable to initialize username template: %w", err)
}
p.usernameProducer = up

_, err = p.usernameProducer.Generate(dbplugin.UsernameMetadata{})
if err != nil {
return dbplugin.InitializeResponse{}, fmt.Errorf("invalid username template: %w", err)
}

resp := dbplugin.InitializeResponse{
Config: newConf,
}
Expand Down Expand Up @@ -224,12 +248,7 @@ func (p *PostgreSQL) NewUser(ctx context.Context, req dbplugin.NewUserRequest) (
p.Lock()
defer p.Unlock()

username, err := credsutil.GenerateUsername(
credsutil.DisplayName(req.UsernameConfig.DisplayName, 8),
credsutil.RoleName(req.UsernameConfig.RoleName, 8),
credsutil.Separator("-"),
credsutil.MaxLength(63),
)
username, err := p.usernameProducer.Generate(req.UsernameConfig)
if err != nil {
return dbplugin.NewUserResponse{}, err
}
Expand Down
244 changes: 228 additions & 16 deletions plugins/database/postgresql/postgresql_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,8 @@ import (
"github.com/hashicorp/vault/helper/testhelpers/postgresql"
dbplugin "github.com/hashicorp/vault/sdk/database/dbplugin/v5"
dbtesting "github.com/hashicorp/vault/sdk/database/dbplugin/v5/testing"
"github.com/hashicorp/vault/sdk/helper/template"
"github.com/stretchr/testify/require"
)

func getPostgreSQL(t *testing.T, options map[string]interface{}) (*PostgreSQL, func()) {
Expand Down Expand Up @@ -77,8 +79,11 @@ func TestPostgreSQL_NewUser(t *testing.T) {
Password: "somesecurepassword",
Expiration: time.Now().Add(1 * time.Minute),
},
expectErr: true,
credsAssertion: assertCredsDoNotExist,
expectErr: true,
credsAssertion: assertCreds(
assertUsernameRegex("^$"),
assertCredsDoNotExist,
),
},
"admin name": {
req: dbplugin.NewUserRequest{
Expand All @@ -98,8 +103,11 @@ func TestPostgreSQL_NewUser(t *testing.T) {
Password: "somesecurepassword",
Expiration: time.Now().Add(1 * time.Minute),
},
expectErr: false,
credsAssertion: assertCredsExist,
expectErr: false,
credsAssertion: assertCreds(
assertUsernameRegex("^v-test-test-[a-zA-Z0-9]{20}-[0-9]{10}$"),
assertCredsExist,
),
},
"admin username": {
req: dbplugin.NewUserRequest{
Expand All @@ -119,8 +127,11 @@ func TestPostgreSQL_NewUser(t *testing.T) {
Password: "somesecurepassword",
Expiration: time.Now().Add(1 * time.Minute),
},
expectErr: false,
credsAssertion: assertCredsExist,
expectErr: false,
credsAssertion: assertCreds(
assertUsernameRegex("^v-test-test-[a-zA-Z0-9]{20}-[0-9]{10}$"),
assertCredsExist,
),
},
"read only name": {
req: dbplugin.NewUserRequest{
Expand All @@ -141,8 +152,11 @@ func TestPostgreSQL_NewUser(t *testing.T) {
Password: "somesecurepassword",
Expiration: time.Now().Add(1 * time.Minute),
},
expectErr: false,
credsAssertion: assertCredsExist,
expectErr: false,
credsAssertion: assertCreds(
assertUsernameRegex("^v-test-test-[a-zA-Z0-9]{20}-[0-9]{10}$"),
assertCredsExist,
),
},
"read only username": {
req: dbplugin.NewUserRequest{
Expand All @@ -163,8 +177,11 @@ func TestPostgreSQL_NewUser(t *testing.T) {
Password: "somesecurepassword",
Expiration: time.Now().Add(1 * time.Minute),
},
expectErr: false,
credsAssertion: assertCredsExist,
expectErr: false,
credsAssertion: assertCreds(
assertUsernameRegex("^v-test-test-[a-zA-Z0-9]{20}-[0-9]{10}$"),
assertCredsExist,
),
},
// https://github.com/hashicorp/vault/issues/6098
"reproduce GH-6098": {
Expand All @@ -182,8 +199,11 @@ func TestPostgreSQL_NewUser(t *testing.T) {
Password: "somesecurepassword",
Expiration: time.Now().Add(1 * time.Minute),
},
expectErr: false,
credsAssertion: assertCredsDoNotExist,
expectErr: false,
credsAssertion: assertCreds(
assertUsernameRegex("^v-test-test-[a-zA-Z0-9]{20}-[0-9]{10}$"),
assertCredsDoNotExist,
),
},
"reproduce issue with template": {
req: dbplugin.NewUserRequest{
Expand All @@ -199,8 +219,11 @@ func TestPostgreSQL_NewUser(t *testing.T) {
Password: "somesecurepassword",
Expiration: time.Now().Add(1 * time.Minute),
},
expectErr: false,
credsAssertion: assertCredsDoNotExist,
expectErr: false,
credsAssertion: assertCreds(
assertUsernameRegex("^v-test-test-[a-zA-Z0-9]{20}-[0-9]{10}$"),
assertCredsDoNotExist,
),
},
"large block statements": {
req: dbplugin.NewUserRequest{
Expand All @@ -214,8 +237,11 @@ func TestPostgreSQL_NewUser(t *testing.T) {
Password: "somesecurepassword",
Expiration: time.Now().Add(1 * time.Minute),
},
expectErr: false,
credsAssertion: assertCredsExist,
expectErr: false,
credsAssertion: assertCreds(
assertUsernameRegex("^v-test-test-[a-zA-Z0-9]{20}-[0-9]{10}$"),
assertCredsExist,
),
},
}

Expand Down Expand Up @@ -578,6 +604,22 @@ func TestDeleteUser(t *testing.T) {

type credsAssertion func(t testing.TB, connURL, username, password string)

func assertCreds(assertions ...credsAssertion) credsAssertion {
return func(t testing.TB, connURL, username, password string) {
t.Helper()
for _, assertion := range assertions {
assertion(t, connURL, username, password)
}
}
}

func assertUsernameRegex(rawRegex string) credsAssertion {
return func(t testing.TB, _, username, _ string) {
t.Helper()
require.Regexp(t, rawRegex, username)
}
}

func assertCredsExist(t testing.TB, connURL, username, password string) {
t.Helper()
err := testCredsExist(t, connURL, username, password)
Expand Down Expand Up @@ -745,3 +787,173 @@ func TestExtractQuotedStrings(t *testing.T) {
})
}
}

func TestUsernameGeneration(t *testing.T) {
type testCase struct {
data dbplugin.UsernameMetadata
expectedRegex string
}

tests := map[string]testCase{
"simple display and role names": {
data: dbplugin.UsernameMetadata{
DisplayName: "token",
RoleName: "myrole",
},
expectedRegex: `v-token-myrole-[a-zA-Z0-9]{20}-[0-9]{10}`,
},
"display name has dash": {
data: dbplugin.UsernameMetadata{
DisplayName: "token-foo",
RoleName: "myrole",
},
expectedRegex: `v-token-fo-myrole-[a-zA-Z0-9]{20}-[0-9]{10}`,
},
"display name has underscore": {
data: dbplugin.UsernameMetadata{
DisplayName: "token_foo",
RoleName: "myrole",
},
expectedRegex: `v-token_fo-myrole-[a-zA-Z0-9]{20}-[0-9]{10}`,
},
"display name has period": {
data: dbplugin.UsernameMetadata{
DisplayName: "token.foo",
RoleName: "myrole",
},
expectedRegex: `v-token.fo-myrole-[a-zA-Z0-9]{20}-[0-9]{10}`,
},
"role name has dash": {
data: dbplugin.UsernameMetadata{
DisplayName: "token",
RoleName: "myrole-foo",
},
expectedRegex: `v-token-myrole-f-[a-zA-Z0-9]{20}-[0-9]{10}`,
},
"role name has underscore": {
data: dbplugin.UsernameMetadata{
DisplayName: "token",
RoleName: "myrole_foo",
},
expectedRegex: `v-token-myrole_f-[a-zA-Z0-9]{20}-[0-9]{10}`,
},
"role name has period": {
data: dbplugin.UsernameMetadata{
DisplayName: "token",
RoleName: "myrole.foo",
},
expectedRegex: `v-token-myrole.f-[a-zA-Z0-9]{20}-[0-9]{10}`,
},
}

for name, test := range tests {
t.Run(fmt.Sprintf("new-%s", name), func(t *testing.T) {
up, err := template.NewTemplate(
template.Template(defaultUserNameTemplate),
)
require.NoError(t, err)

for i := 0; i < 1000; i++ {
username, err := up.Generate(test.data)
require.NoError(t, err)
require.Regexp(t, test.expectedRegex, username)
}
})
}
}

func TestNewUser_CustomUsername(t *testing.T) {
cleanup, connURL := postgresql.PrepareTestContainer(t, "latest")
defer cleanup()

type testCase struct {
usernameTemplate string
newUserData dbplugin.UsernameMetadata
expectedRegex string
}

tests := map[string]testCase{
"default template": {
usernameTemplate: "",
newUserData: dbplugin.UsernameMetadata{
DisplayName: "displayname",
RoleName: "longrolename",
},
expectedRegex: "^v-displayn-longrole-[a-zA-Z0-9]{20}-[0-9]{10}$",
},
"explicit default template": {
usernameTemplate: defaultUserNameTemplate,
newUserData: dbplugin.UsernameMetadata{
DisplayName: "displayname",
RoleName: "longrolename",
},
expectedRegex: "^v-displayn-longrole-[a-zA-Z0-9]{20}-[0-9]{10}$",
},
"unique template": {
usernameTemplate: "foo-bar",
newUserData: dbplugin.UsernameMetadata{
DisplayName: "displayname",
RoleName: "longrolename",
},
expectedRegex: "^foo-bar$",
},
"custom prefix": {
usernameTemplate: "foobar-{{.DisplayName | truncate 8}}-{{.RoleName | truncate 8}}-{{random 20}}-{{unix_time}}",
newUserData: dbplugin.UsernameMetadata{
DisplayName: "displayname",
RoleName: "longrolename",
},
expectedRegex: "^foobar-displayn-longrole-[a-zA-Z0-9]{20}-[0-9]{10}$",
},
"totally custom template": {
usernameTemplate: "foobar_{{random 10}}-{{.RoleName | uppercase}}.{{unix_time}}x{{.DisplayName | truncate 5}}",
newUserData: dbplugin.UsernameMetadata{
DisplayName: "displayname",
RoleName: "longrolename",
},
expectedRegex: `^foobar_[a-zA-Z0-9]{10}-LONGROLENAME\.[0-9]{10}xdispl$`,
},
}

for name, test := range tests {
t.Run(name, func(t *testing.T) {
initReq := dbplugin.InitializeRequest{
Config: map[string]interface{}{
"connection_url": connURL,
"username_template": test.usernameTemplate,
},
VerifyConnection: true,
}

db := new()

ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
defer cancel()

_, err := db.Initialize(ctx, initReq)
require.NoError(t, err)

newUserReq := dbplugin.NewUserRequest{
UsernameConfig: test.newUserData,
Statements: dbplugin.Statements{
Commands: []string{`
CREATE ROLE "{{name}}" WITH
LOGIN
PASSWORD '{{password}}'
VALID UNTIL '{{expiration}}';
GRANT ALL PRIVILEGES ON ALL TABLES IN SCHEMA public TO "{{name}}";`,
},
},
Password: "myReally-S3curePassword",
Expiration: time.Now().Add(1 * time.Hour),
}
ctx, cancel = context.WithTimeout(context.Background(), 2*time.Second)
defer cancel()

newUserResp, err := db.NewUser(ctx, newUserReq)
require.NoError(t, err)

require.Regexp(t, test.expectedRegex, newUserResp.Username)
})
}
}
Loading

0 comments on commit cf85a86

Please sign in to comment.