From c32d40636d65477355220aeb94f8044dd285b263 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Du=C5=A1an=20Borov=C4=8Danin?=
 <dusan.borovcanin@abstractmachines.fr>
Date: Wed, 25 Dec 2024 23:44:57 +0100
Subject: [PATCH] NOISSUE - Add token type check (#2621)

Signed-off-by: Dusan Borovcanin <borovcanindusan1@gmail.com>
---
 auth/jwt/token_test.go | 23 +++++++++++++++++++++++
 auth/jwt/tokenizer.go  |  9 ++++++++-
 auth/keys.go           |  4 ++++
 3 files changed, 35 insertions(+), 1 deletion(-)

diff --git a/auth/jwt/token_test.go b/auth/jwt/token_test.go
index 7e83cb3371..e150a3ac3d 100644
--- a/auth/jwt/token_test.go
+++ b/auth/jwt/token_test.go
@@ -104,6 +104,18 @@ func TestIssue(t *testing.T) {
 			},
 			err: nil,
 		},
+		{
+			desc: "issue token without type",
+			key: auth.Key{
+				ID:       testsutil.GenerateUUID(t),
+				Type:     auth.KeyType(auth.InvitationKey + 1),
+				Subject:  testsutil.GenerateUUID(t),
+				User:     testsutil.GenerateUUID(t),
+				Domain:   testsutil.GenerateUUID(t),
+				IssuedAt: time.Now().Add(-10 * time.Second).Round(time.Second),
+			},
+			err: nil,
+		},
 		{
 			desc: "issue token without a domain and subject",
 			key: auth.Key{
@@ -156,6 +168,11 @@ func TestParse(t *testing.T) {
 	emptySubjectToken, err := tokenizer.Issue(emptySubjectKey)
 	require.Nil(t, err, fmt.Sprintf("issuing user key expected to succeed: %s", err))
 
+	emptyTypeKey := key()
+	emptyTypeKey.Type = auth.KeyType(auth.InvitationKey + 1)
+	emptyTypeToken, err := tokenizer.Issue(emptyTypeKey)
+	require.Nil(t, err, fmt.Sprintf("issuing user key expected to succeed: %s", err))
+
 	emptyKey := key()
 	emptyKey.Domain = ""
 	emptyKey.Subject = ""
@@ -218,6 +235,12 @@ func TestParse(t *testing.T) {
 			token: emptySubjectToken,
 			err:   nil,
 		},
+		{
+			desc:  "parse token with empty type",
+			key:   emptyTypeKey,
+			token: emptyTypeToken,
+			err:   errors.ErrAuthentication,
+		},
 		{
 			desc:  "parse token with empty domain and subject",
 			key:   emptyKey,
diff --git a/auth/jwt/tokenizer.go b/auth/jwt/tokenizer.go
index d035f70630..3c2007ff22 100644
--- a/auth/jwt/tokenizer.go
+++ b/auth/jwt/tokenizer.go
@@ -17,7 +17,10 @@ import (
 )
 
 var (
+	// errInvalidIssuer is returned when the issuer is not supermq.auth.
 	errInvalidIssuer = errors.New("invalid token issuer value")
+	// errInvalidType is returned when there is no type field.
+	errInvalidType = errors.New("invalid token type")
 	// errJWTExpiryKey is used to check if the token is expired.
 	errJWTExpiryKey = errors.New(`"exp" not satisfied`)
 	// ErrSignJWT indicates an error in signing jwt token.
@@ -127,12 +130,16 @@ func toKey(tkn jwt.Token) (auth.Key, error) {
 
 	tType, ok := tkn.Get(tokenType)
 	if !ok {
-		return auth.Key{}, err
+		return auth.Key{}, errInvalidType
 	}
 	ktype, err := strconv.ParseInt(fmt.Sprintf("%v", tType), 10, 64)
 	if err != nil {
 		return auth.Key{}, err
 	}
+	kt := auth.KeyType(ktype)
+	if !kt.Validate() {
+		return auth.Key{}, errInvalidType
+	}
 
 	key.ID = tkn.JwtID()
 	key.Type = auth.KeyType(ktype)
diff --git a/auth/keys.go b/auth/keys.go
index e273119014..8fcc42039d 100644
--- a/auth/keys.go
+++ b/auth/keys.go
@@ -36,6 +36,10 @@ const (
 	InvitationKey
 )
 
+func (kt KeyType) Validate() bool {
+	return AccessKey <= kt && kt <= InvitationKey
+}
+
 func (kt KeyType) String() string {
 	switch kt {
 	case AccessKey: