tokens.go

// Copyright (c) 2021-2023, R.I. Pienaar and the Choria Project contributors
//
// SPDX-License-Identifier: Apache-2.0

package tokens

import (
	"bytes"
	"context"
	"crypto/ed25519"
	"crypto/rsa"
	"crypto/tls"
	"encoding/base64"
	"encoding/hex"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"os"
	"regexp"
	"strings"
	"time"

	"github.com/golang-jwt/jwt/v4"
	"github.com/segmentio/ksuid"
	"github.com/sirupsen/logrus"
)

var (
	algRS256          = jwt.SigningMethodRS256.Alg()
	algRS384          = jwt.SigningMethodRS384.Alg()
	algRS512          = jwt.SigningMethodRS512.Alg()
	algEdDSA          = jwt.SigningMethodEdDSA.Alg()
	validMethods      = []string{algRS256, algRS384, algRS512, algEdDSA}
	hexEncodedMatcher = regexp.MustCompile("^[0-9a-fA-F]+$")
)

const (
	rsaKeyHeader      = "-----BEGIN RSA PRIVATE KEY"
	certHeader        = "-----BEGIN CERTIFICATE"
	pkHeader          = "-----BEGIN PUBLIC KEY"
	keyHeader         = "-----BEGIN PRIVATE KEY"
	defaultOrg        = "choria"
	OrgIssuerPrefix   = "I-"
	ChainIssuerPrefix = "C-"
	DefaultValidity   = time.Hour
)

var defaultIssuer = "Choria Tokens Package"

// Purpose indicates what kind of token a JWT is and helps us parse it into the right data structure
type Purpose string

// ErrorNotSignedByIssuer is raised when a token did not pass validation against the issuer chain
var ErrorNotSignedByIssuer = errors.New("not signed by issuer")

const (
	// UnknownPurpose is a JWT that does not have a purpose set
	UnknownPurpose Purpose = ""

	// ClientIDPurpose indicates a JWT is a ClientIDClaims JWT
	ClientIDPurpose Purpose = "choria_client_id"

	// ProvisioningPurpose indicates a JWT is a ProvisioningClaims JWT
	ProvisioningPurpose Purpose = "choria_provisioning"

	// ServerPurpose indicates a JWT is a ServerClaims JWT
	ServerPurpose Purpose = "choria_server"
)

// MapClaims are free form map claims
type MapClaims jwt.MapClaims

// ParseToken parses token into claims and verify the token is valid using the pk,
// if the token is signed by a chain issuer then pk must be the org issuer pk and
// the chain will be verified
func ParseToken(token string, claims jwt.Claims, pk any) error {
	if pk == nil {
		return fmt.Errorf("invalid public key")
	}

	_, err := jwt.ParseWithClaims(token, claims, func(t *jwt.Token) (any, error) {
		switch t.Method.Alg() {
		case algRS256, algRS512, algRS384:
			pk, ok := pk.(*rsa.PublicKey)
			if !ok {
				return nil, fmt.Errorf("rsa public key required")
			}
			return pk, nil

		case algEdDSA:
			pk, ok := pk.(ed25519.PublicKey)
			if !ok {
				return nil, fmt.Errorf("ed25519 public key required")
			}

			var sc *StandardClaims

			// if it's a client and from a chain we will verify it using the chain issuer pubk
			client, ok := claims.(*ClientIDClaims)
			if ok && strings.HasPrefix(client.Issuer, ChainIssuerPrefix) {
				sc = &client.StandardClaims
			}

			// if it's a server and from a chain we will verify it using the chain issuer pubk
			server, ok := claims.(*ServerClaims)
			if ok && strings.HasPrefix(server.Issuer, ChainIssuerPrefix) {
				sc = &server.StandardClaims
			}

			if sc != nil {
				valid, signerPk, err := sc.IsSignedByIssuer(pk)
				if err != nil {
					return nil, fmt.Errorf("%w: %s", ErrorNotSignedByIssuer, err)
				}
				if !valid {
					return nil, ErrorNotSignedByIssuer
				}
				pk = signerPk
			}

			return pk, nil

		default:
			return nil, fmt.Errorf("unsupported signing method %v in token", t.Method)
		}
	}, jwt.WithValidMethods(validMethods))
	if err != nil {
		return err
	}

	return nil
}

