sec-jwt-library-rs256

Answer

Complete Implementation — jwt.go

The library uses only Go's standard library (crypto/rsa, crypto/sha256, crypto/x509, encoding/json, etc.) to create and verify RS256-signed JWTs. No third-party dependencies.

Key design decisions:

package jwt

import (
    "crypto"
    "crypto/rand"
    "crypto/rsa"
    "crypto/sha256"
    "crypto/x509"
    "encoding/base64"
    "encoding/json"
    "encoding/pem"
    "errors"
    "fmt"
    "strings"
    "time"
)

const algorithm = "RS256"

var (
    ErrInvalidToken      = errors.New("jwt: invalid token")
    ErrInvalidSignature  = errors.New("jwt: invalid signature")
    ErrTokenExpired      = errors.New("jwt: token has expired")
    ErrUnexpectedSigning = errors.New("jwt: unexpected signing method")
    ErrKeyNotFound       = errors.New("jwt: key not found")
)

type KeyPair struct {
    ID         string
    PrivateKey *rsa.PrivateKey
    PublicKey  *rsa.PublicKey
}

type JWT struct {
    keys        map[string]*KeyPair
    activeKeyID string
}

// New creates a JWT instance seeded with one key pair.
func New(keyID string, privateKeyPEM []byte) (*JWT, error) {
    j := &JWT{keys: make(map[string]*KeyPair)}
    if err := j.AddKey(keyID, privateKeyPEM); err != nil {
        return nil, err
    }
    j.activeKeyID = keyID
    return j, nil
}

// AddKey registers a PEM-encoded RSA private key for verification (key rotation).
func (j *JWT) AddKey(keyID string, privateKeyPEM []byte) error {
    block, _ := pem.Decode(privateKeyPEM)
    if block == nil {
        return fmt.Errorf("jwt: failed to decode PEM block")
    }
    var priv *rsa.PrivateKey
    priv, err := x509.ParsePKCS1PrivateKey(block.Bytes)
    if err != nil {
        priv2, err2 := x509.ParsePKCS8PrivateKey(block.Bytes)
        if err2 != nil {
            return fmt.Errorf("jwt: failed to parse private key (PKCS1: %v, PKCS8: %v)", err, err2)
        }
        var ok bool
        priv, ok = priv2.(*rsa.PrivateKey)
        if !ok {
            return fmt.Errorf("jwt: key is not an RSA private key")
        }
    }
    j.keys[keyID] = &KeyPair{ID: keyID, PrivateKey: priv, PublicKey: &priv.PublicKey}
    return nil
}

// SetActiveKey sets the key used for signing new tokens.
func (j *JWT) SetActiveKey(keyID string) error {
    if _, ok := j.keys[keyID]; !ok {
        return ErrKeyNotFound
    }
    j.activeKeyID = keyID
    return nil
}

// RemoveKey removes a verification key (cannot remove the active signing key).
func (j *JWT) RemoveKey(keyID string) error {
    if _, ok := j.keys[keyID]; !ok {
        return ErrKeyNotFound
    }
    if keyID == j.activeKeyID {
        return fmt.Errorf("jwt: cannot remove the active signing key")
    }
    delete(j.keys, keyID)
    return nil
}

// Create generates a signed JWT with the given claims and an exp claim.
func (j *JWT) Create(claims map[string]interface{}, expiry time.Duration) (string, error) {
    if j.activeKeyID == "" {
        return "", fmt.Errorf("jwt: no active signing key")
    }
    key, ok := j.keys[j.activeKeyID]
    if !ok {
        return "", ErrKeyNotFound
    }
    claimsCopy := make(map[string]interface{}, len(claims)+1)
    for k, v := range claims {
        claimsCopy[k] = v
    }
    claimsCopy["exp"] = time.Now().Add(expiry).Unix()
    return j.compact(claimsCopy, key)
}

