From a1fa3bfe518cc53d5e95407f9d8ce9293dcbd594 Mon Sep 17 00:00:00 2001
From: ArkaSaha30 <arkasaha30@gmail.com>
Date: Thu, 6 Apr 2023 15:36:07 +0530
Subject: [PATCH] Add test cases for malformed jwt fix

Signed-off-by: ArkaSaha30 <arkasaha30@gmail.com>
---
 server/auth/jwt_test.go | 75 +++++++++++++++++++++++++++++++++++++++++
 1 file changed, 75 insertions(+)

diff --git a/server/auth/jwt_test.go b/server/auth/jwt_test.go
index a3983cc5a..e7c4840b4 100644
--- a/server/auth/jwt_test.go
+++ b/server/auth/jwt_test.go
@@ -18,7 +18,10 @@ import (
 	"context"
 	"fmt"
 	"testing"
+	"time"
 
+	"github.com/golang-jwt/jwt/v4"
+	"github.com/stretchr/testify/require"
 	"go.uber.org/zap"
 )
 
@@ -132,6 +135,78 @@ func testJWTInfo(t *testing.T, opts map[string]string) {
 	}
 }
 
+func TestJWTTokenWithMissingFields(t *testing.T) {
+	testCases := []struct {
+		name        string
+		username    string // An empty string means not present
+		revision    uint64 // 0 means not present
+		expectValid bool
+	}{
+		{
+			name:        "valid token",
+			username:    "hello",
+			revision:    100,
+			expectValid: true,
+		},
+		{
+			name:        "no username",
+			username:    "",
+			revision:    100,
+			expectValid: false,
+		},
+		{
+			name:        "no revision",
+			username:    "hello",
+			revision:    0,
+			expectValid: false,
+		},
+	}
+
+	for _, tc := range testCases {
+		tc := tc
+		optsMap := map[string]string{
+			"priv-key":    jwtRSAPrivKey,
+			"sign-method": "RS256",
+			"ttl":         "1h",
+		}
+
+		t.Run(tc.name, func(t *testing.T) {
+			// prepare claims
+			claims := jwt.MapClaims{
+				"exp": time.Now().Add(time.Hour).Unix(),
+			}
+			if tc.username != "" {
+				claims["username"] = tc.username
+			}
+			if tc.revision != 0 {
+				claims["revision"] = tc.revision
+			}
+
+			// generate a JWT token with the given claims
+			var opts jwtOptions
+			err := opts.ParseWithDefaults(optsMap)
+			require.NoError(t, err)
+			key, err := opts.Key()
+			require.NoError(t, err)
+
+			tk := jwt.NewWithClaims(opts.SignMethod, claims)
+			token, err := tk.SignedString(key)
+			require.NoError(t, err)
+
+			// verify the token
+			jwtProvider, err := newTokenProviderJWT(zap.NewNop(), optsMap)
+			require.NoError(t, err)
+			ai, ok := jwtProvider.info(context.TODO(), token, 123)
+
+			require.Equal(t, tc.expectValid, ok)
+			if ok {
+				require.Equal(t, tc.username, ai.Username)
+				require.Equal(t, tc.revision, ai.Revision)
+			}
+		})
+	}
+}
+
 func TestJWTBad(t *testing.T) {
 
 	var badCases = map[string]map[string]string{