12 - Store Refresh Tokens in Redis

Howdy ho, neighbor!

Last time we cleaned up our app by:

• making some structs package private
• using the proper request context, *gin.Context.Request(), to forward down the call chain
• reading in all environment variables in package main
• making sure to send the user a clear error message if they send a non-JSON body.

As there were a lot of changes, you may want to check our the Github repository's branch called lesson-11, to make sure you have properly updated all of your code!

We're in for another big one today, but hopefully the patterns we established while learning to signup users will become even clearer as we store tokens.

You might also find the video useful!

Why We'll Store Refresh Tokens

We recently implemented signing up a user (see parts 8-10 of this series). However, we omitted an important detail in our application - storing refresh tokens. You can see these details below in the calls from NewPairFromUser to the token repository. In the diagram, I've grayed out and checked off the completed portions required to sign up a user.

Why store these tokens at all? Because this gives us the ability to invalidate user's long-lived tokens.

Recall that we set the duration to expiration to 3 days for refresh tokens and 15 minutes for ID tokens. This allows a user with an expired idToken to reach out to this account application with their refreshToken to get a new idToken. The idToken will allow the user to access various applications in our domain.

At the end of the day, we do this to make accessing our applications more pleasant, since it would be a pain in the arse to need to login every 15 minutes.

However, there is a danger in creating a token that lives for 3 days. If the token is stolen, or if a user wants to sign out of all of their devices, they could still access the application for up to 3 days!

For this reason, we're going to create a Redis database that stores valid refresh tokens for users. We use Redis because it's in-memory database (fast), has an easy API for storing key-value pairs, and allows us to easily remove the tokens from the cache after 3 days. A diagram of our storage approach is shown below.

We call this store a white list of refresh tokens. When a user receives an idToken/refreshToken pair, the refresh token will always be stored in Redis. Later on, we'll add a token refresh route and logic to our application. In this case, the user already has a refresh token, which is required to get a new idToken. Nevertheless, we'll expire their old refreshToken in Redis and provide them a fresh one. This is of a type of "revolving" refresh token technique.

I'll admit there are probably more infrastructure-friendly ways to do this, but I think we'll still learn a lot implementing the white list approach, and you can adjust the approach to your needs.

I see that many people create a "black list" of invalidated tokens, instead of a "white list" of valid tokens. In the blacklist approach, if a user wants to reset their password, or log out of all devices, then all of their current refresh tokens would need to be moved to Redis from the main database.

But wait, isn't this the same as a white list? No, because these tokens might originally be stored in the main database, and only sent to the in-memory cache when tokens must be invalidated. This may be a better architectural decision in your case. You'll likely save money using this latter approach, as you would be storing fewer tokens inside of Redis (in-memory is pricier, to the best of my knowledge). You will probably add some complexity, though, by needing to push all of the user's tokens from your main DB to Redis, though that shouldn't be too difficult.

Topics We'll Cover Today

We'll need to perform the following steps to update our application for token storage.

1. Create a Redis Container in docker-compose.yml.
2. Instantiate a connection in package main. Recall that we have a data_sources.go file for initializing our data.
3. Add a TokenRepository interface to our ~/model/interfaces.go and specify the method signatures for SetRefreshToken and DeleteRefreshToken.
4. Update our TokenService struct and NewPairFromUser method so that we can store the user's current refresh token, and delete their previous refresh token (will be used in a refresh token endpoint will implement down the road).
5. Add a TokenRepository mock and update the unit test for NewPairFromUser.
6. Create a RedisTokenRepository and implement SetRefreshToken and DeleteRefreshToken.
7. Make sure to inject our RedisTokenRepository into TokenService and run the application.

Add Redis to docker-compose.yaml

Add the following service at the same indentation level as other services (e.g., postgres-account) in docker-compose.yml, located at the project root.

redis-account:
    image: "redis:alpine"
    ports:
      - "6379:6379"
    volumes:
      - "redisdata:/data"

We'll also need to make sure to add redis-account to the depends_on list of account

account:
  # ... content omitted here ...
  depends_on:
    - postgres-account
    - redis-account