// ParseTokenUnverified parses token into claims and DOES not verify the token validity in any way
func ParseTokenUnverified(token string) (jwt.MapClaims, error) {
	parser := new(jwt.Parser)
	claims := new(jwt.MapClaims)
	_, _, err := parser.ParseUnverified(token, claims)
	return *claims, err
}

// TokenPurpose parses, without validating, token and checks for a Purpose field in it
func TokenPurpose(token string) Purpose {
	parser := new(jwt.Parser)
	claims := StandardClaims{}
	parser.ParseUnverified(token, &claims)

	if claims.Purpose == UnknownPurpose {
		if claims.RegisteredClaims.Subject == string(ProvisioningPurpose) {
			return ProvisioningPurpose
		}
	}

	return claims.Purpose
}

// TokenPurposeBytes called TokenPurpose with a bytes input
func TokenPurposeBytes(token []byte) Purpose {
	return TokenPurpose(string(token))
}

// TokenSigningAlgorithmBytes determines the signing algorithm used for a token
func TokenSigningAlgorithmBytes(token []byte) (string, error) {
	return TokenSigningAlgorithm(string(token))
}

// TokenSigningAlgorithm determines the signing algorithm used for a token
func TokenSigningAlgorithm(token string) (string, error) {
	parser := new(jwt.Parser)
	claims := StandardClaims{}
	t, _, err := parser.ParseUnverified(token, &claims)
	if err != nil {
		return "", err
	}

	return t.Method.Alg(), nil
}

// SignTokenWithKeyFile signs a JWT using an RSA Private Key in PEM format
func SignTokenWithKeyFile(claims jwt.Claims, pkFile string) (string, error) {
	keydat, err := os.ReadFile(pkFile)
	if err != nil {
		return "", fmt.Errorf("could not read signing key: %s", err)
	}

	if bytes.HasPrefix(keydat, []byte(rsaKeyHeader)) || bytes.HasPrefix(keydat, []byte(keyHeader)) {
		key, err := jwt.ParseRSAPrivateKeyFromPEM(keydat)
		if err != nil {
			return "", fmt.Errorf("could not parse signing key: %s", err)
		}

		return SignToken(claims, key)
	}

	if len(keydat) == ed25519.PrivateKeySize {
		seed, err := hex.DecodeString(string(keydat))
		if err != nil {
			return "", fmt.Errorf("invalid ed25519 seed file: %v", err)
		}

		return SignToken(claims, ed25519.NewKeyFromSeed(seed))
	}

	return "", fmt.Errorf("unsupported key in %v", pkFile)
}

// SignToken signs a JWT using an RSA Private Key
func SignToken(claims jwt.Claims, pk any) (string, error) {
	var stoken string
	var err error

	switch pri := pk.(type) {
	case ed25519.PrivateKey:
		token := jwt.NewWithClaims(jwt.SigningMethodEdDSA, claims)
		stoken, err = token.SignedString(pri)

	case *rsa.PrivateKey:
		token := jwt.NewWithClaims(jwt.SigningMethodRS256, claims)
		stoken, err = token.SignedString(pri)

	default:
		return "", fmt.Errorf("unsupported private key")
	}

	if err != nil {
		return "", fmt.Errorf("could not sign token using key: %s", err)
	}

	return stoken, nil
}

// SaveAndSignTokenWithKeyFile signs a token using SignTokenWithKeyFile and saves it to outFile
func SaveAndSignTokenWithKeyFile(claims jwt.Claims, pkFile string, outFile string, perm os.FileMode) error {
	token, err := SignTokenWithKeyFile(claims, pkFile)
	if err != nil {
		return err
	}

	return os.WriteFile(outFile, []byte(token), perm)
}

