DEV Community

Clavin June
Clavin June

Posted on • Originally published at clavinjune.dev on

Handle Database Credential Rotation in Golang

Photo by @picoftasty on Unsplash

Introduction

Many of you might have heard of the database credential rotation. It is a common solution for securing your database. Many providers have implemented this solution such as Hashicorp Vault, AWS Secret Manager, etc. By automatically rotating your database credential, no one can really know how to connect to the database. Indeed it is secure, but it leads to another problem. By rotating your database credential, your application might not be able to connect to the database after some times. Especially when you declare your database maximum connection lifetime. In this article, you will learn how to rotate your database credential programmatically in golang.

Custom Driver

To implement this rotator as a generic solution, you will need to create a customDriver struct which holds your base database driver (e.g postgres, sqlite3, mysql, etc) and a Fetcher interface which will be used to fetch the database credential.

type Fetcher interface {
    Fetch() (string, error)
}

type FetcherFunc func() (string, error)
func (f FetcherFunc) Fetch() (string, error) {
    return f()
}

// customDriver implements the `sql.Driver` interface.
type customDriver struct {
    // base is the base database driver
    base      driver.Driver

    // fetcherFn is the function that will be used to fetch the database credential
    fetcherFn FetcherFunc
}

func (d *customDriver) Open(_ string) (driver.Conn, error) {
    // fetch the database credential
    dsn, err := d.fetcherFn()
    if err != nil {
        return nil, err
    }

    // open the database connection using the fetched credential and base driver
    return d.base.Open(dsn)
}
Enter fullscreen mode Exit fullscreen mode

Register the Driver and Open the Connection

Now to be able to use the driver, you need to register the driver and open the connection.

func OpenWithRotator(name string, base driver.Driver, fetcher Fetcher) (*sql.DB, error) {
    sql.Register(name, &customDriver{
        base:      base,
        fetcherFn: fetcher.Fetch,
    })

    // you don't need to fill the dsn, it will be fetched from the fetcher.
    return sql.Open(name, "")
}
Enter fullscreen mode Exit fullscreen mode

Implement Fetcher Interface

Let's implement the Fetcher interface by using a simple function. Let's assume that you want to connect to a different database for every time your database lifetime is up.

var counter int

func simpleFetcher() (string, error) {
    log.Println("fetcher called")

    // add your custom logic
    // e.g. fetching from vault / config / etc.

    counter++
    return fmt.Sprintf("file:foobar-%d.sqlite", counter), nil
}
Enter fullscreen mode Exit fullscreen mode

Open the Connection with Rotator

Open and set the maximum lifetime of the connection to 2 seconds. simpleFetcher will be called every 2 seconds.

If your database credential is rotated faster than your connection lifetime, Golang still can use the old connection.

// you can adjust the `&sqlite3.SQLiteDriver{}` accordingly (e.g. &pq.Driver{}, etc.)
db, err := OpenWithRotator("foobar", &sqlite3.SQLiteDriver{}, FetcherFunc(simpleFetcher))
if err != nil {
    log.Fatal(err)
}
defer db.Close()
db.SetConnMaxLifetime(2 * time.Second)
Enter fullscreen mode Exit fullscreen mode

Test the Connection

To simply test whether the connection is working, you can use the Ping method every second.

for range time.Tick(time.Second) {
    if err := db.Ping(); err != nil {
        log.Fatal(err)
    }
}
Enter fullscreen mode Exit fullscreen mode

Now when you run the program, you will see the following output:

$ go run .
2022/08/29 16:56:16 fetcher called
2022/08/29 16:56:19 fetcher called
2022/08/29 16:56:22 fetcher called
2022/08/29 16:56:25 fetcher called
2022/08/29 16:56:28 fetcher called
2022/08/29 16:56:31 fetcher called
Enter fullscreen mode Exit fullscreen mode

Here's the sqlite3 database files:

$ ls | grep foobar-
foobar-1.sqlite
foobar-2.sqlite
foobar-3.sqlite
foobar-4.sqlite
foobar-5.sqlite
foobar-6.sqlite
Enter fullscreen mode Exit fullscreen mode

Here's the full code:

package main

import (
    "database/sql"
    "database/sql/driver"
    "fmt"
    "log"
    "time"

    "github.com/mattn/go-sqlite3"
)

type Fetcher interface {
    Fetch() (string, error)
}

type FetcherFunc func() (string, error)
func (f FetcherFunc) Fetch() (string, error) {
    return f()
}

// customDriver implements the `sql.Driver` interface.
type customDriver struct {
    // base is the base database driver
    base      driver.Driver

    // fetcherFn is the function that will be used to fetch the database credential
    fetcherFn FetcherFunc
}

func (d *customDriver) Open(_ string) (driver.Conn, error) {
    // fetch the database credential
    dsn, err := d.fetcherFn()
    if err != nil {
        return nil, err
    }

    // open the database connection using the fetched credential and base driver
    return d.base.Open(dsn)
}

func OpenWithRotator(name string, base driver.Driver, fetcher Fetcher) (*sql.DB, error) {
    sql.Register(name, &customDriver{
        base:      base,
        fetcherFn: fetcher.Fetch,
    })

    // you don't need to fill the dsn, it will be fetched from the fetcher.
    return sql.Open(name, "")
}

var counter int
func simpleFetcher() (string, error) {
    log.Println("fetcher called")

    // add your custom logic
    // e.g. fetching from vault / config / etc.

    counter++
    return fmt.Sprintf("file:foobar-%d.sqlite", counter), nil
}

func main() {
    // you can adjust the `&sqlite3.SQLiteDriver{}` accordingly (e.g. &pq.Driver{}, etc.)
    db, err := OpenWithRotator("foobar", &sqlite3.SQLiteDriver{}, FetcherFunc(simpleFetcher))
    if err != nil {
        log.Fatal(err)
    }
    defer db.Close()
    db.SetConnMaxLifetime(2 * time.Second)

    for range time.Tick(time.Second) {
        if err := db.Ping(); err != nil {
            log.Fatal(err)
        }
    }
}
Enter fullscreen mode Exit fullscreen mode

Conclusion

This is the simple implementation of a database rotator. You might want to implement your fetcher accordingly depends on your use case.

Thank you for reading!

Top comments (0)