Also include the volume redisdata under the volumes key at the bottom of the file.

volumes:
  pgdata_account:
  redisdata:

That's all we need to create our Redis container!

Establish Redis Connection

We'll establish the application's connection to Redis inside of ~/data_sources.go of package main. We'll be using the Redis client package github.com/go-redis/redis/v8. Make sure to add this package to go.mod by running go get github.com/go-redis/redis/v8 from inside of the account folder.

Then, update the dataSources struct as follows:

type dataSources struct {
    DB          *sqlx.DB
    RedisClient *redis.Client
}

Then, inside of initDs(), establish a the connection as follows, and make sure to return this connection in dataSources.

// Initialize redis connection
    redisHost := os.Getenv("REDIS_HOST")
    redisPort := os.Getenv("REDIS_PORT")

    log.Printf("Connecting to Redis\n")
    rdb := redis.NewClient(&redis.Options{
        Addr:     fmt.Sprintf("%s:%s", redisHost, redisPort),
        Password: "",
        DB:       0,
    })

    // verify redis connection

    _, err = rdb.Ping(context.Background()).Result()

    if err != nil {
        return nil, fmt.Errorf("error connecting to redis: %w", err)
    }

    return &dataSources{
        DB:          db,
        RedisClient: rdb,
    }, nil

Notice that we're making use of the REDIS_HOST and REDIS_PORT environment variables. Therefore, we need to add the host (inside of Docker networking, this will be the name of the service) and the default Redis port to our ~/account/env.dev file.

REDIS_HOST=redis-account
REDIS_PORT=6379

Finally, we want to make sure we close the connection in our close method at the bottom of the file.

// close to be used in graceful server shutdown
func (d *dataSources) close() error {
    // ...code omitted

    if err := d.RedisClient.Close(); err != nil {
        return fmt.Errorf("error closing Redis Client: %w", err)
    }

    return nil
}

You should now be able spin up the application by running docker-compose up inside of the project root folder! Hopefully you'll see a log showing your Redis container is up and running and that the application has established a connection!

Create Token Repository Interface

We have yet to define expectations for a token repository. We can do so inside ~/model/interfaces.go of the model layer. We'll also add the first two methods for this repository. We'll soon see how these methods are used.

// TokenRepository defines methods it expects a repository
// it interacts with to implement
type TokenRepository interface {
    SetRefreshToken(ctx context.Context, userID string, tokenID string, expiresIn time.Duration) error
    DeleteRefreshToken(ctx context.Context, userID string, prevTokenID string) error
}

Access TokenRepository from tokenService

We will access a TokenRepository interface from inside of our tokenService implementation. Currently this service has a single method, NewPairFromUser, which will access the repository methods we just created.

Let's make sure to add model.TokenRepository as a dependency to tokenService in ~/service/token_service.go as follows:

// tokenService used for injecting an implementation of TokenRepository
// for use in service methods along with keys and secrets for
// signing JWTs
type tokenService struct {
    TokenRepository       model.TokenRepository
    PrivKey               *rsa.PrivateKey
    PubKey                *rsa.PublicKey
    RefreshSecret         string
    IDExpirationSecs      int64
    RefreshExpirationSecs int64
}

// TSConfig will hold repositories that will eventually be injected into this
// this service layer
type TSConfig struct {
    TokenRepository       model.TokenRepository
    PrivKey               *rsa.PrivateKey
    PubKey                *rsa.PublicKey
    RefreshSecret         string
    IDExpirationSecs      int64
    RefreshExpirationSecs int64
}

// NewTokenService is a factory function for
// initializing a UserService with its repository layer dependencies
func NewTokenService(c *TSConfig) model.TokenService {
    return &tokenService{
        TokenRepository:       c.TokenRepository,
        PrivKey:               c.PrivKey,
        PubKey:                c.PubKey,
        RefreshSecret:         c.RefreshSecret,
        IDExpirationSecs:      c.IDExpirationSecs,
        RefreshExpirationSecs: c.RefreshExpirationSecs,
    }
}

