Using Refreshed IdP Tokens (Before Expiry) to Connect to CockroachDB via Connection Pools

Add refreshed tokens from Okta to your connection pool seamlessly. Check token expiration, or use your own TTL schedule.

Overview

In this blog we will be expanding on the concepts covered in my previous blog, Cockroach Cloud SQL Connection via SSO with Okta IdP - Golang.

Now that we have some of the basic concepts in place, we can expand on our knowledge of refresh tokens to use them in a more realistic manner. In this article we will simulate a long running service that utilizes connection pooling to connect to CockroachDB. We will be updating the connections in the connection pool ~30 seconds by setting/checking an expiration flag in our app.

High Level Steps

Get initial ID token and refresh token from Okta IdP using username and password. This simulates initial login by a user.
Save refresh token as environment variable to be used later at our discretion.
Construct configuration for a connection pool, using our initial token.
Before a connection, check expiration
(Optional) Before acquiring a connection from a pool, check expiration
(Optional) Before releasing a connection back to the pool, check expiration
Run a constant ticker in the background to force a refresh on our ID token every 30 seconds
Execute a "long running" workload that goes through multiple iterations of token refreshes while running

As you can see in steps 4-6, there are multiple places you can inject extra functionality into your config. These may not all be necessary for you, but I've added them for example purposes. Feel free to read more about these configs here.

The other main thing to call out is step 7. In this example I am using a global bool variable to keep track of whether I want to refresh a token or not. You could also leverage the API to check token expirations and determine your own refresh period based off these values.

Pre-reqs

This article will require the same pre-reqs as the previous article. CockroachDB cluster with settings enabled, Okta application, environment variables, etc. I will also link my colleagues first article and second article for further reading.

We will be beginning with the code base from the first article. Here is the GitHub link.

And here is the repo for all completed code after the changes below

Global Variables

The first is pretty self explanatory. This is the pool connection that will now be used throughout our app. isTokenExpired is a boolean that we will be toggling in order to quickly test refreshing tokens in the connection pool. More on setting these later.

var dbPool *pgxpool.Pool
var isTokenExpired = false

Create Config for Connection Pool

The bulk of what we care about will be in this function. Here we will build up our connection pool config.

Method Signature

The method signature should look pretty familiar to the previous blog, since we'll need to leverage most of the same information. The main difference is that we'll be returning a pgxpool.Config.

