From b6d5bef56c48692509ce50c773a80799d0c7a460 Mon Sep 17 00:00:00 2001
From: MrMelon54 <github@mrmelon54.com>
Date: Sat, 10 Feb 2024 16:23:07 +0000
Subject: [PATCH] Use JWT access tokens

---
 server/jwt.go    | 35 +++++++++++++++++++++++++++++++++++
 server/server.go |  3 +--
 2 files changed, 36 insertions(+), 2 deletions(-)
 create mode 100644 server/jwt.go

diff --git a/server/jwt.go b/server/jwt.go
new file mode 100644
index 0000000..cce69f7
--- /dev/null
+++ b/server/jwt.go
@@ -0,0 +1,35 @@
+package server
+
+import (
+	"context"
+	"crypto/sha256"
+	"encoding/base64"
+	"github.com/1f349/mjwt"
+	"github.com/1f349/mjwt/auth"
+	"github.com/go-oauth2/oauth2/v4"
+	"github.com/golang-jwt/jwt/v4"
+	"github.com/google/uuid"
+	"strings"
+)
+
+type JWTAccessGenerate struct {
+	signer mjwt.Signer
+}
+
+func NewJWTAccessGenerate(signer mjwt.Signer) *JWTAccessGenerate {
+	return &JWTAccessGenerate{signer}
+}
+
+var _ oauth2.AccessGenerate = &JWTAccessGenerate{}
+
+func (j JWTAccessGenerate) Token(ctx context.Context, data *oauth2.GenerateBasic, isGenRefresh bool) (access, refresh string, err error) {
+	access, err = j.signer.GenerateJwt(data.UserID, "", jwt.ClaimStrings{data.Client.GetID()}, data.TokenInfo.GetAccessExpiresIn(), auth.AccessTokenClaims{})
+
+	if isGenRefresh {
+		t := uuid.NewHash(sha256.New(), uuid.New(), []byte(access), 5).String()
+		refresh = base64.URLEncoding.EncodeToString([]byte(t))
+		refresh = strings.ToUpper(strings.TrimRight(refresh, "="))
+	}
+
+	return
+}
diff --git a/server/server.go b/server/server.go
index ae6bc5a..b79b820 100644
--- a/server/server.go
+++ b/server/server.go
@@ -14,7 +14,6 @@ import (
 	scope2 "github.com/1f349/tulip/scope"
 	"github.com/1f349/tulip/theme"
 	"github.com/go-oauth2/oauth2/v4/errors"
-	"github.com/go-oauth2/oauth2/v4/generates"
 	"github.com/go-oauth2/oauth2/v4/manage"
 	"github.com/go-oauth2/oauth2/v4/server"
 	"github.com/go-oauth2/oauth2/v4/store"
@@ -86,7 +85,7 @@ func NewHttpServer(conf Conf, db *database.DB, signingKey mjwt.Signer) *http.Ser
 
 	oauthManager.SetAuthorizeCodeTokenCfg(manage.DefaultAuthorizeCodeTokenCfg)
 	oauthManager.MustTokenStorage(store.NewMemoryTokenStore())
-	oauthManager.MapAccessGenerate(generates.NewAccessGenerate())
+	oauthManager.MapAccessGenerate(NewJWTAccessGenerate(hs.signingKey))
 	oauthManager.MapClientStorage(clientStore.New(db))
 
 	oauthSrv.SetResponseErrorHandler(func(re *errors.Response) {