We'll now add the logic for storing and deleting the tokens in NewPairFromUser. We do this after the tokens have been generated. Check out part 9 if you want to learn more about generating tokens.

// set freshly minted refresh token to valid list
    if err := s.TokenRepository.SetRefreshToken(ctx, u.UID.String(), refreshToken.ID, refreshToken.ExpiresIn); err != nil {
        log.Printf("Error storing tokenID for uid: %v. Error: %v\n", u.UID, err.Error())
        return nil, apperrors.NewInternal()
    }

    // delete user's current refresh token (used when refreshing idToken)
    if prevTokenID != "" {
        if err := s.TokenRepository.DeleteRefreshToken(ctx, u.UID.String(), prevTokenID); err != nil {
            log.Printf("Could not delete previous refreshToken for uid: %v, tokenID: %v\n", u.UID.String(), prevTokenID)
        }
    }

First, we store the token, which is derived from a userID (as a string), refreshToken.ID and refreshToken.expiresIn. The user is passed to the function from the handler layer, and the refreshToken is generated in a utility function called generateRefreshToken. We return an internal server error should any error occur.

Next, we delete the user's current refresh token if one exists. This functionality is not used when signing up a user. In ~/handler/signup.go, we call NewPairFromUser with an empty string for the final parameter, which is prevTokenID. Note that if there is an error deleting an old refresh token, we do not return an error, but merely log this for internal use.

Here is the call in ~/handler/signup.go.

  // create token pair as strings
  tokens, err := h.TokenService.NewPairFromUser(ctx, u, "")

Even though we're not yet deleting refresh tokens, I wanted to add this logic now so that we can complete our NewPairFromUser method. We'll call this method again when we add token refreshing.

Testing Newly Added Functionality

We will now update our ~/service/token_service_test.go to account for the updated logic to NewPairFromUser. As always, we need to mock the behavior or our TokenRepository so that this service doesn't depend on any concrete repository implementation. Let's add a mock, leveraging the testify package as always!

Add Mock Token Repository

Create a ~/model/mocks/token_repository.go file with the following content for mocking a TokenRepository. These mocks will be fairly straightforward as both methods only return errors. As always, please check out the testify documentation or previous videos to better understand these mocks!

package mocks

// ... imports omitted

// MockTokenRepository is a mock type for model.TokenRepository
type MockTokenRepository struct {
    mock.Mock
}

// SetRefreshToken is a mock of model.TokenRepository SetRefreshToken
func (m *MockTokenRepository) SetRefreshToken(ctx context.Context, userID string, tokenID string, expiresIn time.Duration) error {
    ret := m.Called(ctx, userID, tokenID, expiresIn)

    var r0 error

    if ret.Get(0) != nil {
        r0 = ret.Get(0).(error)
    }

    return r0
}

// DeleteRefreshToken is a mock of model.TokenRepository DeleteRefreshToken
func (m *MockTokenRepository) DeleteRefreshToken(ctx context.Context, userID string, prevTokenID string) error {
    ret := m.Called(ctx, userID, prevTokenID)

    var r0 error

    if ret.Get(0) != nil {
        r0 = ret.Get(0).(error)
    }

    return r0
}

Update Unit Test

Let's make the following updates to ~/service/token_service_test.go:

1. assert that TokenRepository.SetRefreshToken gets called inside of the current test ("Returns a token pair with proper values"). Also make sure that TokenRepository.DeleteRefreshToken is called if we pass prevTokenID to NewPairFromUser
2. assert TokenService.NewPairFromUser returns an error if there is an error setting the refresh token.
3. assert that DeleteRefreshToken is not called if we supply an empty string, "", as the prevTokenID.

I want to setup the mock before the t.Run blocks. Previously, we set up the mocks inside of each t.Run block.

To do this, import and instantiate a MockTokenRepository, making sure to supply it to the NewTokenService factory.

  mockTokenRepository := new(mocks.MockTokenRepository)

    // instantiate a common token service to be used by all tests
    tokenService := NewTokenService(&TSConfig{
        TokenRepository:       mockTokenRepository,
        PrivKey:               privKey,
        PubKey:                pubKey,
        RefreshSecret:         secret,
        IDExpirationSecs:      idExp,
        RefreshExpirationSecs: refreshExp,
    })