// Verify decodes and validates a JWT: checks structure, algorithm, signature, and expiry.
func (j *JWT) Verify(tokenString string) (map[string]interface{}, error) {
    tok, err := j.parse(tokenString)
    if err != nil {
        return nil, err
    }
    if tok.Header.Alg != algorithm {
        return nil, ErrUnexpectedSigning
    }
    kid := tok.Header.Kid
    if kid == "" {
        for _, k := range j.keys {
            kid = k.ID
            break
        }
    }
    key, ok := j.keys[kid]
    if !ok {
        return nil, ErrKeyNotFound
    }
    parts := strings.Split(tokenString, ".")
    signingInput := parts[0] + "." + parts[1]
    if err := verifyRS256(key.PublicKey, []byte(signingInput), tok.Signature); err != nil {
        return nil, ErrInvalidSignature
    }
    if err := verifyTime(tok.Claims); err != nil {
        return nil, err
    }
    return tok.Claims, nil
}

// --- internal helpers -------------------------------------------------------

func base64URLEncode(src []byte) string { return base64.RawURLEncoding.EncodeToString(src) }

func base64URLDecode(s string) ([]byte, error) {
    switch len(s) % 4 {
    case 2:
        s += "=="
    case 3:
        s += "="
    }
    return base64.URLEncoding.DecodeString(s)
}

func signRS256(priv *rsa.PrivateKey, data []byte) ([]byte, error) {
    h := sha256.Sum256(data)
    return rsa.SignPKCS1v15(rand.Reader, priv, crypto.SHA256, h[:])
}

func verifyRS256(pub *rsa.PublicKey, data, sig []byte) error {
    h := sha256.Sum256(data)
    return rsa.VerifyPKCS1v15(pub, crypto.SHA256, h[:], sig)
}

type header struct {
    Alg string `json:"alg"`
    Kid string `json:"kid,omitempty"`
    Typ string `json:"typ,omitempty"`
}

type jwtToken struct {
    Header    header
    Claims    map[string]interface{}
    Signature []byte
}

func (j *JWT) compact(claims map[string]interface{}, key *KeyPair) (string, error) {
    h := header{Alg: algorithm, Kid: key.ID, Typ: "JWT"}
    hdr, _ := json.Marshal(h)
    clm, _ := json.Marshal(claims)
    hdrB64 := base64URLEncode(hdr)
    clmB64 := base64URLEncode(clm)
    sig, err := signRS256(key.PrivateKey, []byte(hdrB64+"."+clmB64))
    if err != nil {
        return "", err
    }
    return hdrB64 + "." + clmB64 + "." + base64URLEncode(sig), nil
}

func (j *JWT) parse(tokenString string) (*jwtToken, error) {
    parts := strings.Split(tokenString, ".")
    if len(parts) != 3 {
        return nil, ErrInvalidToken
    }
    hdr, err := base64URLDecode(parts[0])
    if err != nil {
        return nil, fmt.Errorf("jwt: decode header: %w", err)
    }
    clm, err := base64URLDecode(parts[1])
    if err != nil {
        return nil, fmt.Errorf("jwt: decode claims: %w", err)
    }
    sig, err := base64URLDecode(parts[2])
    if err != nil {
        return nil, fmt.Errorf("jwt: decode signature: %w", err)
    }
    var h header
    json.Unmarshal(hdr, &h)
    var c map[string]interface{}
    json.Unmarshal(clm, &c)
    return &jwtToken{Header: h, Claims: c, Signature: sig}, nil
}

func verifyTime(claims map[string]interface{}) error {
    expClaim, ok := claims["exp"]
    if !ok {
        return nil
    }
    var exp int64
    switch v := expClaim.(type) {
    case float64:
        exp = int64(v)
    default:
        return fmt.Errorf("jwt: invalid exp claim")
    }
    if time.Now().Unix() > exp {
        return ErrTokenExpired
    }
    return nil
}