package main

/*
  DESCRIPTION:
    - built on top of golang-jwt
    - Supports multiple audiences and issuers
    - Has h.RequireRole("role")-Middleware that ensures certain role availability

  USAGE:
  
  // Setup the Bearer information
  auth := &JwtBearerAuthHandler{
    symmetricKey: []byte("your-256-bit-secret"),
    // optional
    validIssuers: map[string]struct{}{
      "jwt.io": {},
    },
    // optional
    validAudiences: map[string]struct{}{
      "company": {},
      "example": {},
    }
    
    // Protect an endpoint by requiring a role
    router.DELETE("/delete",
        auth.RequireRole("admin"),
        func(c *gin.Context) {
            result, err := t.Exec(c.Request.Context(), "DELETE FROM Todos")
            if err != nil {
              c.AbortWithError(http.StatusInternalServerError, err)
              return
            }
            c.JSON(http.StatusOK, result)
          })
*/

import (
	"errors"
	"net/http"
	"strings"

	"github.com/gin-gonic/gin"
	"github.com/golang-jwt/jwt/v5"
)

type AuthHandler interface {
	RequireRole(requiredRole string) gin.HandlerFunc
}

type JwtBearerAuthHandler struct {
	symmetricKey   []byte
	validIssuers   map[string]struct{}
	validAudiences map[string]struct{}
}

type roleClaims struct {
	jwt.RegisteredClaims
	// Should be either string or []interface{} containing strings.
	// values like "role": "admin" and "role": ["admin", "user"] are supported
	Role  jwt.ClaimStrings `json:"role"`
	Scope string           `json:"scope"`
}

const bearerTokenHandlerKey = "bearer_token"

func (h *JwtBearerAuthHandler) ensureToken(c *gin.Context) bool {
	if _, ok := c.Get(bearerTokenHandlerKey); ok {
		return true
	}

	const BearerPrefix = "Bearer "
	authorization := c.GetHeader("Authorization")
	if len(authorization) < len(BearerPrefix) {
		return false
	}
	if !strings.EqualFold(authorization[:7], BearerPrefix) {
		return false
	}
	jwtToken := strings.TrimSpace(authorization[7:])
	token, err := jwt.ParseWithClaims(jwtToken,
		&roleClaims{},
		func(token *jwt.Token) (interface{}, error) {
			if _, ok := token.Method.(*jwt.SigningMethodHMAC); !ok {
				return nil, errors.New("unexpected signing method received")
			}
			return h.symmetricKey, nil
		})
	if err != nil || token == nil || !token.Valid {
		return false
	}

	if !h.validateIssuer(token) {
		return false
	}
	if !h.validateAudience(token) {
		return false
	}

	c.Set(bearerTokenHandlerKey, token)
	return true
}

func (h *JwtBearerAuthHandler) validateIssuer(token *jwt.Token) bool {
	if h.validIssuers == nil {
		return true
	}
	mapped := token.Claims.(*roleClaims)
	_, found := h.validIssuers[mapped.Issuer]
	return found
}

func (h *JwtBearerAuthHandler) validateAudience(token *jwt.Token) bool {
	if h.validAudiences == nil {
		return true
	}
	mapped := token.Claims.(*roleClaims)

	for _, audience := range mapped.Audience {
		if _, found := h.validAudiences[audience]; found {
			return true
		}
	}
	return true
}

func (h *JwtBearerAuthHandler) RequireRole(requiredRole string) gin.HandlerFunc {
	return func(c *gin.Context) {
		if !h.ensureToken(c) {
			c.AbortWithStatus(http.StatusUnauthorized)
			return
		}
		tokenIface, ok := c.Get(bearerTokenHandlerKey)
		if !ok {
			c.AbortWithStatus(http.StatusUnauthorized)
			return
		}

		token := tokenIface.(*jwt.Token)
		mapped := token.Claims.(*roleClaims)

		for _, role := range mapped.Role {
			if role == requiredRole {
				return
			}
		}

		c.AbortWithStatus(http.StatusForbidden)
	}
}