feat: jwk endpoint

internal/bootstrap/router_bootstrap.go

@@ -113,9 +113,7 @@ func (app *BootstrapApp) setupRouter() (*gin.Engine, error) {
 
 	healthController.SetupRoutes()
 
-	wellknownController := controller.NewWellKnownController(controller.WellKnownControllerConfig{
-		OpenIDConnectIssuer: app.services.oidcService.GetIssuer(),
-	}, engine)
+	wellknownController := controller.NewWellKnownController(controller.WellKnownControllerConfig{}, app.services.oidcService, engine)
 
 	wellknownController.SetupRoutes()
 

internal/controller/oidc_controller.go

@@ -150,7 +150,7 @@ func (controller *OIDCController) Authorize(c *gin.Context) {
 	}
 
 	// We also need a snapshot of the user that authorized this (skip if no openid scope)
-	if slices.Contains(strings.Split(req.Scope, " "), "openid") {
+	if slices.Contains(strings.Fields(req.Scope), "openid") {
 		err = controller.oidc.StoreUserinfo(c, sub, userContext, req)
 
 		if err != nil {

internal/controller/well_known_controller.go

@@ -2,6 +2,7 @@ package controller
 
 import (
 	"fmt"
+	"net/http"
 
 	"github.com/gin-gonic/gin"
 	"github.com/steveiliop56/tinyauth/internal/service"
@@ -23,33 +24,35 @@ type OpenIDConnectConfiguration struct {
 	ServiceDocumentation              string   `json:"service_documentation"`
 }
 
-type WellKnownControllerConfig struct {
-	OpenIDConnectIssuer string
-}
+type WellKnownControllerConfig struct{}
 
 type WellKnownController struct {
 	config WellKnownControllerConfig
 	engine *gin.Engine
+	oidc   *service.OIDCService
 }
 
-func NewWellKnownController(config WellKnownControllerConfig, engine *gin.Engine) *WellKnownController {
+func NewWellKnownController(config WellKnownControllerConfig, oidc *service.OIDCService, engine *gin.Engine) *WellKnownController {
 	return &WellKnownController{
 		config: config,
+		oidc:   oidc,
 		engine: engine,
 	}
 }
 
 func (controller *WellKnownController) SetupRoutes() {
 	controller.engine.GET("/.well-known/openid-configuration", controller.OpenIDConnectConfiguration)
+	controller.engine.GET("/.well-known/jwks.json", controller.JWKS)
 }
 
 func (controller *WellKnownController) OpenIDConnectConfiguration(c *gin.Context) {
+	issuer := controller.oidc.GetIssuer()
 	c.JSON(200, OpenIDConnectConfiguration{
-		Issuer:                            controller.config.OpenIDConnectIssuer,
-		AuthorizationEndpoint:             fmt.Sprintf("%s/authorize", controller.config.OpenIDConnectIssuer),
-		TokenEndpoint:                     fmt.Sprintf("%s/api/oidc/token", controller.config.OpenIDConnectIssuer),
-		UserinfoEndpoint:                  fmt.Sprintf("%s/api/oidc/userinfo", controller.config.OpenIDConnectIssuer),
-		JwksUri:                           fmt.Sprintf("%s/api/oidc/jwks", controller.config.OpenIDConnectIssuer),
+		Issuer:                            issuer,
+		AuthorizationEndpoint:             fmt.Sprintf("%s/authorize", issuer),
+		TokenEndpoint:                     fmt.Sprintf("%s/api/oidc/token", issuer),
+		UserinfoEndpoint:                  fmt.Sprintf("%s/api/oidc/userinfo", issuer),
+		JwksUri:                           fmt.Sprintf("%s/.well-known/jwks.json", issuer),
 		ScopesSupported:                   service.SupportedScopes,
 		ResponseTypesSupported:            service.SupportedResponseTypes,
 		GrantTypesSupported:               service.SupportedGrantTypes,
@@ -60,3 +63,23 @@ func (controller *WellKnownController) OpenIDConnectConfiguration(c *gin.Context
 		ServiceDocumentation:              "https://tinyauth.app/docs/reference/openid",
 	})
 }
+
+func (controller *WellKnownController) JWKS(c *gin.Context) {
+	jwks, err := controller.oidc.GetJWK()
+
+	if err != nil {
+		c.JSON(500, gin.H{
+			"status":  "500",
+			"message": "failed to get JWK",
+		})
+		return
+	}
+
+	c.Header("content-type", "application/json")
+
+	c.Writer.WriteString(`{"keys":[`)
+	c.Writer.Write(jwks)
+	c.Writer.WriteString(`]}`)
+
+	c.Status(http.StatusOK)
+}

internal/service/oidc_service.go

@@ -8,6 +8,7 @@ import (
 	"crypto/sha256"
 	"crypto/x509"
 	"database/sql"
+	"encoding/json"
 	"encoding/pem"
 	"errors"
 	"fmt"
@@ -17,14 +18,12 @@ import (
 	"time"
 
 	"github.com/gin-gonic/gin"
+	"github.com/go-jose/go-jose/v4"
 	"github.com/steveiliop56/tinyauth/internal/config"
 	"github.com/steveiliop56/tinyauth/internal/repository"
 	"github.com/steveiliop56/tinyauth/internal/utils"
 	"github.com/steveiliop56/tinyauth/internal/utils/tlog"
 	"golang.org/x/exp/slices"
-
-	// Should probably switch to another package but for now this works
-	"golang.org/x/oauth2/jws"
 )
 
 var (
@@ -40,6 +39,14 @@ var (
 	ErrTokenExpired  = errors.New("token_expired")
 )
 
+type ClaimSet struct {
+	Iss string `json:"iss"`
+	Aud string `json:"aud"`
+	Sub string `json:"sub"`
+	Iat int64  `json:"iat"`
+	Exp int64  `json:"exp"`
+}
+
 type UserinfoResponse struct {
 	Sub               string   `json:"sub"`
 	Name              string   `json:"name"`
@@ -333,7 +340,21 @@ func (service *OIDCService) generateIDToken(client config.OIDCClientConfig, sub
 	createdAt := time.Now().Unix()
 	expiresAt := time.Now().Add(time.Duration(service.config.SessionExpiry) * time.Second).Unix()
 
-	claims := jws.ClaimSet{
+	signer, err := jose.NewSigner(jose.SigningKey{
+		Algorithm: jose.RS256,
+		Key:       service.privateKey,
+	}, &jose.SignerOptions{
+		ExtraHeaders: map[jose.HeaderKey]any{
+			"typ": "jwt",
+			"jku": fmt.Sprintf("%s/.well-known/jwks.json", service.issuer),
+		},
+	})
+
+	if err != nil {
+		return "", err
+	}
+
+	claims := ClaimSet{
 		Iss: service.issuer,
 		Aud: client.ClientID,
 		Sub: sub,
@@ -341,12 +362,19 @@ func (service *OIDCService) generateIDToken(client config.OIDCClientConfig, sub
 		Exp: expiresAt,
 	}
 
-	header := jws.Header{
-		Algorithm: "RS256",
-		Typ:       "JWT",
+	payload, err := json.Marshal(claims)
+
+	if err != nil {
+		return "", err
 	}
 
-	token, err := jws.Encode(&header, &claims, service.privateKey)
+	object, err := signer.Sign(payload)
+
+	if err != nil {
+		return "", err
+	}
+
+	token, err := object.CompactSerialize()
 
 	if err != nil {
 		return "", err
@@ -595,3 +623,13 @@ func (service *OIDCService) Cleanup() {
 		}
 	}
 }
+
+func (service *OIDCService) GetJWK() ([]byte, error) {
+	jwk := jose.JSONWebKey{
+		Key:       service.privateKey,
+		Algorithm: string(jose.RS256),
+		Use:       "sig",
+	}
+
+	return jwk.Public().MarshalJSON()
+}