From f23b14a8c2010a42fabbf965d9e4db0f77ab3aef Mon Sep 17 00:00:00 2001 From: Becca Petrin Date: Tue, 17 Apr 2018 16:31:09 -0700 Subject: [PATCH] Release database resources on each iteration of a loop (#4305) --- builtin/logical/mssql/path_creds_create.go | 9 +-- builtin/logical/mssql/secret_creds.go | 10 +-- builtin/logical/mysql/path_role_create.go | 9 +-- .../logical/postgresql/path_role_create.go | 10 ++- builtin/logical/postgresql/secret_creds.go | 19 ++---- helper/dbtxn/dbtxn.go | 63 +++++++++++++++++++ plugins/database/hana/hana.go | 17 ++--- plugins/database/mssql/mssql.go | 37 +++-------- plugins/database/mysql/mysql.go | 14 ++--- plugins/database/postgresql/postgresql.go | 47 ++++---------- 10 files changed, 112 insertions(+), 123 deletions(-) create mode 100644 helper/dbtxn/dbtxn.go diff --git a/builtin/logical/mssql/path_creds_create.go b/builtin/logical/mssql/path_creds_create.go index 7e26937016..1a954c87b4 100644 --- a/builtin/logical/mssql/path_creds_create.go +++ b/builtin/logical/mssql/path_creds_create.go @@ -6,6 +6,7 @@ import ( "strings" "github.com/hashicorp/go-uuid" + "github.com/hashicorp/vault/helper/dbtxn" "github.com/hashicorp/vault/helper/strutil" "github.com/hashicorp/vault/logical" "github.com/hashicorp/vault/logical/framework" @@ -90,15 +91,11 @@ func (b *backend) pathCredsCreateRead(ctx context.Context, req *logical.Request, continue } - stmt, err := tx.Prepare(Query(query, map[string]string{ + m := map[string]string{ "name": username, "password": password, - })) - if err != nil { - return nil, err } - defer stmt.Close() - if _, err := stmt.Exec(); err != nil { + if err := dbtxn.ExecuteTxQuery(ctx, tx, m, query); err != nil { return nil, err } } diff --git a/builtin/logical/mssql/secret_creds.go b/builtin/logical/mssql/secret_creds.go index ce93643eec..5edac67be5 100644 --- a/builtin/logical/mssql/secret_creds.go +++ b/builtin/logical/mssql/secret_creds.go @@ -6,6 +6,7 @@ import ( "fmt" "github.com/hashicorp/errwrap" + "github.com/hashicorp/vault/helper/dbtxn" "github.com/hashicorp/vault/logical" "github.com/hashicorp/vault/logical/framework" ) @@ -130,16 +131,11 @@ func (b *backend) secretCredsRevoke(ctx context.Context, req *logical.Request, d // many permissions as possible right now var lastStmtError error for _, query := range revokeStmts { - stmt, err := db.Prepare(query) - if err != nil { + + if err := dbtxn.ExecuteDBQuery(ctx, db, nil, query); err != nil { lastStmtError = err continue } - defer stmt.Close() - _, err = stmt.Exec() - if err != nil { - lastStmtError = err - } } // can't drop if not all database users are dropped diff --git a/builtin/logical/mysql/path_role_create.go b/builtin/logical/mysql/path_role_create.go index ae184bdb34..135587575d 100644 --- a/builtin/logical/mysql/path_role_create.go +++ b/builtin/logical/mysql/path_role_create.go @@ -6,6 +6,7 @@ import ( "strings" "github.com/hashicorp/go-uuid" + "github.com/hashicorp/vault/helper/dbtxn" "github.com/hashicorp/vault/helper/strutil" "github.com/hashicorp/vault/logical" "github.com/hashicorp/vault/logical/framework" @@ -103,15 +104,11 @@ func (b *backend) pathRoleCreateRead(ctx context.Context, req *logical.Request, continue } - stmt, err := tx.Prepare(Query(query, map[string]string{ + m := map[string]string{ "name": username, "password": password, - })) - if err != nil { - return nil, err } - defer stmt.Close() - if _, err := stmt.Exec(); err != nil { + if err := dbtxn.ExecuteTxQuery(ctx, tx, m, query); err != nil { return nil, err } } diff --git a/builtin/logical/postgresql/path_role_create.go b/builtin/logical/postgresql/path_role_create.go index 7fed904667..113e3ab896 100644 --- a/builtin/logical/postgresql/path_role_create.go +++ b/builtin/logical/postgresql/path_role_create.go @@ -7,6 +7,7 @@ import ( "time" "github.com/hashicorp/go-uuid" + "github.com/hashicorp/vault/helper/dbtxn" "github.com/hashicorp/vault/helper/strutil" "github.com/hashicorp/vault/logical" "github.com/hashicorp/vault/logical/framework" @@ -106,16 +107,13 @@ func (b *backend) pathRoleCreateRead(ctx context.Context, req *logical.Request, continue } - stmt, err := tx.Prepare(Query(query, map[string]string{ + m := map[string]string{ "name": username, "password": password, "expiration": expiration, - })) - if err != nil { - return nil, err } - defer stmt.Close() - if _, err := stmt.Exec(); err != nil { + + if err := dbtxn.ExecuteTxQuery(ctx, tx, m, query); err != nil { return nil, err } } diff --git a/builtin/logical/postgresql/secret_creds.go b/builtin/logical/postgresql/secret_creds.go index 87748d1ff9..d00beacb65 100644 --- a/builtin/logical/postgresql/secret_creds.go +++ b/builtin/logical/postgresql/secret_creds.go @@ -8,6 +8,7 @@ import ( "time" "github.com/hashicorp/errwrap" + "github.com/hashicorp/vault/helper/dbtxn" "github.com/hashicorp/vault/helper/strutil" "github.com/hashicorp/vault/logical" "github.com/hashicorp/vault/logical/framework" @@ -211,14 +212,7 @@ func (b *backend) secretCredsRevoke(ctx context.Context, req *logical.Request, d // 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 { + if err := dbtxn.ExecuteDBQuery(ctx, db, nil, query); err != nil { lastStmtError = err } } @@ -258,15 +252,10 @@ func (b *backend) secretCredsRevoke(ctx context.Context, req *logical.Request, d continue } - stmt, err := tx.Prepare(Query(query, map[string]string{ + m := map[string]string{ "name": username, - })) - if err != nil { - return nil, err } - defer stmt.Close() - - if _, err := stmt.Exec(); err != nil { + if err := dbtxn.ExecuteTxQuery(ctx, tx, m, query); err != nil { return nil, err } } diff --git a/helper/dbtxn/dbtxn.go b/helper/dbtxn/dbtxn.go new file mode 100644 index 0000000000..3337bd97b2 --- /dev/null +++ b/helper/dbtxn/dbtxn.go @@ -0,0 +1,63 @@ +package dbtxn + +import ( + "context" + "database/sql" + "fmt" + "strings" +) + +// ExecuteDBQuery handles executing one single statement, while properly releasing its resources. +// - ctx: Required +// - db: Required +// - config: Optional, may be nil +// - query: Required +func ExecuteDBQuery(ctx context.Context, db *sql.DB, params map[string]string, query string) error { + + parsedQuery := parseQuery(params, query) + + stmt, err := db.PrepareContext(ctx, parsedQuery) + if err != nil { + return err + } + defer stmt.Close() + + return execute(ctx, stmt) +} + +// ExecuteTxQuery handles executing one single statement, while properly releasing its resources. +// - ctx: Required +// - tx: Required +// - config: Optional, may be nil +// - query: Required +func ExecuteTxQuery(ctx context.Context, tx *sql.Tx, params map[string]string, query string) error { + + parsedQuery := parseQuery(params, query) + + stmt, err := tx.PrepareContext(ctx, parsedQuery) + if err != nil { + return err + } + defer stmt.Close() + + return execute(ctx, stmt) +} + +func execute(ctx context.Context, stmt *sql.Stmt) error { + if _, err := stmt.ExecContext(ctx); err != nil { + return err + } + return nil +} + +func parseQuery(m map[string]string, tpl string) string { + + if m == nil || len(m) <= 0 { + return tpl + } + + for k, v := range m { + tpl = strings.Replace(tpl, fmt.Sprintf("{{%s}}", k), v, -1) + } + return tpl +} diff --git a/plugins/database/hana/hana.go b/plugins/database/hana/hana.go index 1fdafe77ad..62e739a669 100644 --- a/plugins/database/hana/hana.go +++ b/plugins/database/hana/hana.go @@ -11,6 +11,7 @@ import ( _ "github.com/SAP/go-hdb/driver" "github.com/hashicorp/vault/api" "github.com/hashicorp/vault/builtin/logical/database/dbplugin" + "github.com/hashicorp/vault/helper/dbtxn" "github.com/hashicorp/vault/helper/strutil" "github.com/hashicorp/vault/plugins" "github.com/hashicorp/vault/plugins/helper/database/connutil" @@ -143,16 +144,12 @@ func (h *HANA) CreateUser(ctx context.Context, statements dbplugin.Statements, u continue } - stmt, err := tx.PrepareContext(ctx, dbutil.QueryHelper(query, map[string]string{ + m := map[string]string{ "name": username, "password": password, "expiration": expirationStr, - })) - if err != nil { - return "", "", err } - defer stmt.Close() - if _, err := stmt.ExecContext(ctx); err != nil { + if err := dbtxn.ExecuteTxQuery(ctx, tx, m, query); err != nil { return "", "", err } } @@ -238,14 +235,10 @@ func (h *HANA) RevokeUser(ctx context.Context, statements dbplugin.Statements, u continue } - stmt, err := tx.PrepareContext(ctx, dbutil.QueryHelper(query, map[string]string{ + m := map[string]string{ "name": username, - })) - if err != nil { - return err } - defer stmt.Close() - if _, err := stmt.ExecContext(ctx); err != nil { + if err := dbtxn.ExecuteTxQuery(ctx, tx, m, query); err != nil { return err } } diff --git a/plugins/database/mssql/mssql.go b/plugins/database/mssql/mssql.go index 84f7e1462a..9b0a78c0ae 100644 --- a/plugins/database/mssql/mssql.go +++ b/plugins/database/mssql/mssql.go @@ -12,6 +12,7 @@ import ( "github.com/hashicorp/errwrap" "github.com/hashicorp/vault/api" "github.com/hashicorp/vault/builtin/logical/database/dbplugin" + "github.com/hashicorp/vault/helper/dbtxn" "github.com/hashicorp/vault/helper/strutil" "github.com/hashicorp/vault/plugins" "github.com/hashicorp/vault/plugins/helper/database/connutil" @@ -129,16 +130,13 @@ func (m *MSSQL) CreateUser(ctx context.Context, statements dbplugin.Statements, continue } - stmt, err := tx.PrepareContext(ctx, dbutil.QueryHelper(query, map[string]string{ + m := map[string]string{ "name": username, "password": password, "expiration": expirationStr, - })) - if err != nil { - return "", "", err } - defer stmt.Close() - if _, err := stmt.ExecContext(ctx); err != nil { + + if err := dbtxn.ExecuteTxQuery(ctx, tx, m, query); err != nil { return "", "", err } } @@ -189,14 +187,10 @@ func (m *MSSQL) RevokeUser(ctx context.Context, statements dbplugin.Statements, continue } - stmt, err := tx.PrepareContext(ctx, dbutil.QueryHelper(query, map[string]string{ + m := map[string]string{ "name": username, - })) - if err != nil { - return err } - defer stmt.Close() - if _, err := stmt.ExecContext(ctx); err != nil { + if err := dbtxn.ExecuteTxQuery(ctx, tx, m, query); err != nil { return err } } @@ -285,14 +279,7 @@ func (m *MSSQL) revokeUserDefault(ctx context.Context, username string) error { // many permissions as possible right now var lastStmtError error for _, query := range revokeStmts { - stmt, err := db.PrepareContext(ctx, query) - if err != nil { - lastStmtError = err - continue - } - defer stmt.Close() - _, err = stmt.ExecContext(ctx) - if err != nil { + if err := dbtxn.ExecuteDBQuery(ctx, db, nil, query); err != nil { lastStmtError = err } } @@ -355,16 +342,12 @@ func (m *MSSQL) RotateRootCredentials(ctx context.Context, statements []string) if len(query) == 0 { continue } - stmt, err := tx.PrepareContext(ctx, dbutil.QueryHelper(query, map[string]string{ + + m := map[string]string{ "username": m.Username, "password": password, - })) - if err != nil { - return nil, err } - - defer stmt.Close() - if _, err := stmt.ExecContext(ctx); err != nil { + if err := dbtxn.ExecuteTxQuery(ctx, tx, m, query); err != nil { return nil, err } } diff --git a/plugins/database/mysql/mysql.go b/plugins/database/mysql/mysql.go index 00fe475045..a36f1a8686 100644 --- a/plugins/database/mysql/mysql.go +++ b/plugins/database/mysql/mysql.go @@ -10,6 +10,7 @@ import ( stdmysql "github.com/go-sql-driver/mysql" "github.com/hashicorp/vault/api" "github.com/hashicorp/vault/builtin/logical/database/dbplugin" + "github.com/hashicorp/vault/helper/dbtxn" "github.com/hashicorp/vault/helper/strutil" "github.com/hashicorp/vault/plugins" "github.com/hashicorp/vault/plugins/helper/database/connutil" @@ -182,10 +183,11 @@ func (m *MySQL) CreateUser(ctx context.Context, statements dbplugin.Statements, return "", "", err } - defer stmt.Close() if _, err := stmt.ExecContext(ctx); err != nil { + stmt.Close() return "", "", err } + stmt.Close() } } @@ -291,16 +293,12 @@ func (m *MySQL) RotateRootCredentials(ctx context.Context, statements []string) if len(query) == 0 { continue } - stmt, err := tx.PrepareContext(ctx, dbutil.QueryHelper(query, map[string]string{ + + m := map[string]string{ "username": m.Username, "password": password, - })) - if err != nil { - return nil, err } - - defer stmt.Close() - if _, err := stmt.ExecContext(ctx); err != nil { + if err := dbtxn.ExecuteTxQuery(ctx, tx, m, query); err != nil { return nil, err } } diff --git a/plugins/database/postgresql/postgresql.go b/plugins/database/postgresql/postgresql.go index c56f9ed02d..36dd0036a9 100644 --- a/plugins/database/postgresql/postgresql.go +++ b/plugins/database/postgresql/postgresql.go @@ -11,6 +11,7 @@ import ( "github.com/hashicorp/errwrap" "github.com/hashicorp/vault/api" "github.com/hashicorp/vault/builtin/logical/database/dbplugin" + "github.com/hashicorp/vault/helper/dbtxn" "github.com/hashicorp/vault/helper/strutil" "github.com/hashicorp/vault/plugins" "github.com/hashicorp/vault/plugins/helper/database/connutil" @@ -139,16 +140,12 @@ func (p *PostgreSQL) CreateUser(ctx context.Context, statements dbplugin.Stateme continue } - stmt, err := tx.PrepareContext(ctx, dbutil.QueryHelper(query, map[string]string{ + m := map[string]string{ "name": username, "password": password, "expiration": expirationStr, - })) - if err != nil { - return "", "", err } - defer stmt.Close() - if _, err := stmt.ExecContext(ctx); err != nil { + if err := dbtxn.ExecuteTxQuery(ctx, tx, m, query); err != nil { return "", "", err } } @@ -157,7 +154,6 @@ func (p *PostgreSQL) CreateUser(ctx context.Context, statements dbplugin.Stateme // Commit the transaction if err := tx.Commit(); err != nil { return "", "", err - } return username, password, nil @@ -198,16 +194,12 @@ func (p *PostgreSQL) RenewUser(ctx context.Context, statements dbplugin.Statemen if len(query) == 0 { continue } - stmt, err := tx.PrepareContext(ctx, dbutil.QueryHelper(query, map[string]string{ + + m := map[string]string{ "name": username, "expiration": expirationStr, - })) - if err != nil { - return err } - - defer stmt.Close() - if _, err := stmt.ExecContext(ctx); err != nil { + if err := dbtxn.ExecuteTxQuery(ctx, tx, m, query); err != nil { return err } } @@ -251,15 +243,10 @@ func (p *PostgreSQL) customRevokeUser(ctx context.Context, username string, revo continue } - stmt, err := tx.PrepareContext(ctx, dbutil.QueryHelper(query, map[string]string{ + m := map[string]string{ "name": username, - })) - if err != nil { - return err } - defer stmt.Close() - - if _, err := stmt.ExecContext(ctx); err != nil { + if err := dbtxn.ExecuteTxQuery(ctx, tx, m, query); err != nil { return err } } @@ -352,14 +339,7 @@ func (p *PostgreSQL) defaultRevokeUser(ctx context.Context, username string) err // many permissions as possible right now var lastStmtError error for _, query := range revocationStmts { - stmt, err := db.PrepareContext(ctx, query) - if err != nil { - lastStmtError = err - continue - } - defer stmt.Close() - _, err = stmt.ExecContext(ctx) - if err != nil { + if err := dbtxn.ExecuteDBQuery(ctx, db, nil, query); err != nil { lastStmtError = err } } @@ -423,16 +403,11 @@ func (p *PostgreSQL) RotateRootCredentials(ctx context.Context, statements []str if len(query) == 0 { continue } - stmt, err := tx.PrepareContext(ctx, dbutil.QueryHelper(query, map[string]string{ + m := map[string]string{ "username": p.Username, "password": password, - })) - if err != nil { - return nil, err } - - defer stmt.Close() - if _, err := stmt.ExecContext(ctx); err != nil { + if err := dbtxn.ExecuteTxQuery(ctx, tx, m, query); err != nil { return nil, err } }