swt.go

package swt

import (
	"context"
	"crypto/sha256"
	"crypto/sha3"
	"crypto/sha512"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"log/slog"
	"net/http"
	"strings"
	"time"

	"github.com/golang-jwt/jwt/v5"
	"github.com/google/uuid"
)

const (
	defaultExpiration    = 5 * time.Minute
	defaultHashAlgorithm = "sha-256"
	defaultMaxBodySize   = 32 << 20 // 32MB
	jwtHeaderTyp         = "SWT"    // Custom type for identifying an JWT as a Secure Webhook Token
	webhookClaim         = "webhook"

	// Supported hash algorithms
	SHA256   = "sha-256"
	SHA384   = "sha-384"
	SHA512   = "sha-512"
	SHA3_256 = "sha3-256"
	SHA3_384 = "sha3-384"
	SHA3_512 = "sha3-512"
)

var (
	ErrInvalidToken             = errors.New("invalid token - please use New or NewWithClaims to properly initialize token")
	ErrMissingClaims            = errors.New("missing claims")
	ErrInvalidOption            = errors.New("invalid option")
	ErrInvalidData              = errors.New("invalid data")
	ErrInvalidTokenHandler      = errors.New("invalid token handler function")
	ErrInvalidHeaderClaim       = errors.New("invalid token header claim")
	ErrUnsupportedHashAlgorithm = errors.New("unsupported hash algorithm")

	// validSigningMethods defines the allowed signing methods.
	// Draft-02 allows both symmetric (HMAC) and asymmetric (RSA, ECDSA, EdDSA) methods.
	validSigningMethods = []string{
		"HS256", "HS384", "HS512",
		"RS256", "RS384", "RS512",
		"ES256", "ES384", "ES512",
		"EdDSA",
		"PS256", "PS384", "PS512",
	}

	// Ensure Claims implements WebhookClaims
	_ Claims = (*WebhookClaims)(nil)
)

// New creates a SecureWebhookToken with the given issuer, event name, and data.
// Can be further customized with additional Option parameters.
// See Option for available functional parameters.
// Returns an error if claims or options are invalid.
func New(issuer, event string, hash *Hash, opts ...Option) (*SecureWebhookToken, error) {
	claims := NewWebhookClaims(issuer, event)
	claims.Webhook.Hash = hash
	return NewWithClaims(&claims, opts...)
}

// NewWithClaims creates a new SecureWebhookToken with the given Claims object.
// Must be a pointer to an instance of WebhookClaims or a custom Claims instance, which embeds the WebhookClaims.
// Returns an error if claims are nil, options are invalid, or signing method is "none".
func NewWithClaims(c Claims, opts ...Option) (*SecureWebhookToken, error) {
	if c == nil {
		return nil, ErrMissingClaims
	}

	swt := &SecureWebhookToken{token: &jwt.Token{Claims: c}}
	for _, opt := range opts {
		if opt == nil {
			return nil, ErrInvalidOption
		}
		opt(swt)
	}

	// In case the signingMethod was set to nil via WithSigningMethod.
	if swt.token.Method == nil {
		swt.token.Method = jwt.SigningMethodHS256
	}

	// Disallow "none" JWT signing algorithm.
	if swt.token.Method.Alg() == "none" {
		return nil, jwt.NoneSignatureTypeDisallowedError
	}
	// Create the JWT
	swt.token = jwt.NewWithClaims(swt.token.Method, swt.token.Claims)
	// Set the SWT specific Typ on the JWT Header
	swt.token.Header["typ"] = jwtHeaderTyp

	return swt, nil
}

// NewWebhookClaims creates WebhookClaims with a given issuer, event name and hash, and the default registered claims.
func NewWebhookClaims(issuer, event string) WebhookClaims {
	return WebhookClaims{
		Webhook: Webhook{
			Event: event,
		},
		RegisteredClaims: newRegClaims(issuer),
	}
}

