ctrl+c to stop test server, fix bug where not reading jwt count correctly

Author: dominicriordan

README.md

@@ -87,13 +87,13 @@ Flags:
   -H, --headers strings          headers to send in request, can have multiple i.e -H 'content-type:application/json' -H' connection:close'
   -h, --help                     help for run
       --jwt-aud string           JWT audience (aud) claim
+      --jwt-claims string        JWT custom claims
       --jwt-header string        JWT header field name
       --jwt-iss string           JWT issuer (iss) claim
       --jwt-key string           JWT signing private key path
       --jwt-kid string           JWT KID
       --jwt-sub string           JWT subject (sub) claim
-      --jwt-claims string        JWT custom claims as a JSON string, ex: {"iat": 1719410063, "browser": "chrome"}
-  -f, --jwts-filename string     File path & name where the JWTs to use are stored
+  -f, --jwts-filename string     File path for pre-generated JWTs, separated by new lines
   -m, --method string            request method (default "GET")
       --mtls-cert string         mTLS cert path
       --mtls-key string          mTLS cert private key path
@@ -222,6 +222,13 @@ https://github.com/domsolutions/gopayloader
 +-----------------------+-------------------------------+
 ```
 
+If you have your own JWTs you want to test, you can supply a file to send the JWTs i.e. `./my-jwts.txt` where each jwt is separated by a new line.
+
+```shell
+./gopayloader run http://localhost:8081 -c 1 -r 1000000 --jwt-header "my-jwt" -f ./my-jwts.txt
+```
+
+
 To remove all generated jwts;
 
 ```shell

cmd/payloader/test-server.go

@@ -11,9 +11,11 @@ import (
 	"log"
 	"net/http"
 	"os"
+	"os/signal"
 	"path/filepath"
 	"strconv"
 	"strings"
+	"syscall"
 	"time"
 )
 
@@ -87,9 +89,25 @@ var runServerCmd = &cobra.Command{
 				},
 			}
 
-			if err := server.ListenAndServe(addr); err != nil {
-				return err
+			errs := make(chan error)
+			go func() {
+				if err := server.ListenAndServe(addr); err != nil {
+					log.Println(err)
+					errs <- err
+				}
+			}()
+
+			c := make(chan os.Signal, 1)
+			signal.Notify(c, os.Interrupt, syscall.SIGTERM)
+
+			select {
+			case <-c:
+				log.Println("User cancelled, shutting down")
+				server.Shutdown()
+			case err := <-errs:
+				log.Printf("Got error from server; %v \n", err)
 			}
+
 			return nil
 		}
 

pkgs/jwt-generator/cache.go

@@ -4,6 +4,7 @@ import (
 	"bufio"
 	"errors"
 	"fmt"
+	"github.com/pterm/pterm"
 	"os"
 	"strconv"
 	"strings"
@@ -31,9 +32,9 @@ func newCache(f *os.File) (*cache, error) {
 			return nil, err
 		}
 
-		s := string(bb)
-		count, err := strconv.ParseInt(s, 10, 64)
+		count, err := getCount(bb)
 		if err != nil {
+			pterm.Error.Printf("Got error reading jwt count from cache; %v", err)
 			return nil, err
 		}
 
@@ -47,6 +48,23 @@ func (c *cache) getJwtCount() int64 {
 	return c.count
 }
 
+func getCount(bb []byte) (int64, error) {
+	num := make([]byte, 0)
+	for _, m := range bb {
+		if m == 0 {
+			break
+		}
+		num = append(num, m)
+	}
+
+	s := string(num)
+	i, err := strconv.ParseInt(s, 10, 64)
+	if err != nil {
+		return 0, err
+	}
+	return i, nil
+}
+
 func (c *cache) get(count int64) (<-chan string, <-chan error) {
 	recv := make(chan string, 1000000)
 	errs := make(chan error, 1)
@@ -78,16 +96,7 @@ func (c *cache) get(count int64) (<-chan string, <-chan error) {
 		return recv, errs
 	}
 
-	num := make([]byte, 0)
-	for _, m := range meta {
-		if m == 0 {
-			break
-		}
-		num = append(num, m)
-	}
-
-	s := string(num)
-	i, err := strconv.ParseInt(s, 10, 64)
+	i, err := getCount(meta)
 	if err != nil {
 		errs <- fmt.Errorf("failed to get jwt count; %v", err)
 		close(errs)

pkgs/jwt-generator/jwt.go

@@ -141,7 +141,7 @@ func (j *JWTGenerator) Generate(reqJwtCount int64, dir string, retrying bool) er
 			return err
 		}
 		f.Close()
-		pterm.Debug.Printf("jwt cache %s file corrupt, attempting to delete and recreate; got error; %v \n", fname, err)
+		pterm.Error.Printf("jwt cache %s file corrupt, attempting to delete and recreate; got error; %v \n", fname, err)
 		if err := os.Remove(fname); err != nil {
 			pterm.Error.Printf("Couldn't remove cache file %s; %v", fname, err)
 			return err

version/version.go

@@ -1,3 +1,3 @@
 package version
 
-const Version = "0.3.2"
+const Version = "0.3.3"