func getVaultIssuerPubKey(ctx context.Context, tlsc *tls.Config, key string, log *logrus.Entry) (ed25519.PublicKey, error) {
	vt := os.Getenv("VAULT_TOKEN")
	va := os.Getenv("VAULT_ADDR")

	if vt == "" || va == "" {
		return nil, fmt.Errorf("requires VAULT_TOKEN and VAULT_ADDR environment variables")
	}

	uri, err := url.Parse(va)
	if err != nil {
		return nil, err
	}

	uri.Path = fmt.Sprintf("/v1/transit/keys/%s", key)
	client := &http.Client{}
	if tlsc != nil {
		client.Transport = &http.Transport{TLSClientConfig: tlsc}
	}

	req, err := http.NewRequestWithContext(ctx, "GET", uri.String(), nil)
	if err != nil {
		return nil, err
	}
	req.Header.Add("X-Vault-Token", vt)

	resp, err := client.Do(req)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()

	body, err := io.ReadAll(resp.Body)
	if err != nil {
		return nil, err
	}

	if resp.StatusCode != 200 {
		return nil, fmt.Errorf("request failed: code: %d: %s", resp.StatusCode, string(body))
	}

	log.Debugf("JSON Response: %s", string(body))

	var vr struct {
		Data struct {
			Keys map[string]struct {
				PublicKey []byte `json:"public_key"`
			} `json:"keys"`
		} `json:"data"`
	}
	err = json.Unmarshal(body, &vr)
	if err != nil {
		return nil, err
	}

	if len(vr.Data.Keys) == 0 {
		return nil, fmt.Errorf("did not receive keys in response")
	}

	pk, ok := vr.Data.Keys["1"]
	if !ok {
		return nil, fmt.Errorf("did not receive keys in response")
	}

	if len(pk.PublicKey) != ed25519.PublicKeySize {
		return nil, fmt.Errorf("did not receive a valid public key in response")
	}

	return pk.PublicKey, nil
}

func signWithVault(ctx context.Context, tlsc *tls.Config, key string, ss []byte, log *logrus.Entry) ([]byte, error) {
	vt := os.Getenv("VAULT_TOKEN")
	va := os.Getenv("VAULT_ADDR")

	if vt == "" || va == "" {
		return nil, fmt.Errorf("requires VAULT_TOKEN and VAULT_ADDR environment variables")
	}

	uri, err := url.Parse(va)
	if err != nil {
		return nil, err
	}

	uri.Path = fmt.Sprintf("/v1/transit/sign/%s", key)

	dat := map[string]any{
		"signature_algorithm": "ed25519",
		"input":               base64.StdEncoding.EncodeToString(ss),
	}
	jdat, err := json.Marshal(dat)
	if err != nil {
		return nil, err
	}
	log.Debugf("JSON Request: %s", string(jdat))

	client := &http.Client{}
	if tlsc != nil {
		client.Transport = &http.Transport{TLSClientConfig: tlsc}
	}

	req, err := http.NewRequestWithContext(ctx, "POST", uri.String(), bytes.NewBuffer(jdat))
	if err != nil {
		return nil, err
	}
	req.Header.Add("X-Vault-Token", vt)

	resp, err := client.Do(req)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()

	body, err := io.ReadAll(resp.Body)
	if err != nil {
		return nil, err
	}

	if resp.StatusCode != 200 {
		return nil, fmt.Errorf("request failed: code: %d: %s", resp.StatusCode, string(body))
	}

	log.Debugf("JSON Response: %s", string(body))

	var vr struct {
		Data struct {
			Sig string `json:"signature"`
		} `json:"data"`
	}
	err = json.Unmarshal(body, &vr)
	if err != nil {
		return nil, err
	}

	if vr.Data.Sig == "" {
		return nil, fmt.Errorf("no signature in response: %s", string(body))
	}

	const vaultSigPrefix = "vault:v1:"

	if !strings.HasPrefix(vr.Data.Sig, vaultSigPrefix) {
		return nil, fmt.Errorf("invalid signature, no vault:v1 prefix")
	}

	signature, err := base64.StdEncoding.DecodeString(strings.TrimPrefix(vr.Data.Sig, vaultSigPrefix))
	if err != nil {
		return nil, fmt.Errorf("could not decode vault response: %w", err)
	}

	return signature, nil
}

// SaveAndSignTokenWithVault signs a token using the named key in a Vault Transit engine.  Requires VAULT_TOKEN and VAULT_ADDR to be set.
func SaveAndSignTokenWithVault(ctx context.Context, claims jwt.Claims, key string, outFile string, perm os.FileMode, tlsc *tls.Config, log *logrus.Entry) error {
	token := jwt.NewWithClaims(jwt.SigningMethodEdDSA, claims)
	ss, err := token.SigningString()
	if err != nil {
		return err
	}

	signature, err := signWithVault(ctx, tlsc, key, []byte(ss), log)
	if err != nil {
		return err
	}

	signed := fmt.Sprintf("%s.%s", ss, strings.TrimRight(base64.RawURLEncoding.EncodeToString(signature), "="))

	return os.WriteFile(outFile, []byte(signed), perm)
}