// ReplayChecker is an interface for checking and recording token IDs to prevent replay attacks.
// Implementations should store the jti (JWT ID) claim and reject tokens with previously seen IDs.
type ReplayChecker interface {
	// CheckAndRecord checks if a token ID has been seen before and records it if not.
	// Returns an error if the token has already been processed (replay attack detected).
	CheckAndRecord(ctx context.Context, jti string) error
}

type HandlerOptions struct {
	MaxBodySize        int64         // Maximum allowed request body size (default: 32MB)
	Logger             *slog.Logger  // Custom logger (default: slog.Default())
	ReturnErrorDetails bool          // Return JSON error details in response body (default: false)
	ReplayChecker      ReplayChecker // Optional replay attack protection
}

// NewHandlerFunc Creates a new http.HandlerFunc which will process an incoming
// webhook request and pass the token, if valid, to the given handleFn.
// Returns an error if handleFn is nil.
func NewHandlerFunc(secret []byte, handleFn func(token *SecureWebhookToken, data []byte) error, opts *HandlerOptions) (http.HandlerFunc, error) {
	var (
		token *SecureWebhookToken
		body  []byte
		err   error
	)

	if handleFn == nil {
		return nil, ErrInvalidTokenHandler
	}

	if opts == nil {
		opts = &HandlerOptions{
			MaxBodySize: defaultMaxBodySize,
		}
	}

	logger := opts.Logger
	if logger == nil {
		logger = slog.Default()
	}
	logger = logger.With("package", "github.com/SecureWebhookToken/swt")

	// Helper for sending error responses
	sendError := func(w http.ResponseWriter, status int, message string) {
		w.WriteHeader(status)
		if opts.ReturnErrorDetails {
			w.Header().Set("Content-Type", "application/json")
			_ = json.NewEncoder(w).Encode(map[string]string{"error": message})
		}
	}

	return func(w http.ResponseWriter, r *http.Request) {
		if r.Method != http.MethodPost {
			logger.Warn(fmt.Sprintf("method not allowed: %s", r.Method), "request", r)
			w.WriteHeader(http.StatusMethodNotAllowed)
			return
		}

		authHeader := r.Header.Get("Authorization")
		res := strings.Split(authHeader, "Bearer ")
		if len(res) != 2 {
			logger.Warn(fmt.Sprintf("invalid authorization header: %s", authHeader), "request", r)
			sendError(w, http.StatusBadRequest, "invalid authorization header")
			return
		}

		tokenStr := strings.TrimSpace(res[1])
		token, err = Parse(tokenStr, secret)
		if err != nil {
			logger.Warn(fmt.Sprintf("parse token error: %v", err), "request", r)
			sendError(w, http.StatusBadRequest, "invalid or expired token")
			return
		}

		// Check for replay attacks if replay checker is provided
		if opts.ReplayChecker != nil {
			jti, _ := token.Claims().GetID()
			if err = opts.ReplayChecker.CheckAndRecord(r.Context(), jti); err != nil {
				logger.Warn("replay attack detected", "jti", jti, "error", err)
				sendError(w, http.StatusBadRequest, "duplicate token detected")
				return
			}
		}

		if r.ContentLength == -1 {
			logger.Error("unknown content length", "request", r)
			sendError(w, http.StatusBadRequest, "unknown content length")
			return
		}

		wh, _ := token.Claims().GetWebhook()

		// Limit the request body to a given size to prevent wasting server resources by malicious clients
		body, err = io.ReadAll(http.MaxBytesReader(w, r.Body, opts.MaxBodySize))

		var maxBytesError *http.MaxBytesError
		if errors.As(err, &maxBytesError) {
			logger.Error("max request size reached", "error", err)
			sendError(w, http.StatusBadRequest, "max request size reached")
			return
		}

		if err != nil {
			logger.Error("failed to read request body", "error", err)
			sendError(w, http.StatusBadRequest, "failed to read request body")
			return
		}

		if len(body) > 0 {
			err = validateBody(&wh, body)
			if err != nil {
				logger.Error("failed to validate request body", "error", err)
				sendError(w, http.StatusBadRequest, "failed to validate request body")
				return
			}
		} else if wh.Hash != nil {
			// Draft-02 Section 11: If the request body is empty, the webhook claim MUST NOT contain a hash field.
			logger.Error("webhook claim contains hash but request body is empty")
			sendError(w, http.StatusBadRequest, "hash claim not allowed for empty body")
			return
		}

		// Pass the valid token to the handler function
		err = handleFn(token, body)
		if err != nil {
			logger.Error("handleFn error!", "error", err, "token", token.String())
			sendError(w, http.StatusBadRequest, err.Error())
			return
		}
		w.WriteHeader(http.StatusNoContent)
	}, nil
}

