From 4478d638add24e642fad2ff203668bc603818a2e Mon Sep 17 00:00:00 2001
From: Roman Loginov <r.loginov@yadro.com>
Date: Fri, 29 Nov 2024 09:48:44 +0300
Subject: [PATCH] [#163] Support JSON bearer token

Signed-off-by: Roman Loginov <r.loginov@yadro.com>
---
 tokens/bearer-token.go      | 18 ++++++++++----
 tokens/bearer-token_test.go | 48 +++++++++++++++++++++++++++++--------
 2 files changed, 51 insertions(+), 15 deletions(-)

diff --git a/tokens/bearer-token.go b/tokens/bearer-token.go
index 880a100..24ffcbe 100644
--- a/tokens/bearer-token.go
+++ b/tokens/bearer-token.go
@@ -82,14 +82,22 @@ func fetchBearerToken(ctx *fasthttp.RequestCtx) (*bearer.Token, error) {
 		tkn = new(bearer.Token)
 	)
 	for _, parse := range []fromHandler{BearerTokenFromHeader, BearerTokenFromCookie} {
-		if buf = parse(&ctx.Request.Header); buf == nil {
+		buf = parse(&ctx.Request.Header)
+		if buf == nil {
 			continue
-		} else if data, err := base64.StdEncoding.DecodeString(string(buf)); err != nil {
+		}
+
+		data, err := base64.StdEncoding.DecodeString(string(buf))
+		if err != nil {
 			lastErr = fmt.Errorf("can't base64-decode bearer token: %w", err)
 			continue
-		} else if err = tkn.Unmarshal(data); err != nil {
-			lastErr = fmt.Errorf("can't unmarshal bearer token: %w", err)
-			continue
+		}
+
+		if err = tkn.Unmarshal(data); err != nil {
+			if err = tkn.UnmarshalJSON(data); err != nil {
+				lastErr = fmt.Errorf("can't unmarshal bearer token: %w", err)
+				continue
+			}
 		}
 
 		return tkn, nil
diff --git a/tokens/bearer-token_test.go b/tokens/bearer-token_test.go
index 6fb3bf4..60e9ea2 100644
--- a/tokens/bearer-token_test.go
+++ b/tokens/bearer-token_test.go
@@ -98,8 +98,14 @@ func TestFetchBearerToken(t *testing.T) {
 	tkn := new(bearer.Token)
 	tkn.ForUser(uid)
 
-	t64 := base64.StdEncoding.EncodeToString(tkn.Marshal())
-	require.NotEmpty(t, t64)
+	jsonToken, err := tkn.MarshalJSON()
+	require.NoError(t, err)
+
+	jsonTokenBase64 := base64.StdEncoding.EncodeToString(jsonToken)
+	binaryTokenBase64 := base64.StdEncoding.EncodeToString(tkn.Marshal())
+
+	require.NotEmpty(t, jsonTokenBase64)
+	require.NotEmpty(t, binaryTokenBase64)
 
 	cases := []struct {
 		name   string
@@ -143,25 +149,47 @@ func TestFetchBearerToken(t *testing.T) {
 			error:  "can't unmarshal bearer token",
 		},
 		{
-			name:   "bad header, but good cookie",
+			name:   "bad header, but good cookie with binary token",
 			header: "dGVzdAo=",
-			cookie: t64,
+			cookie: binaryTokenBase64,
 			expect: tkn,
 		},
 		{
-			name:   "bad cookie, but good header",
-			header: t64,
+			name:   "bad cookie, but good header with binary token",
+			header: binaryTokenBase64,
 			cookie: "dGVzdAo=",
 			expect: tkn,
 		},
 		{
-			name:   "ok for header",
-			header: t64,
+			name:   "bad header, but good cookie with json token",
+			header: "dGVzdAo=",
+			cookie: jsonTokenBase64,
 			expect: tkn,
 		},
 		{
-			name:   "ok for cookie",
-			cookie: t64,
+			name:   "bad cookie, but good header with json token",
+			header: jsonTokenBase64,
+			cookie: "dGVzdAo=",
+			expect: tkn,
+		},
+		{
+			name:   "ok for header with binary token",
+			header: binaryTokenBase64,
+			expect: tkn,
+		},
+		{
+			name:   "ok for cookie with binary token",
+			cookie: binaryTokenBase64,
+			expect: tkn,
+		},
+		{
+			name:   "ok for header with json token",
+			header: jsonTokenBase64,
+			expect: tkn,
+		},
+		{
+			name:   "ok for cookie with json token",
+			cookie: jsonTokenBase64,
 			expect: tkn,
 		},
 	}