And just before the t.Run, we add the following:

  // code after instantiating a tokenService

  // include password to make sure it is not serialized
    // since json tag is "-"
    uid, _ := uuid.NewRandom()
    u := &model.User{
        UID:      uid,
        Email:    "bob@bob.com",
        Password: "blarghedymcblarghface",
    }

    // Setup mock call responses in setup before t.Run statements
    uidErrorCase, _ := uuid.NewRandom()
    uErrorCase := &model.User{
        UID:      uidErrorCase,
        Email:    "failure@failure.com",
        Password: "blarghedymcblarghface",
    }
    prevID := "a_previous_tokenID"

    setSuccessArguments := mock.Arguments{
        mock.AnythingOfType("*context.emptyCtx"),
        u.UID.String(),
        mock.AnythingOfType("string"),
        mock.AnythingOfType("time.Duration"),
    }

    setErrorArguments := mock.Arguments{
        mock.AnythingOfType("*context.emptyCtx"),
        uErrorCase.UID.String(),
        mock.AnythingOfType("string"),
        mock.AnythingOfType("time.Duration"),
    }

    deleteWithPrevIDArguments := mock.Arguments{
        mock.AnythingOfType("*context.emptyCtx"),
        u.UID.String(),
        prevID,
    }

    // mock call argument/responses
    mockTokenRepository.On("SetRefreshToken", setSuccessArguments...).Return(nil)
    mockTokenRepository.On("SetRefreshToken", setErrorArguments...).Return(fmt.Errorf("Error setting refresh token"))
    mockTokenRepository.On("DeleteRefreshToken", deleteWithPrevIDArguments...).Return(nil)

Since we'll be setting up our mock responses before the t.Run blocks, I want to be explicit about the arguments that we pass to various method calls. That is why you see arguments stored in mock.Arguments structs. This will then make it easier to check that the methods were called with those specific arguments without having to rewrite all of the arguments in test-case assertions.

Let's now update our existing test to make sure that SetRefreshToken and DeleteRefreshToken are called with correct parameters. In order to test the call of DeleteRefreshToken, I make sure to pass prevID (defined in the test setup) to NewPairFromUser.

  ctx := context.Background()                                    // updated from context.TODO()
  tokenPair, err := tokenService.NewPairFromUser(ctx, u, prevID) // replaced "" with prevID from setup
  assert.NoError(t, err)

  // SetRefreshToken should be called with setSuccessArguments
  mockTokenRepository.AssertCalled(t, "SetRefreshToken", setSuccessArguments...)
  // DeleteRefreshToken should not be called since prevID is ""
  mockTokenRepository.AssertCalled(t, "DeleteRefreshToken", deleteWithPrevIDArguments...)

  // rest of test is unchanged