// Parse parses a given token string, verifies it and returns a SecureWebhookToken if successful.
func Parse(tokenStr string, key any) (*SecureWebhookToken, error) {
	return ParseWithClaims(tokenStr, &WebhookClaims{}, key)
}

// ParseWithClaims parses a given token string and given claims, verifies it and returns a SecureWebhookToken if successful.
func ParseWithClaims(tokenStr string, claims Claims, key any) (*SecureWebhookToken, error) {
	var (
		swt = &SecureWebhookToken{}
		err error
	)

	if claims == nil {
		return nil, jwt.ErrTokenInvalidClaims
	}

	if swt.token, err = jwt.ParseWithClaims(
		tokenStr,
		claims,
		func(token *jwt.Token) (any, error) {
			return key, nil
		},
		jwt.WithValidMethods(validSigningMethods),
		jwt.WithExpirationRequired(),
		jwt.WithIssuedAt(),
	); err != nil {
		return nil, err
	}

	if typ, ok := swt.token.Header["typ"]; !ok || typ != jwtHeaderTyp {
		return nil, fmt.Errorf("%w: header claim typ=SWT expected", ErrInvalidHeaderClaim)
	}

	return swt, nil
}

// ParseWithContext parses a given token string with context support, verifies it and returns a SecureWebhookToken if successful.
// The context can be used for cancellation and timeout control.
func ParseWithContext(ctx context.Context, tokenStr string, key any) (*SecureWebhookToken, error) {
	return ParseWithClaimsContext(ctx, tokenStr, &WebhookClaims{}, key)
}

// ParseWithClaimsContext parses a given token string and given claims with context support, verifies it and returns a SecureWebhookToken if successful.
// The context can be used for cancellation and timeout control.
func ParseWithClaimsContext(ctx context.Context, tokenStr string, claims Claims, key any) (*SecureWebhookToken, error) {
	// Check context before doing any work
	if err := ctx.Err(); err != nil {
		return nil, err
	}

	return ParseWithClaims(tokenStr, claims, key)
}

// SecureWebhookToken is a structure for secure webhook tokens.
type SecureWebhookToken struct {
	token *jwt.Token
}

// Algorithm returns the used SigningMethod algorithm.
func (swt *SecureWebhookToken) Algorithm() string {
	if swt.token == nil {
		panic(ErrInvalidToken)
	}
	return swt.token.Method.Alg()
}

// Claims return all token Claims.
// If no custom WebhookClaims have been set with NewWithClaims(), then the underlying type will be *WebhookClaims.
func (swt *SecureWebhookToken) Claims() Claims {
	if swt.token == nil {
		panic(ErrInvalidToken)
	}
	return swt.token.Claims.(Claims)
}

// ID convenient method for accessing the ID of the SecureWebhookToken.
func (swt *SecureWebhookToken) ID() string {
	return swt.Claims().wc().ID
}

// Issuer convenient method for accessing the Issuer claim of the SecureWebhookToken.
func (swt *SecureWebhookToken) Issuer() string {
	return swt.Claims().wc().Issuer
}

// Webhook convenient method for accessing the Webhook claim of the SecureWebhookToken.
func (swt *SecureWebhookToken) Webhook() Webhook {
	return swt.Claims().wc().Webhook
}

