From d2f64e88db70e4edd5bc89918d25d73ba14d4b1e Mon Sep 17 00:00:00 2001
From: Kevin Wooten <kevin@outfoxx.io>
Date: Thu, 16 Nov 2023 17:17:42 -0700
Subject: [PATCH] Add stress test

---
 .gitignore                  |   4 ++
 Makefile                    |   4 ++
 test/Stress-Dockerfile      |  27 +++++++++
 test/stress-test.sh         | 106 ++++++++++++++++++++++++++++++++++++
 test/testdata/claims1.json  |   5 ++
 test/testdata/claims10.json |   5 ++
 test/testdata/claims2.json  |   5 ++
 test/testdata/claims3.json  |   5 ++
 test/testdata/claims4.json  |   5 ++
 test/testdata/claims5.json  |   5 ++
 test/testdata/claims6.json  |   5 ++
 test/testdata/claims7.json  |   5 ++
 test/testdata/claims8.json  |   5 ++
 test/testdata/claims9.json  |   5 ++
 14 files changed, 191 insertions(+)
 create mode 100644 test/Stress-Dockerfile
 create mode 100755 test/stress-test.sh
 create mode 100644 test/testdata/claims1.json
 create mode 100644 test/testdata/claims10.json
 create mode 100644 test/testdata/claims2.json
 create mode 100644 test/testdata/claims3.json
 create mode 100644 test/testdata/claims4.json
 create mode 100644 test/testdata/claims5.json
 create mode 100644 test/testdata/claims6.json
 create mode 100644 test/testdata/claims7.json
 create mode 100644 test/testdata/claims8.json
 create mode 100644 test/testdata/claims9.json

diff --git a/.gitignore b/.gitignore
index fccedb1..7f7834d 100644
--- a/.gitignore
+++ b/.gitignore
@@ -1,3 +1,7 @@
 /vault-plugin-secrets-jwt
 
 dist/
+
+.DS_Store
+.idea
+
diff --git a/Makefile b/Makefile
index a02cd8d..4cf717f 100644
--- a/Makefile
+++ b/Makefile
@@ -34,6 +34,10 @@ endif
 functional:
 	@docker build --no-cache -f test/Dockerfile -t vault-jwt-e2e-test .
 
+# stress runs an end-to-end stress test in docker.
+stress:
+	@docker build --no-cache -f test/Stress-Dockerfile -t vault-jwt-e2e-test .
+
 # fmt formats the files according to go recommended style
 fmt:
 	@gofmt -w $(GOFMT_FILES)