Take note that I updated context.TODO to context.Background to conform to our mock argument expectations (context.emptyCtx is the type of context.Background, and this is what we've been using in other tests).

Next, we'll add a test for when SetRefreshToken returns an error. In this test, we make sure an error is returned of the correct type, and that DeleteRefreshToken is not called, since and error will be returned from NewPairFromUser before DeleteRefreshToken can be called.

  t.Run("Error setting refresh token", func(t *testing.T) {
        ctx := context.Background()
        _, err := tokenService.NewPairFromUser(ctx, uErrorCase, "")
        assert.Error(t, err) // should return an error

        // SetRefreshToken should be called with setErrorArguments
        mockTokenRepository.AssertCalled(t, "SetRefreshToken", setErrorArguments...)
        // DeleteRefreshToken should not be since SetRefreshToken causes method to return
        mockTokenRepository.AssertNotCalled(t, "DeleteRefreshToken")
    })

Finally, let's add a test for the case when no prevTokenID is provided. We'll only test to make sure DeleteRefreshToken is not called. We've already tested for the response body in the first test.

  t.Run("Empty string provided for prevID", func(t *testing.T) {
        ctx := context.Background()
        _, err := tokenService.NewPairFromUser(ctx, u, "")
        assert.NoError(t, err)

        // SetRefreshToken should be called with setSuccessArguments
        mockTokenRepository.AssertCalled(t, "SetRefreshToken", setSuccessArguments...)
        // DeleteRefreshToken should not be called since prevID is ""
        mockTokenRepository.AssertNotCalled(t, "DeleteRefreshToken")
    })

You should now retest this method by running the following from the account directory:

go test -v ./service -run NewPairFromUser

RedisTokenRepository Implementation

We start by creating a ~/account/repository/redis_token_repository.go file.

As with other repositories and services, we start by defining a dependency struct and a factory function to initialize this struct.

Make sure that you import redis/v8 as auto-import may not import the correct version.

package repository

//... imports omitted

// redisTokenRepository is data/repository implementation
// of service layer TokenRepository
type redisTokenRepository struct {
    Redis *redis.Client
}

// NewTokenRepository is a factory for initializing User Repositories
func NewTokenRepository(redisClient *redis.Client) model.TokenRepository {
    return &redisTokenRepository{
        Redis: redisClient,
    }
}

Next, we'll implement our current two methods on the model.TokenRepository.

// SetRefreshToken stores a refresh token with an expiry time
func (r *redisTokenRepository) SetRefreshToken(ctx context.Context, userID string, tokenID string, expiresIn time.Duration) error {
    // We'll store userID with token id so we can scan (non-blocking)
    // over the user's tokens and delete them in case of token leakage
    key := fmt.Sprintf("%s:%s", userID, tokenID)
    if err := r.Redis.Set(ctx, key, 0, expiresIn).Err(); err != nil {
        log.Printf("Could not SET refresh token to redis for userID/tokenID: %s/%s: %v\n", userID, tokenID, err)
        return apperrors.NewInternal()
    }
    return nil
}

// DeleteRefreshToken used to delete old  refresh tokens
// Services my access this to revolve tokens
func (r *redisTokenRepository) DeleteRefreshToken(ctx context.Context, userID string, tokenID string) error {
    key := fmt.Sprintf("%s:%s", userID, tokenID)
    if err := r.Redis.Del(ctx, key).Err(); err != nil {
        log.Printf("Could not delete refresh token to redis for userID/tokenID: %s/%s: %v\n", userID, tokenID, err)
        return apperrors.NewInternal()
    }

    return nil
}

I've added some comments about the storage of these tokens, but let's make sure it's clear what we're doing!

We store the tokens with a key {userID}:{tokenID}. We append the token to the userID (as opposed to storing the tokenID alone) because this will allow us to use a Redis operation called scan if we need to invalidate all of a user's refresh tokens. This would be necessary if the user wanted to reset a password or sign out of all devices. We'll be able to scan for and delete all keys starting with the user's id. Scan does this in a non-blocking fashion so other operations can be performed on the database.

The DeleteRefreshToken is used to revolve refresh tokens. If a user wants a new idToken, we'll also give them a new refreshToken and invalidate their former token. But I think I'm repeating myself. 😉

Inject RedisTokenRepository and Run Application

Let's take our Redis data source, create a TokenRepository from its RedisClient field, and inject this into our TokenService implementation. The TokenService has already been injected into the handler. We make these updates in ~/injection.go of package main.

/*
     * repository layer
     */
    userRepository := repository.NewUserRepository(d.DB)
    tokenRepository := repository.NewTokenRepository(d.RedisClient)

  tokenService := service.NewTokenService(&service.TSConfig{
        TokenRepository:       tokenRepository,
        PrivKey:               privKey,
        PubKey:                pubKey,
        RefreshSecret:         refreshSecret,
        IDExpirationSecs:      idExp,
        RefreshExpirationSecs: refreshExp,
    })

Use Redis Client to View Cache

Let's re-run the application, sign up a new user, and then check to make sure that user has been created with a token. Note that I've deleted all users from Postgres before running these commands.

curl --location --request POST 'http://malcorp.test/api/account/signup' \
--header 'Content-Type: application/json' \
--data-raw '{
    "email": "jacob@jacob.com",
    "password": "avalidpassword123"
}'

