How to Use K6 for Load Testing IAP Protected Sites?

When developing sites and APIs, we want to know what load they can handle. For this measurement, we are using K6. In this blog, I will show you how you can load test your apps protected by Identity Aware Proxy (IAP).

Site protection

Usually, you don't want to do load testing on your production infrastructure, but also you want something close to it. We at Ackee have this staging environment with the same or soon-to-be-released version of frontend and backend APIs. This environment shouldn't be public as there can still be some errors or shortcuts for easier testing. There are several options which we considered or used to do this.

1) Local server

It is one of the most straightforward solutions (at least on first look). You don't need to worry too much about limiting access to the app. However, it becomes more problematic as soon as you want to expose the app to other people outside your local network (e.g. developers working from home or external partners). Also, this doesn't work if you are hosting your whole infrastructure in the cloud (as we do).

2) Cloudflare Access

Since we are hosting our apps in the cloud and using Cloudflare for DNS, this was a good solution. Set it up, add your office IP and partners' office IPs to the allowlist, and you can still have 30 free users outside this range. Perfect, no additional cost, right?

Yes and no. It was ideal for our K6 load tests; the runner was located in the office, so it could access the stage environment without any additional config, and it just worked. But when we began to have more than 30 users outside of the offices simultaneously, we had to start paying for the service (on a per user basis). So we began looking for some other solution.

3) Google Identity Aware Proxy

As I have already mentioned, our apps run in the cloud, specifically in Google Cloud. After some probing, we decided to use Google IAP, which requires all users to be logged in before they can access the resources. It was nearly perfect, the configuration was quite simple, and when using it, most of us didn't even notice it since we are using Google Accounts for everything. But there was one thing: we couldn't make any IP exceptions, so our tests had to be adjusted.

K6 and Google IAP

So now the question is how to adjust the K6 tests to authenticate against IAP and still work without user interaction.

A good starting point could be looking at the Node.js example in Google documentation. It looks good, but when you try to use it in your K6 tests, it doesn't work. The issue is that the K6 executor is not a full-fledged JS runtime and only supports a limited subset of JS APIs.

K6 extensions

Fortunately, K6 supports writing custom extensions in Go. And now, since we have access to all APIs of Go, we can use the example from the already mentioned documentation, but this time for Go.

After combining these two examples, we end up, through some slight modifications, with the following:

func init() {
      modules.Register("k6/x/googleIap", new(GoogleIap))
}

type GoogleIap struct{}

func (*GoogleIap) GetToken(audience string, saKey string) string {
      idToken := GetIdToken(audience, saKey)
      if idToken == nil {
           return ""
      }

     return idToken.AccessToken
}

func GetIdToken(audience string, credentials string) *oauth2.Token {
      ctx := context.Background()

      tokenSource, err := idtoken.NewTokenSource(ctx, audience, idtoken.WithCredentialsJSON([]byte(credentials)))
     if err != nil {
            fmt.Printf("[k6/x/googleIapToken][error]: %v\n", err.Error())
           return nil
      }

      token, err := tokenSource.Token()
      if err != nil {
            fmt.Printf("[k6/x/googleIapToken][error]: %v\n", err.Error()) 
           return nil
      }
  
      return token
}

Nice. This K6 extensions module has one function called getToken, which accepts two parameters:

1. Audience: which is the Client ID of your IAP
2. Service Account key: string with a JSON object containing information about the account, which K6 will use for authentication.

And now how to use it? It is quite simple:

import googleIap from 'k6/x/googleIap';
import { check } from 'k6';
import http from 'k6/http';

const client_id = __ENV.CLIENT_ID;
const sa_key = __ENV.SA_KEY;
// URL protected by Google IAP
const url = __ENV.TARGET_URL; 

export default function () {
    const token = googleIap.getToken(client_id, sa_key);

    const params = {
        headers: {
            'Proxy-Authorization': `Bearer ${token}`,
        },
    };

    const response = http.get(url, params);
    console.log(response.url)

    check(response, { "URL doesn't begin with accounts.google.com": response => !response.url.startsWith("https://accounts.google.com") })
}

It is good to point out that we need an Authorization header for our app, so we can't use it for IAP. Thankfully you can use the Proxy-Authorization header precisely in this situation.

We have the extension, but now we need to run the test with it. Unsurprisingly, it is not so difficult. You can build it like this:

$ xk6 build --with github.com/AckeeCZ/xk6-google-iap  

And then run it:

$ ./k6 -e SA_KEY=**"SERVICE_ACCOUNT_KEY"** -e CLIENT_ID=**"OAUTH_CLIENT_ID"** -e TARGET_URL=**"YOUR_WEBSITE_URL"** script.js

Bonus: Typescript hints

In our K6 tests, we mostly use Typescript, so it is quite nice to have some type hints even for this small module:

declare module "k6/x/googleIap" {
    function getToken(clientId: string, serviceAccountKey: string): string
}

(note: Don't forget to add the path to the folder with type hints to your tsconfig.json)

Conclusion

This pretty simple extension provides means for authentication against Google IAP, but it shouldn't be hard to add or replace it with other providers.

However, this extension is not by any means perfect. For example, it doesn't support refreshing tokens, so please be mindful of that.

You can find the source code used in this blog and the whole module on Github.