diff --git a/test/Stress-Dockerfile b/test/Stress-Dockerfile
new file mode 100644
index 0000000..f1de019
--- /dev/null
+++ b/test/Stress-Dockerfile
@@ -0,0 +1,27 @@
+# Install vault
+FROM alpine as vault-installer
+WORKDIR /vault
+RUN wget https://releases.hashicorp.com/vault/1.15.2/vault_1.15.2_linux_amd64.zip -O vault.zip
+RUN unzip vault.zip && chmod +x vault
+
+# Build the addon and the test helper
+FROM golang:1.19-alpine as plugin-builder
+COPY go.mod go.sum ${GOPATH}/src/github.com/outfoxx/vault-plugin-secrets-jwt/
+COPY cmd/vault-plugin-secrets-jwt/main.go ${GOPATH}/src/github.com/outfoxx/vault-plugin-secrets-jwt/cmd/vault-plugin-secrets-jwt/
+COPY plugin/ ${GOPATH}/src/github.com/outfoxx/vault-plugin-secrets-jwt/plugin/
+COPY test/jwtverify/jwtverify.go ${GOPATH}/src/github.com/outfoxx/vault-plugin-secrets-jwt/test/
+WORKDIR ${GOPATH}/src/github.com/outfoxx/vault-plugin-secrets-jwt
+RUN go build -o /vault/plugins/vault-plugin-secrets-jwt cmd/vault-plugin-secrets-jwt/main.go
+RUN go install test/jwtverify.go
+
+# Test environment
+FROM alpine
+RUN apk add bash jq
+COPY --from=vault-installer /vault /usr/local/bin/
+COPY test/config.hcl /vault/
+COPY test/testdata/* test/stress-test.sh /test/
+COPY --from=plugin-builder /vault/plugins /vault/plugins/
+COPY --from=plugin-builder /go/bin/jwtverify /usr/local/bin/
+
+WORKDIR /test
+RUN chmod +x /test/stress-test.sh
diff --git a/test/stress-test.sh b/test/stress-test.sh
new file mode 100755
index 0000000..1e5c5d6
--- /dev/null
+++ b/test/stress-test.sh
@@ -0,0 +1,106 @@
+#!/bin/bash
+
+# Configure vault
+vault server -dev -dev-root-token-id="root" -config=/vault/config.hcl &
+VAULT_PROC=$!
+
+export VAULT_ADDR='http://127.0.0.1:8200'
+
+pid=$$
+
+fail() {
+  pkill -P $pid
+}
+
+expect_equal() {
+    # Usage: expect_equal op1 op2 message
+    if [[ ! "$1" = "$2" ]]; then
+        echo "$3: $1 != $2"
+        fail
+    fi
+}
+
+expect_match() {
+    # Usage: expect_match str pattern message
+    if [[ ! $1 =~ $2 ]]; then
+        echo "$3: $1 does not match $2"
+        fail
+    fi
+}
+
+SHASUM=$(sha256sum "/vault/plugins/vault-plugin-secrets-jwt" | cut -d " " -f1)
+
+vault login root
+
+set -e
+
+echo -e "\n### Register plugin"
+vault plugin register -sha256 $SHASUM vault-plugin-secrets-jwt
+
+echo -e "\n### Enable JWT engine at /jwt1 path"
+vault secrets enable -path=jwt1 vault-plugin-secrets-jwt
+
+echo -e "\n### Change the expiry time and make a pattern to check subjects against"
+vault write jwt1/config "sig_alg=RS256" "key_ttl=3s" "jwt_ttl=40s"
+
+echo -e "\n### Enable JWT engine at /jwt2 path"
+vault secrets enable -path=jwt2 vault-plugin-secrets-jwt
+
+echo -e "\n### Change the expiry time and make a pattern to check subjects against"
+vault write jwt2/config "sig_alg=RS256" "key_ttl=3s" "jwt_ttl=40s"
+
+stress() {
+
+  echo -e "### [${1}] Adding role test${1}"
+  if ! vault write jwt${2}/roles/test${1} issuer="DOOP"; then
+    echo "Failed to add role"
+    fail
+  fi
+
+  expected_sub=$(cat claims${3}.json | jq -r '.claims.sub')
+
+  for i in {1..1000}; do
+    echo -e "### [${1}] <${i}> Generating a token"
+    if ! vault write -field=token jwt${2}/sign/test${1} @claims${3}.json > jwt-${1}-${i}.txt; then
+      echo -e "##############################################"
+      echo -e "### [${1}] <${i}> Failed to generate token ###"
+      echo -e "##############################################"
+      fail
+    fi
+
+    START_TIME="$(date -u +%s)"
+    echo -e "### [${1}] <${i}> Validating 100 times"
+    for j in {1..100}; do
+#      echo -e "### [${1}] <${i}:${j}> Verify that the token is formatted as expected"
+      if ! jwtverify "$(cat jwt-${1}-${i}.txt)" $VAULT_ADDR/v1/jwt${2}/jwks > decoded-${1}-${i}-${j}.txt; then
+        echo -e "### [${1}] <${i}:${j}> Failed to verify token"
+        fail
+      fi
+
+      expect_equal "$(cat decoded-${1}-${i}-${j}.txt | jq -r '.sub')" "${expected_sub}" "Wrong subject"
+      expect_match "$(cat decoded-${1}-${i}-${j}.txt | jq '.exp')" "[0-9]+" "Invalid 'exp' claim"
+      expect_match "$(cat decoded-${1}-${i}-${j}.txt | jq '.iat')" "[0-9]+" "Invalid 'iat' claim"
+      expect_match "$(cat decoded-${1}-${i}-${j}.txt | jq '.nbf')" "[0-9]+" "Invalid 'nbf' claim"
+    done
+    END_TIME="$(date -u +%s)"
+
+    ELAPSED_TIME="$(($END_TIME-$START_TIME))"
+    if [[ $ELAPSED_TIME -gt 30 ]]; then
+      echo -e "############################################################"
+      echo -e "### [${1}] <${i}> Elapsed time: ${ELAPSED_TIME} seconds"
+      echo -e "############################################################"
+      fail
+    fi
+
+  done
+}
+
+for i in {1..10}; do
+  stress $i "1" $i &
+  sleep 1
+done
+
+for i in {1..10}; do
+  stress $((i+10)) "2" $i &
+  sleep 1
+done
diff --git a/test/testdata/claims1.json b/test/testdata/claims1.json
new file mode 100644
index 0000000..a22a117
--- /dev/null
+++ b/test/testdata/claims1.json
@@ -0,0 +1,5 @@
+{
+    "claims": {
+        "sub": "Zapp Brannigan"
+    }
+}
\ No newline at end of file
diff --git a/test/testdata/claims10.json b/test/testdata/claims10.json
new file mode 100644
index 0000000..c1676a2
--- /dev/null
+++ b/test/testdata/claims10.json
@@ -0,0 +1,5 @@
+{
+    "claims": {
+        "sub": "Scruffy Scruffington"
+    }
+}
diff --git a/test/testdata/claims2.json b/test/testdata/claims2.json
new file mode 100644
index 0000000..7d78394
--- /dev/null
+++ b/test/testdata/claims2.json
@@ -0,0 +1,5 @@
+{
+    "claims": {
+        "sub": "Kif Kroker"
+    }
+}
diff --git a/test/testdata/claims3.json b/test/testdata/claims3.json
new file mode 100644
index 0000000..69f8e48
--- /dev/null
+++ b/test/testdata/claims3.json
@@ -0,0 +1,5 @@
+{
+    "claims": {
+        "sub": "Philip J. Fry"
+    }
+}
diff --git a/test/testdata/claims4.json b/test/testdata/claims4.json
new file mode 100644
index 0000000..c345fb2
--- /dev/null
+++ b/test/testdata/claims4.json
@@ -0,0 +1,5 @@
+{
+    "claims": {
+        "sub": "Turanga Leela"
+    }
+}
diff --git a/test/testdata/claims5.json b/test/testdata/claims5.json
new file mode 100644
index 0000000..c7197bf
--- /dev/null
+++ b/test/testdata/claims5.json
@@ -0,0 +1,5 @@
+{
+    "claims": {
+        "sub": "Bender Bending Rodriguez"
+    }
+}
diff --git a/test/testdata/claims6.json b/test/testdata/claims6.json
new file mode 100644
index 0000000..f137fbe
--- /dev/null
+++ b/test/testdata/claims6.json
@@ -0,0 +1,5 @@
+{
+    "claims": {
+        "sub": "Professor Farnsworth"
+    }
+}
diff --git a/test/testdata/claims7.json b/test/testdata/claims7.json
new file mode 100644
index 0000000..f88a553
--- /dev/null
+++ b/test/testdata/claims7.json
@@ -0,0 +1,5 @@
+{
+    "claims": {
+        "sub": "Amy Wong"
+    }
+}
diff --git a/test/testdata/claims8.json b/test/testdata/claims8.json
new file mode 100644
index 0000000..bf7ea91
--- /dev/null
+++ b/test/testdata/claims8.json
@@ -0,0 +1,5 @@
+{
+    "claims": {
+        "sub": "Hermes Conrad"
+    }
+}
diff --git a/test/testdata/claims9.json b/test/testdata/claims9.json
new file mode 100644
index 0000000..14d0293
--- /dev/null
+++ b/test/testdata/claims9.json
@@ -0,0 +1,5 @@
+{
+    "claims": {
+        "sub": "Doctor Zoidberg"
+    }
+}