You should get a response with the idToken and refreshToken stored as JWTs:

{
  "tokens":{
  "idToken": "eyJhbGciOiJSUzI1NiIsInR5cCI6IkpXVCJ9.eyJ1c2VyIjp7InVpZCI6ImIyODRmZGQ5LTU2ZjItNDgyNC1hN2IwLTU1NWQ0NWQwYTFiNCIsImVtYWlsIjoiYm9iQGJvYi5jb20iLCJuYW1lIjoiIiwiaW1hZ2VVcmwiOiIiLCJ3ZWJzaXRlIjoiIn0sImV4cCI6MTYwNzM4OTQ5MywiaWF0IjoxNjA3Mzg4NTkzfQ.jeT0WyRN0rN9RuudXr7fG1rRmVtBw_t-efBRq8E3iaLyEgPMgG5SC6y16mheeMtXQ74ksVKAlDVV7siS9D94pNBh2PoknA7H2Pa_9k5RFSxJPw4g-qk2NIMnOQIJ6NBrpbc2g1HozndGYpX7wnoCMDxlJvuzGg3mpMXIXausQdUG7nr5VJH_izksybRdhmW_vaK4ZushH8oFJT8XpW8zmylI3tz7g7ICSeRd7wIFybW3XbMudmIw_NOhV-dxd_UGgELQfZG4WRvoBOgVFhd0NmWhwqtyVb1CL1NeJJ1zItfX5WriXGvLesGfbDBSgRThirWlWs2tRBFNVRi6kJYu0g",
  "refreshToken": "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJ1aWQiOiJiMjg0ZmRkOS01NmYyLTQ4MjQtYTdiMC01NTVkNDVkMGExYjQiLCJleHAiOjE2MDc2NDc3OTMsImp0aSI6ImMyMzYwMjJkLWE5MTAtNGE4NS1iMjdkLTljNGE1MTQ0NDhjNCIsImlhdCI6MTYwNzM4ODU5M30.DDuKOfQBPqXdSDAwfqVLy0vECcGJtUfScoWYRUFH4Gk"
  }
}

You can then copy and paste the refreshToken inside of the jwt.io debugger or in another way that you prefer.

I get the following body for the token:

{
  "uid": "b284fdd9-56f2-4824-a7b0-555d45d0a1b4",
  "exp": 1607647793,
  "jti": "c236022d-a910-4a85-b27d-9c4a514448c4",
  "iat": 1607388593
}

The jti key is where we store the refresh token ID. You can then check Redis for a token with the uid and jti separated by a colon: b284fdd9-56f2-4824-a7b0-555d45d0a1b4:c236022d-a910-4a85-b27d-9c4a514448c4.

There are various options to look at a Redis DB, including a GUI version.

For this tutorial, I'll show how to use the CLI version.

You can check for the existence of the key as follows (we're using the default Redis port, so you need not do anything special to connect from your local machine):

redis-cli get b284fdd9-56f2-4824-a7b0-555d45d0a1b4:c236022d-a910-4a85-b27d-9c4a514448c4

And you should receive a response of "0", the value we stored under the key (because I believe you have to store a value for each key).

redis-cli get b284fdd9-56f2-4824-a7b0-555d45d0a1b4:c236022d-a910-4a85-b27d-9c4a514448c4
"0"

You can also check the expiration on the key with the TTL (time to live) command.

➜ redis-cli TTL b284fdd9-56f2-4824-a7b0-555d45d0a1b4:c236022d-a910-4a85-b27d-9c4a514448c4
(integer) 258634

It looks like our time is just shy of 3 days, in seconds!

Conclusion

That was another big one! Thanks for hanging around! We're starting to reap the benefits of our application's architecture. Sure, it was a pain in the arse to set it up, but now I'm quite enjoying connecting everything!

Next time, we'll add middleware to set a handler timeout for each individual request. The tutorial won't involve modifying so many files, but it will be interesting in terms of working with Go and address some concurrency principles!

Hasta pronto!