// SignedString returns the encoded and signed SecureWebhookToken as a JWT string.
func (swt *SecureWebhookToken) SignedString(key any) (string, error) {
	if swt.token == nil {
		return "", ErrInvalidToken
	}

	// Validate claims before signing
	if err := swt.Validate(); err != nil {
		return "", err
	}
	return swt.token.SignedString(key)
}

// Valid returns true only if the SecureWebhookToken has been created
// via Parse method and successfully been validated.
func (swt *SecureWebhookToken) Valid() bool {
	if swt.token == nil {
		return false
	}
	return swt.token.Valid
}

// Validate validates the Claims of the SecureWebhookToken.
func (swt *SecureWebhookToken) Validate() error {
	v := jwt.NewValidator(
		jwt.WithValidMethods(validSigningMethods),
		jwt.WithExpirationRequired(),
		jwt.WithIssuedAt(),
	)
	return v.Validate(swt.Claims())
}

// String implements the Stringer interface for returning the SecureWebhookToken as a JSON string.
func (swt *SecureWebhookToken) String() string {
	tokenData := struct {
		Header    map[string]any `json:"header"`
		Payload   jwt.Claims     `json:"payload"`
		Signature string         `json:"signature"`
		Validated bool           `json:"validated"`
	}{
		Header:    swt.token.Header,
		Payload:   swt.token.Claims,
		Signature: swt.token.EncodeSegment(swt.token.Signature),
		Validated: swt.token.Valid,
	}
	out, _ := json.Marshal(tokenData)

	return string(out)
}

type Webhook struct {
	Event      string `json:"event"` // Event name. Should have the form EVENT_NAME.ACTIVITY (e.g., user.created or pull_request.merged)
	Hash       *Hash  `json:"hash,omitempty"`
	RetryCount uint   `json:"retry_count,omitempty"`
}

type Hash string

func NewHash(alg string, data []byte) (*Hash, error) {
	if alg == "" {
		alg = defaultHashAlgorithm
	}
	if data != nil {
		hSum, err := HashSum(alg, data)
		if err != nil {
			return nil, err
		}
		h := Hash(strings.ToLower(alg) + ":" + hSum)
		return &h, nil
	}
	return nil, nil
}

func NewHashFromString(input string) *Hash {
	h := Hash(strings.ToLower(input))
	return &h
}

func (h *Hash) parts() (algo, sum string) {
	res := strings.Split(string(*h), ":")
	if len(res) != 2 {
		return "", ""
	}
	return res[0], res[1]
}

func (h *Hash) Algorithm() string {
	alg, _ := h.parts()
	return alg
}

func (h *Hash) Sum() string {
	_, sum := h.parts()
	return sum
}

func (h *Hash) String() string {
	alg, sum := h.parts()
	if alg == "" && sum == "" {
		return ""
	}
	return alg + ":" + sum
}

type Claims interface {
	GetID() (string, error)
	GetWebhook() (Webhook, error)
	wc() *WebhookClaims
	jwt.Claims
}

type WebhookClaims struct {
	Webhook Webhook `json:"webhook"`
	jwt.RegisteredClaims
}

// the `jti` (JWT ID) claim. See https://datatracker.ietf.org/doc/html/rfc7519#section-4.1.7
func (wc *WebhookClaims) GetID() (string, error) {
	return wc.ID, nil
}

// the `webhook` (SWT) claim. See https://www.ietf.org/archive/id/draft-knauer-secure-webhook-token-02.html#name-json-web-token-claims-regis
func (wc *WebhookClaims) GetWebhook() (Webhook, error) {
	return wc.Webhook, nil
}