func getConfig(response OktaResponse, oktaUrl string, clientID string, clientSecret string) (*pgxpool.Config, error) {

Connection String

Next we can build up our connection string just as before, and use pgxpool.ParseConfig() to do the extracting. This code block is near identical, and is currently living in the executeWorkload() func. You can pull it from that func completely.

// Create connection string with initial ID token
// Update the next 3 variables in order to complete your DB connection string
sqlUser := "sqlUser"
host := "host"
cert := "/ca.cert"
dbURL := "postgresql://" + sqlUser + ":" + response.IdToken + "@" + host + ":26257/defaultdb?sslmode=verify-full&sslrootcert=" + cert + "&options=--crdb:jwt_auth_enabled=true"

// Set initial connection pool configuration
config, err := pgxpool.ParseConfig(dbURL)
if err != nil {
    log.Fatal("error configuring the pool: ", err)
}

Additional Configs

These values are arbitrary in this demo. You can read more about sizing connection pools here.

config.MaxConns = 4
config.MaxConnLifetime = (120 * time.Second)

BeforeConnect (Our Bread and Butter)

This is the key piece our new functionality.

BeforeConnect is called before a new connection is made. It is passed a copy of the underlying pgx.ConnConfig and will not impact any existing open connections.

First we will check to see if isTokenExpired is true. If so, we know we need to refresh our token and update our config. If not, we can proceed as normal with the token we already set as the password a few lines earlier. Again, for us we just check our boolean, but this could just as easily be a function to fit your specific needs.

After learning a token is expired we can grab our saved refresh token (we'll get there) to request a new ID token and refresh token. The useRefreshToken function call should look familiar since we used it in the first article :)

We'll then update the config.Password with this new token, update with our latest refresh token, and then set isTokenExpired back to false again.

func getConfig(...) {
...
...
...

config.BeforeConnect = func(ctx context.Context, config *pgx.ConnConfig) error {
    // check global variable
    if isTokenExpired {
        fmt.Println("ID Token is expired, issuing refresh")

        currentRefreshToken := os.Getenv("REFRESH_TOKEN")
        idToken, refreshToken := useRefreshToken(currentRefreshToken, oktaUrl, clientID, clientSecret)
        config.Password = idToken
        os.Setenv("REFRESH_TOKEN", refreshToken)
        isTokenExpired = false
    }

    return nil
}

return config, nil

This concludes the getConfig func, which ends up being pretty straight forward once you decide when and how you want to determine an expiration.

Lay the Groundwork

The rest of our work we will throw in our main() function.

Get Env Variables

This is the left untouched from before.

oktaUrl := os.Getenv("OKTA_URL")
clientID := os.Getenv("CLIENT_ID")
clientSecret := os.Getenv("CLIENT_SECRET")
oktaUsername := os.Getenv("OKTA_USERNAME")
oktaPassword := os.Getenv("PASSWORD")

Get Initial ID Token and Refresh Token

This simulates an initial user login with a username and password. We will leverage the function just as before, although this time the whole OktaResponse object is returned.

In addition, we will set an environment variable that tracks our refresh token. Remember, we grabbed this variable later in our BeforeConnect func.

resp := getTokens(oktaUrl, clientID, clientSecret, oktaUsername, oktaPassword)
fmt.Println("Received initial ID Token")

// Set env variable that tracks the refresh token to use
os.Setenv("REFRESH_TOKEN", resp.RefreshToken)

Create Connection Pool

Now we can call and leverage our getConfig function from earlier. We will pass our initial response from Okta that contains the ID token we know is good.

Construct Config

// Construct config for db pool
config, err := getConfig(resp, oktaUrl, clientID, clientSecret)
if err != nil {
    log.Fatal("error setting pgxpool config: ", err)
}

Additional Config Settings (Optional)

I also added additional config settings here to check token expiration. These may or may not be necessary for your particular use case, but I found them useful to close off the connection the second my timer ran out.

You can read more about these additional configs here.

// Ensure token is not expired before acquiring a connection from the pool
config.BeforeAcquire = func(ctx context.Context, c *pgx.Conn) bool {
    return !isTokenExpired
}

// Ensure token is not expired before releasing connection back to the pool
config.AfterRelease = func(c *pgx.Conn) bool {
    return !isTokenExpired
}

Initialize Connection Pool (Finally)

With all of our configs in place, we can finally initialize our connection pool dbPool.

dbPool, err = pgxpool.ConnectConfig(context.Background(), config)
if err != nil {
    log.Fatal("error creating pool: ", err)
}
defer dbPool.Close()

Update References to `conn`

In our code theres many references to conn. We can now update this to dbPool. *pgx.Conn can also be changed to *pgxpool.Pool.

Create Ticker

The ticker is what we will use to toggle our boolean and indicate that we need to refresh our token. I decided to set it for 30 seconds so we could easily see the results in action.

The main piece to notice here is the 30 second marker, as well as toggling isTokenExpired to true.

ticker := time.NewTicker(30 * time.Second)
quit := make(chan struct{})
go func() {
    for {
        select {
        case <-ticker.C:
            fmt.Println("Ticker set ID token to expired")
            isTokenExpired = true
        case <-quit:
            ticker.Stop()
            return
        }
    }
}()

Run the Workload

I kept this piece pretty simple. We will be leveraging the same executeWorkload() function we already have in place that will perform CRUD operations on the accounts table in our database.

The loop will execute the workload, and then pause 6 seconds before running again. I am pausing just so we can see our refresh happen easier. You could run this continually, or in multiple threads, and see the same results.

for i := 0; i < 10; i++ {
    executeWorkload()
    time.Sleep(6 * time.Second)
}

The workload runs for ~1 minute, so you should see a token refresh 1-2 times. Feel free to increase the loop time, decrease the ticker time, etc.

Results

Conclusion and Next Steps

Refreshing tokens from and IdP, in conjunction with connection pooling, is a very powerful new tool at our disposal with CockroachDB. I'm looking forward to exploring more use cases.

Next steps include:

Keying off a more realistic expiration. Currently JWT tokens in Okta are hard coded to 60 mins. Maybe refreshing 15 minutes before ID token expiration?
Keeping an eye on refresh token expiration. Default expiration is 100 days (this is configurable). If a user doesn't log into the app during that time, we need to prompt for username and password again.
Store the refresh token in a more secure fashion

Using Refreshed IdP Tokens (Before Expiry) to Connect to CockroachDB via Connection Pools - Go