bug: rework Tokenserver load tests for local OAuth verification

.gitignore

@@ -32,6 +32,7 @@ service-account.json
 .sentryclirc
 
 config/local.toml
+tools/tokenserver/loadtests/*.pem
+tools/tokenserver/loadtests/*.pub
 venv
-mozilla-rust-sdk
 .vscode/settings.json

tools/tokenserver/loadtests/README.md

@@ -3,42 +3,85 @@
 This directory contains everything needed to run the suite of load tests for Tokenserver.
 
 ## Building and Running
+
 1. Install the load testing dependencies:
-```sh
-pip3 install -r requirements.txt
-```
-2. Set up a mock OAuth verifier, with which Tokenserver will verify OAuth tokens. The subdirectory [mock-fxa-server/](./mock-fxa-server) includes code deployable as a GCP Cloud Function that acts as a mock FxA server, "verifying" OAuth tokens. You can deploy your own Cloud Function by running the following command in this directory:
-```sh
-gcloud functions deploy mock_fxa_server --runtime=python39 --trigger-http --source=mock-fxa-server
-```
-You can stand up a local copy of the Cloud Function by running the following in this directory:
-```sh
-functions-framework --target mock_fxa_server --debug
-```
-Note that you'll need to install `functions-framework` via `pip3 install functions-framework`. The load tests will use FxA stage to verify BrowserID assertions.
-3. Configure Tokenserver to verify OAuth tokens through the mock FxA service and BrowserID assertions through FxA stage. This is done by setting the following environment variables:
-```
-SYNC_TOKENSERVER__FXA_BROWSERID_AUDIENCE=https://token.stage.mozaws.net
-SYNC_TOKENSERVER__FXA_BROWSERID_ISSUER=api-accounts.stage.mozaws.net
-SYNC_TOKENSERVER__BROWSERID_VERIFIER_URL=https://verifier.stage.mozaws.net/v2
 
-# This variable should be set to point to the host and port of the moack OAuth verifier created in step 2
-SYNC_TOKENSERVER__FXA_OAUTH_SERVER_URL=http://localhost:6000
-```
-4. Tokenserver uses [locust](https://locust.io/) for load testing. To run the load tests, simply run the following command in this directory:
-```sh
-locust
-```
-Next, navigate your browser to http://localhost:8090, where you'll find the locust GUI. Enter the following information:
-* Number of users: The peak number of simultaneous connections to Tokenserver
-* Spawn rate: The rate at which new connections are created
-* Host: The URL of the server to be load tested. Note that this URL must include the protocol (e.g. "http://").
+   ```sh
+   pip3 install -r requirements.txt
+   ```
+
+1. Run the `generate-keys.sh` script to generate an RSA keypair and derive the public JWK:
+
+   ```sh
+   ./generate-keys.sh
+   ```
+
+   This script will output two files:
+
+   - `load_test.pem`: The private key to be used by the load tests to create OAuth tokens
+   - `jwk.json`: The public JWK associated with the private key. This is a key of the form
+
+     ```json
+     {
+        "n": ...,
+        "e": ...,
+        "kty": "RSA"
+     }
+     ```
 
-Click the "Start swarming" button to begin the load tests.
+1. Set the following environment variables/settings on Tokenserver:
+
+   ```sh
+   # Should be set to the "n" component of the JWK
+   SYNC_TOKENSERVER__FXA_OAUTH_PRIMARY_JWK_N
+   # Should be set to the "e" component of the JWK (this value should almost always be "AQAB")
+   SYNC_TOKENSERVER__FXA_OAUTH_PRIMARY_JWK_E
+   SYNC_TOKENSERVER__FXA_OAUTH_PRIMARY_JWK_KTY=RSA
+   SYNC_TOKENSERVER__FXA_OAUTH_PRIMARY_JWK_USE=sig
+   SYNC_TOKENSERVER__FXA_OAUTH_PRIMARY_JWK_ALG=RS256
+
+   # These two environment variables don't affect the load tests, but they need to be set:
+   SYNC_TOKENSERVER__FXA_OAUTH_PRIMARY_JWK_KID=""
+   SYNC_TOKENSERVER__FXA_OAUTH_PRIMARY_JWK_FXA_CREATED_AT=0
+   ```
+
+1. Configure Tokenserver to verify BrowserID assertions through FxA stage. This is done by setting the following environment variables:
+
+   ```sh
+   # The exact value of this environment variable is not important as long as it matches the `BROWSERID_AUDIENCE` environment variable set on the machine running the load tests, as described below
+   SYNC_TOKENSERVER__FXA_BROWSERID_SERVER_URL=https://verifier.stage.mozaws.net/v2
+
+   SYNC_TOKENSERVER__FXA_BROWSERID_AUDIENCE=https://token.stage.mozaws.net
+   SYNC_TOKENSERVER__FXA_BROWSERID_ISSUER=mockmyid.s3-us-west-2.amazonaws.com
+   ```
+
+   Note that, because we have cached the JWK used to verify OAuth tokens, no verification requests will be made to FxA, so the value of `SYNC_TOKENSERVER__FXA_OAUTH_VERIFIER_URL` does not matter; however, Tokenserver expects it to be set, so setting it to something like `http://localhost` will suffice.
+
+1. Set the following environment variables on the machine that will be running the load tests:
+
+   - `OAUTH_PEM_FILE` should be set to the location of the private RSA key generated in a previous step
+   - `BROWSERID_AUDIENCE` should be set to match the `SYNC_TOKENSERVER__FXA_BROWSERID_AUDIENCE` environment variable on Tokenserver
+
+1. Tokenserver uses [locust](https://locust.io/) for load testing. To run the load tests, simply run the following command in this directory:
+
+   ```sh
+   locust
+   ```
+
+1. Navigate your browser to <http://localhost:8090>, where you'll find the locust GUI. Enter the following information:
+
+   - Number of users: The peak number of Tokenserver users to be used during the load tests
+   - Spawn rate: The rate at which new users are spawned
+   - Host: The URL of the server to be load tested. Note that this URL must include the protocol (e.g. "http://")
+
+1. Click the "Start swarming" button to begin the load tests.
 
 ## Populating the Database
+
 This directory includes an optional `populate_db.py` script that can be used to add test users to the database en masse. The script can be run like so:
+
 ```sh
 python3 populate_db.py <sqlurl> <nodes> <number of users>
 ```
-where `sqluri` is the URL of the Tokenserver database, `nodes` is a comma-separated list of nodes **that are already present in the database** to which the users will be randomly assigned, and `number of users` is the number of users to be created. 
+
+where `sqluri` is the URL of the Tokenserver database, `nodes` is a comma-separated list of nodes **that are already present in the database** to which the users will be randomly assigned, and `number of users` is the number of users to be created.

tools/tokenserver/loadtests/generate-keys.sh

@@ -0,0 +1,11 @@
+#!/bin/sh
+
+# Generate a private RSA key
+openssl genrsa -out load_test.pem 2048
+
+# Derive the public key from the private key
+openssl rsa -in load_test.pem -pubout > load_test.pub
+
+# Derive and print the JWK from the public key
+python3 get_jwk.py load_test.pub > jwk.json
+rm load_test.pub

tools/tokenserver/loadtests/get_jwk.py

@@ -0,0 +1,6 @@
+import sys
+from authlib.jose import JsonWebKey
+
+raw_public_key = open(sys.argv[1], "rb").read()
+public_key = JsonWebKey.import_key(raw_public_key, {"kty":  "RSA"})
+print(public_key.as_json())

tools/tokenserver/loadtests/locustfile.py

@@ -1,14 +1,27 @@
 from base64 import urlsafe_b64encode as b64encode
 import binascii
-import json
+import jwt
+import os
 import time
 
 import browserid
 import browserid.jwt
 from browserid.tests.support import make_assertion
+from cryptography.hazmat.primitives import serialization
+from cryptography.hazmat.primitives.asymmetric import rsa
 from locust import HttpUser, task, between
 
+BROWSERID_AUDIENCE = os.environ['BROWSERID_AUDIENCE']
 DEFAULT_OAUTH_SCOPE = 'https://identity.mozilla.com/apps/oldsync'
+
+# To create an invalid token, we sign the JWT with a private key that doesn't
+# correspond with the public key set on Tokenserver. To accomplish this, we
+# just generate a new private key with every run of the load tests.
+INVALID_OAUTH_PRIVATE_KEY = rsa.generate_private_key(
+    public_exponent=65537,
+    key_size=2048,
+)
+
 # We use a custom mockmyid site to synthesize valid assertions.
 # It's hosted in a static S3 bucket so we don't swamp the live mockmyid server.
 MOCKMYID_DOMAIN = "mockmyid.s3-us-west-2.amazonaws.com"
@@ -32,6 +45,13 @@
 ONE_YEAR = 60 * 60 * 24 * 365
 TOKENSERVER_PATH = '/1.0/sync/1.5'
 
+# This is a private key used to "forge" valid tokens. The associated public
+# key must be set using the SYNC_TOKENSERVER__FXA_PRIMARY_JWK_* environment
+# variables on Tokenserver.
+VALID_OAUTH_PRIVATE_KEY = private_key = serialization.load_pem_private_key(
+    open(os.environ['OAUTH_PEM_FILE'], "rb").read(), password=None,
+)
+
 
 class TokenserverTestUser(HttpUser):
     # An instance of this class represents a single Tokenserver user. Instances
@@ -60,15 +80,18 @@ def test_oauth_success(self):
 
     @task(100)
     def test_invalid_oauth(self):
-        token = self._make_oauth_token(status=400)
+        token = self._make_oauth_token(
+            self.email,
+            key=INVALID_OAUTH_PRIVATE_KEY
+        )
 
         self._do_token_exchange_via_oauth(token, status=401)
 
     @task(100)
     def test_invalid_oauth_scope(self):
         token = self._make_oauth_token(
-            user=str(self.fxa_uid),
-            scope=["unrelated", "scopes"],
+            self.email,
+            scope="unrelated scopes",
         )
 
         self._do_token_exchange_via_oauth(token, status=401)
@@ -132,31 +155,27 @@ def test_browserid_invalid_issuer_priv_key(self):
 
         self._do_token_exchange_via_browserid(assertion, status=401)
 
-    def _make_oauth_token(self, user=None, status=200, **fields):
+    def _make_oauth_token(self, email, key=VALID_OAUTH_PRIVATE_KEY, **fields):
         # For mock oauth tokens, we bundle the desired status code
         # and response body into a JSON blob for the mock verifier
         # to echo back to us.
         body = {}
-        if status < 400:
-            if user is None:
-                raise ValueError("Must specify user for valid oauth token")
-            if "scope" not in fields:
-                fields["scope"] = [DEFAULT_OAUTH_SCOPE]
-            if "client_id" not in fields:
-                fields["client_id"] = "x"
-        if user is not None:
-            parts = user.split("@", 1)
-            if len(parts) == 1:
-                body["user"] = user
-            else:
-                body["user"] = parts[0]
-                body["issuer"] = parts[1]
+        if "scope" not in fields:
+            fields["scope"] = DEFAULT_OAUTH_SCOPE
+        if "client_id" not in fields:
+            fields["client_id"] = "x"
+        sub, issuer = email.split("@", 1)
+        body["sub"] = sub
+        body["issuer"] = issuer
         body['fxa-generation'] = self.generation_counter
         body.update(fields)
-        return json.dumps({
-          "status": status,
-          "body": body
-        })
+
+        return jwt.encode(
+            body,
+            key,
+            algorithm="RS256",
+            headers={'typ': 'application/at+jwt'}
+        )
 
     def _make_x_key_id_header(self):
         # In practice, the generation number and keys_changed_at may not be
@@ -170,7 +189,7 @@ def _make_x_key_id_header(self):
 
     def _make_browserid_assertion(self, email, **kwds):
         if "audience" not in kwds:
-            kwds["audience"] = self.client.base_url
+            kwds["audience"] = BROWSERID_AUDIENCE
         if "exp" not in kwds:
             kwds["exp"] = int((time.time() + ONE_YEAR) * 1000)
         if "issuer" not in kwds:

tools/tokenserver/loadtests/requirements.txt

@@ -1,3 +1,6 @@
+authlib
+cryptography
+jwt
 locust
 pybrowserid
 sqlalchemy