// MapClaims will convert a WebhookClaims object into jwt.MapClaims
// Can easily be used with the golang-jwt package for creating JWTs directly if desired.
func (wc *WebhookClaims) MapClaims() jwt.MapClaims {
	mc := map[string]any{
		"aud":     wc.Audience,
		"sub":     wc.Subject,
		"jti":     wc.ID,
		"iss":     wc.Issuer,
		"webhook": wc.Webhook,
	}
	if wc.ExpiresAt != nil {
		mc["exp"] = float64(wc.ExpiresAt.Unix())
	}
	if wc.IssuedAt != nil {
		mc["iat"] = float64(wc.IssuedAt.Unix())
	}
	if wc.NotBefore != nil {
		mc["nbf"] = float64(wc.NotBefore.Unix())
	}
	return mc
}

// Validate implements the jwt.ClaimsValidator interface to perform further required validation on specific claims.
func (wc *WebhookClaims) Validate() error {
	if wc.Issuer == "" {
		return fmt.Errorf("%w: iss must not be empty", jwt.ErrTokenInvalidIssuer)
	}

	if wc.IssuedAt == nil {
		return fmt.Errorf("%w: iat", jwt.ErrTokenRequiredClaimMissing)
	}

	if wc.ExpiresAt == nil {
		return fmt.Errorf("%w: exp", jwt.ErrTokenRequiredClaimMissing)
	}

	if wc.NotBefore == nil {
		return fmt.Errorf("%w: nbf", jwt.ErrTokenRequiredClaimMissing)
	}

	if wc.Webhook.Event == "" {
		return fmt.Errorf("%w: %s must contain an event", jwt.ErrTokenInvalidClaims, webhookClaim)
	}

	if wc.ID == "" {
		return fmt.Errorf("%w: jti must not be empty! a uuid or similar is recommended", jwt.ErrTokenInvalidId)
	}

	return nil
}

// wc is a helper method for overriding the default claims with functional options.
func (wc *WebhookClaims) wc() *WebhookClaims {
	return wc
}

// HashSum computes the hash sum for a given hashAlg and the body to be sent via http.MethodPost request.
func HashSum(hashAlg string, body []byte) (string, error) {
	var hash string

	if hashAlg == "" {
		hashAlg = defaultHashAlgorithm
	}

	switch strings.ToLower(hashAlg) {
	case SHA256:
		hash = fmt.Sprintf("%x", sha256.Sum256(body))
	case SHA384:
		hash = fmt.Sprintf("%x", sha512.Sum384(body))
	case SHA512:
		hash = fmt.Sprintf("%x", sha512.Sum512(body))
	case SHA3_256:
		hash = fmt.Sprintf("%x", sha3.Sum256(body))
	case SHA3_384:
		hash = fmt.Sprintf("%x", sha3.Sum384(body))
	case SHA3_512:
		hash = fmt.Sprintf("%x", sha3.Sum512(body))
	default:
		return "", fmt.Errorf("%w: %s", ErrUnsupportedHashAlgorithm, hashAlg)
	}

	return hash, nil
}

// validateBody validates the given body against the used hash function and signature.
func validateBody(wh *Webhook, body []byte) error {
	if wh == nil {
		return fmt.Errorf("%w: webhook cannot be nil", jwt.ErrTokenRequiredClaimMissing)
	}
	if wh.Hash == nil {
		return fmt.Errorf("%w: webhook.hash cannot be nil", jwt.ErrTokenRequiredClaimMissing)
	}

	hs, err := HashSum(wh.Hash.Algorithm(), body)
	if err != nil {
		return err
	}
	if hs != wh.Hash.Sum() {
		return ErrInvalidData
	}

	return nil
}

// newRegClaims creates jwt.RegisteredClaims with sensible defaults and the required issuer (iss) claim.
func newRegClaims(issuer string) jwt.RegisteredClaims {
	now := time.Now()
	return jwt.RegisteredClaims{
		ID:        uuid.NewString(),
		Issuer:    issuer,
		IssuedAt:  jwt.NewNumericDate(now),
		ExpiresAt: jwt.NewNumericDate(now.Add(defaultExpiration)),
		NotBefore: jwt.NewNumericDate(now),
	}
}