func newStandardClaims(issuer string, purpose Purpose, validity time.Duration, setSubject bool) (*StandardClaims, error) {
	if issuer == "" {
		issuer = defaultIssuer
	}

	now := jwt.NewNumericDate(time.Now().UTC())
	id, err := ksuid.NewRandomWithTime(now.Time)
	if err != nil {
		return nil, err
	}

	if validity == 0 {
		validity = DefaultValidity
	}

	claims := &StandardClaims{
		Purpose: purpose,
		RegisteredClaims: jwt.RegisteredClaims{
			ID:        id.String(),
			Issuer:    issuer,
			IssuedAt:  now,
			NotBefore: now,
			ExpiresAt: jwt.NewNumericDate(now.Add(validity)),
		},
	}

	if setSubject {
		claims.Subject = string(purpose)
	}

	return claims, nil
}

func readRSAOrED25519PublicData(dat []byte) (any, error) {
	var pk any
	var err error

	if bytes.HasPrefix(dat, []byte(certHeader)) || bytes.HasPrefix(dat, []byte(pkHeader)) {
		pk, err = jwt.ParseRSAPublicKeyFromPEM(dat)
		if err != nil {
			return nil, fmt.Errorf("could not parse validation certificate: %s", err)
		}
	} else {
		edpk, err := hex.DecodeString(string(dat))
		if err != nil {
			return nil, fmt.Errorf("could not parse ed25519 public data: %v", err)
		}
		if len(edpk) != ed25519.PublicKeySize {
			return nil, fmt.Errorf("invalid ed25519 public key size")
		}

		pk = ed25519.PublicKey(edpk)
	}

	return pk, nil
}

// NatsConnectionHelpers constructs token based private inbox and helpers for the nats.UserJWT() function. Only Server and Client tokens are supported.
func NatsConnectionHelpers(token string, collective string, seedFile string, log *logrus.Entry) (inbox string, jwth func() (string, error), sigh func([]byte) ([]byte, error), err error) {
	if collective == "" {
		return "", nil, nil, fmt.Errorf("collective is required")
	}

	if seedFile == "" {
		return "", nil, nil, fmt.Errorf("seedfile is required")
	}

	purpose := TokenPurpose(token)

	var uid string
	var isExp func() bool
	var exp time.Time

	switch purpose {
	case ClientIDPurpose:
		client, err := ParseClientIDTokenUnverified(token)
		if err != nil {
			return "", nil, nil, err
		}
		_, uid = client.UniqueID()
		isExp = client.IsExpired
		exp = client.ExpireTime()

	case ServerPurpose:
		server, err := ParseServerTokenUnverified(token)
		if err != nil {
			return "", nil, nil, err
		}
		_, uid = server.UniqueID()
		isExp = server.IsExpired
		exp = server.ExpireTime()

	default:
		return "", nil, nil, fmt.Errorf("unsupported token purpose: %v", purpose)
	}

	inbox = fmt.Sprintf("%s.reply.%s", collective, uid)

	jwth = func() (string, error) {
		if isExp() {
			log.Errorf("Cannot sign connection NONCE: token is expired by %v", time.Since(exp))
			return "", fmt.Errorf("token expired")
		}
		return token, nil
	}

	sigh = func(n []byte) ([]byte, error) {
		if isExp() {
			log.Errorf("Cannot sign connection NONCE: token is expired by %v", time.Since(exp))
			return nil, fmt.Errorf("token expired")
		}
		log.Debugf("Signing nonce using seed file %s", seedFile)
		return ed25519SignWithSeedFile(seedFile, n)
	}

	return inbox, jwth, sigh, nil
}

// IsEncodedEd25519Key determines if b holds valid characters for a hex encoded public key or seed
func IsEncodedEd25519Key(b []byte) bool {
	if len(b) != 64 {
		return false
	}

	return hexEncodedMatcher.Match(b)
}