NWJWT/token.go

package nwjwt

import (
	"bytes"
	"crypto/hmac"
	"crypto/sha256"
	"encoding/base64"
	"encoding/json"
	"strings"
	"time"
)

type Token struct {
	header    Header
	claims    ClaimSet
	signature []byte
}

func NewToken() *Token {
	return &Token{
		header: DefaultHeader(),
		claims: ClaimSet{},
	}
}

func TokenFromString(tokenAsString string) *Token {
	parts := strings.SplitN(tokenAsString, ".", 3)
	if len(parts) != 3 {
		return nil
	}

	var header Header
	headerData, _ := base64.RawURLEncoding.DecodeString(parts[0])
	json.Unmarshal(headerData, &header)

	var claims map[string]any
	claimsData, _ := base64.RawURLEncoding.DecodeString(parts[1])
	json.Unmarshal(claimsData, &claims)

	signature, _ := base64.RawURLEncoding.DecodeString(parts[2])

	return &Token{
		header:    header,
		claims:    claims,
		signature: signature,
	}
}

func (t *Token) Sign(secret []byte) {
	t.signature = t.sign(secret)
}

func (t *Token) SetExpiry(time time.Time) {
	t.claims[Expiry] = time.Unix()
}

func (t *Token) SetIssuedAt() {
	t.claims[IssuedAt] = time.Now().Unix()
}

func (t *Token) SetNotBefore(time time.Time) {
	t.claims[NotBefore] = time.Unix()
}

func (t *Token) SetClaim(key string, value any) {
	t.claims[key] = value
}

func (t Token) Claim(key string) any {
	return t.claims[key]
}

func (t Token) Validate(secret []byte, conditions []Condition) error {
	if !bytes.Equal(t.sign(secret), t.signature) {
		return ErrSignatureInvalid
	}

	if exp, hasExp := t.claims[Expiry]; hasExp {
		if int64(exp.(float64)) < time.Now().Unix() {
			return ErrExpired
		}
	}

	if iat, hasIat := t.claims[IssuedAt]; hasIat {
		if int64(iat.(float64)) > time.Now().Unix() {
			return ErrBeforeIssuedAt
		}
	}

	if nbf, hasNbf := t.claims[NotBefore]; hasNbf {
		if int64(nbf.(float64)) > time.Now().Unix() {
			return ErrBeforeNotBefore
		}
	}

	//Check extra conditions
	for _, condition := range conditions {
		if claim, hasClaim := t.claims[condition.Key]; hasClaim {
			switch condition.Op {
			case FilterOpEqual:
				if claim == condition.Value {
					continue
				}
			case FilterOpNotEqual:
				if claim != condition.Value {
					continue
				}
			case FilterOpGreaterThan:
				claimAsNum, claimIsNum := anyToNum(claim)
				conditionAsNum, conditionIsNum := anyToNum(condition.Value)
				if claimIsNum && conditionIsNum && conditionAsNum > claimAsNum {
					continue
				}
			case FilterOpLessThan:
				claimAsNum, claimIsNum := anyToNum(claim)
				conditionAsNum, conditionIsNum := anyToNum(condition.Value)
				if claimIsNum && conditionIsNum && conditionAsNum < claimAsNum {
					continue
				}
			case FilterOpGreaterOrEqual:
				claimAsNum, claimIsNum := anyToNum(claim)
				conditionAsNum, conditionIsNum := anyToNum(condition.Value)
				if claimIsNum && conditionIsNum && conditionAsNum >= claimAsNum {
					continue
				}
			case FilterOpLessOrEqual:
				claimAsNum, claimIsNum := anyToNum(claim)
				conditionAsNum, conditionIsNum := anyToNum(condition.Value)
				if claimIsNum && conditionIsNum && conditionAsNum <= claimAsNum {
					continue
				}
			}

			return ErrConditionsDoNotMatch
		}
	}

	return nil
}

func (t Token) String() string {
	buf := new(bytes.Buffer)
	be := base64.NewEncoder(base64.RawURLEncoding, buf)

	buf.Write(t.encodeHeaderAndClaims())
	buf.WriteByte('.')
	be.Write(t.signature)
	be.Close()
	return buf.String()
}

func (t Token) sign(secret []byte) []byte {
	hash := hmac.New(sha256.New, secret)
	hash.Write(t.encodeHeaderAndClaims())
	return hash.Sum(nil)
}

func (t Token) encodeHeaderAndClaims() []byte {
	buf := new(bytes.Buffer)
	be := base64.NewEncoder(base64.RawURLEncoding, buf)

	headerData, _ := json.Marshal(t.header)
	be.Write(headerData)
	buf.WriteByte('.')

	claimsData, _ := json.Marshal(t.claims)
	be.Write(claimsData)
	be.Close()
	return buf.Bytes()
}