Repository: appleboy/gin-jwt
Pull request #335 (2024-07-14 14:34:42)

refactor: Provided an Option to Set Cookie using *GinJWTMiddleware
### Problem

1. **DRY Code:** The `SetCookie` code is redundant we’ve duplicated the logic for setting cookies in both the `LoginHandler` and `RefreshToken`.
2. We do have the `func (*GinJWTMiddleware) TokenGenerator` method for clients to generate `JWT` tokens independently of the `LoginHandler` or `RefreshToken` methods. However, without a dedicated `SetCookie` method, cookie setting remains tightly coupled with these methods. Introducing a `SetCookie` method would enhance scalability by enabling cookies to be set independently, eliminating the need for manual cookie handling where such functionality is currently absent.

**Context:** 

This library looks very useful for us. However at the moment I don't think it supports `oauth`? Basically I don't want my users to register with a username/password, but instead I'd like them to use SSO, maybe provided by Facebook, Google or our built inhouse elitmus.com. Once the oauth flow is complete, I'd like to use my own JWT tokens though. I'd like to use the user info I get from SSO to build the token.

So far I have something like this:

I am using the `TokenGenerator` method to generate the `jwt` token

```go
tokenString, expire, err := auth.TokenGenerator(&User)
```

And in our case we want to store the token in the cookie and we don't have an method to do that. As that method is tightly coupled with the `LoginHandler` method.

To make the method to set the cookie for us I have written the method in our application code.

```go
type GinJWTMiddleware struct {
	*jwt.GinJWTMiddleware
}

func NewAuthMiddleware() (*jwt.GinJWTMiddleware, error) {
  ...
}

func (mw *GinJWTMiddleware) SetCookie(c *gin.Context, tokenString string) {
	if mw.SendCookie {
		expireCookie := mw.TimeFunc().Add(mw.CookieMaxAge)
		maxage := int(expireCookie.Unix() - mw.TimeFunc().Unix())

		if mw.CookieSameSite != 0 {
			c.SetSameSite(mw.CookieSameSite)
		}

		c.SetCookie(
			mw.CookieName,
			tokenString,
			maxage,
			"/",
			mw.CookieDomain,
			mw.SecureCookie,
			mw.CookieHTTPOnly,
		)
	}
}
```

**Solution Implementation:**

1. `SetCookie` Method: This method is added to the `*GinJWTMiddleware` and is responsible for setting the cookie in the response. It takes the context, token as parameters.
2. Refactored `LoginHandler` and `RefreshToken`: Both methods now call `SetCookie` to set the cookie, eliminating the duplicated logic and centralizing cookie handling in one place.

Fixes: #336 

maniSHarma7575: @appleboy Could you please review this one?

Patch:
diff --git a/auth_jwt.go b/auth_jwt.go
index 1804494..4ca1809 100644
--- a/auth_jwt.go
+++ b/auth_jwt.go
@@ -517,25 +517,7 @@ func (mw *GinJWTMiddleware) LoginHandler(c *gin.Context) {
 		return
 	}
 
-	// set cookie
-	if mw.SendCookie {
-		expireCookie := mw.TimeFunc().Add(mw.CookieMaxAge)
-		maxage := int(expireCookie.Unix() - mw.TimeFunc().Unix())
-
-		if mw.CookieSameSite != 0 {
-			c.SetSameSite(mw.CookieSameSite)
-		}
-
-		c.SetCookie(
-			mw.CookieName,
-			tokenString,
-			maxage,
-			"/",
-			mw.CookieDomain,
-			mw.SecureCookie,
-			mw.CookieHTTPOnly,
-		)
-	}
+	mw.SetCookie(c, tokenString)
 
 	mw.LoginResponse(c, http.StatusOK, tokenString, expire)
 }
@@ -609,25 +591,7 @@ func (mw *GinJWTMiddleware) RefreshToken(c *gin.Context) (string, time.Time, err
 		return "", time.Now(), err
 	}
 
-	// set cookie
-	if mw.SendCookie {
-		expireCookie := mw.TimeFunc().Add(mw.CookieMaxAge)
-		maxage := int(expireCookie.Unix() - time.Now().Unix())
-
-		if mw.CookieSameSite != 0 {
-			c.SetSameSite(mw.CookieSameSite)
-		}
-
-		c.SetCookie(
-			mw.CookieName,
-			tokenString,
-			maxage,
-			"/",
-			mw.CookieDomain,
-			mw.SecureCookie,
-			mw.CookieHTTPOnly,
-		)
-	}
+	mw.SetCookie(c, tokenString)
 
 	return tokenString, expire, nil
 }
@@ -845,3 +809,26 @@ func GetToken(c *gin.Context) string {
 
 	return token.(string)
 }
+
+// SetCookie help to set the token in the cookie
+func (mw *GinJWTMiddleware) SetCookie(c *gin.Context, token string) {
+	// set cookie
+	if mw.SendCookie {
+		expireCookie := mw.TimeFunc().Add(mw.CookieMaxAge)
+		maxage := int(expireCookie.Unix() - mw.TimeFunc().Unix())
+
+		if mw.CookieSameSite != 0 {
+			c.SetSameSite(mw.CookieSameSite)
+		}
+
+		c.SetCookie(
+			mw.CookieName,
+			token,
+			maxage,
+			"/",
+			mw.CookieDomain,
+			mw.SecureCookie,
+			mw.CookieHTTPOnly,
+		)
+	}
+}


Test patch:
diff --git a/auth_jwt_test.go b/auth_jwt_test.go
index b239a3a..1abd564 100644
--- a/auth_jwt_test.go
+++ b/auth_jwt_test.go
@@ -5,6 +5,7 @@ import (
 	"fmt"
 	"log"
 	"net/http"
+	"net/http/httptest"
 	"os"
 	"reflect"
 	"strings"
@@ -1322,3 +1323,39 @@ func TestLogout(t *testing.T) {
 			assert.Equal(t, fmt.Sprintf("%s=; Path=/; Domain=%s; Max-Age=0", cookieName, cookieDomain), r.HeaderMap.Get("Set-Cookie"))
 		})
 }
+
+func TestSetCookie(t *testing.T) {
+	w := httptest.NewRecorder()
+	c, _ := gin.CreateTestContext(w)
+
+	mw, _ := New(&GinJWTMiddleware{
+		Realm:          "test zone",
+		Key:            key,
+		Timeout:        time.Hour,
+		Authenticator:  defaultAuthenticator,
+		SendCookie:     true,
+		CookieName:     "jwt",
+		CookieMaxAge:   time.Hour,
+		CookieDomain:   "example.com",
+		SecureCookie:   false,
+		CookieHTTPOnly: true,
+		TimeFunc: func() time.Time {
+			return time.Now()
+		},
+	})
+
+	token := makeTokenString("HS384", "admin")
+
+	mw.SetCookie(c, token)
+
+	cookies := w.Result().Cookies()
+
+	assert.Len(t, cookies, 1)
+
+	cookie := cookies[0]
+	assert.Equal(t, "jwt", cookie.Name)
+	assert.Equal(t, token, cookie.Value)
+	assert.Equal(t, "/", cookie.Path)
+	assert.Equal(t, "example.com", cookie.Domain)
+	assert.Equal(t, true, cookie.HttpOnly)
+}