<commit_before>package token

import (
	"context"
	"crypto"
	"crypto/rsa"
	"crypto/x509"
	"encoding/base64"
	"encoding/pem"
	"errors"
	"strings"
	"time"

	_ "golang.org/x/crypto/blake2b" // depends on registering blake2b hash
	"gopkg.in/square/go-jose.v2"
	"gopkg.in/square/go-jose.v2/jwt"
	"xorkevin.dev/governor"
)

const (
	// ScopeAll grants all scopes to a token
	ScopeAll = "all"
)

const (
	// KindAccess is an access token kind
	KindAccess = "access"
	// KindRefresh is a refresh token kind
	KindRefresh = "refresh"
	// KindOAuthAccess is an oauth access token kind
	KindOAuthAccess = "oauth:access"
	// KindOAuthRefresh is an oauth refresh token kind
	KindOAuthRefresh = "oauth:refresh"
	// KindOAuthID is an openid id token kind
	KindOAuthID = "oauth:id"
)

const (
	pemBlockType = "PRIVATE KEY"
)

type (
	// Claims is a set of fields to describe a user
	Claims struct {
		jwt.Claims
		Kind     string `json:"kind"`
		AuthTime int64  `json:"auth_time,omitempty"`
		Scope    string `json:"scope,omitempty"`
		Key      string `json:"key,omitempty"`
	}

	// Tokenizer is a token generator
	Tokenizer interface {
		GetJWKS() *jose.JSONWebKeySet
		Generate(kind string, userid string, duration int64, id string, authTime int64, scope string, key string) (string, *Claims, error)
		GenerateExt(kind string, issuer string, userid string, audience []string, duration int64, id string, authTime int64, claims interface{}) (string, error)
		Validate(kind string, tokenString string) (bool, *Claims)
		GetClaims(kind string, tokenString string) (bool, *Claims)
		GetClaimsExt(kind string, tokenString string, audience []string, claims interface{}) (bool, *Claims)
	}

	// Service is a Tokenizer and governor.Service
	Service interface {
		governor.Service
		Tokenizer
	}

	service struct {
		secret     []byte
		privateKey crypto.Signer
		publicKey  crypto.PublicKey
		issuer     string
		audience   string
		signer     jose.Signer
		keySigner  jose.Signer
		jwk        *jose.JSONWebKey
		logger     governor.Logger
	}

	ctxKeyTokenizer struct{}
)

// GetCtxTokenizer returns a Tokenizer from the context
func GetCtxTokenizer(inj governor.Injector) Tokenizer {
	v := inj.Get(ctxKeyTokenizer{})
	if v == nil {
		return nil
	}
	return v.(Tokenizer)
}

// setCtxTokenizer sets a Tokenizer in the context
func setCtxTokenizer(inj governor.Injector, t Tokenizer) {
	inj.Set(ctxKeyTokenizer{}, t)
}

// New creates a new Tokenizer
func New() Service {
	return &service{
		secret:   nil,
		issuer:   "",
		audience: "",
		signer:   nil,
	}
}

func (s *service) Register(inj governor.Injector, r governor.ConfigRegistrar, jr governor.JobRegistrar) {
	setCtxTokenizer(inj, s)

	r.SetDefault("tokensecret", "")
	r.SetDefault("issuer", "governor")
	r.SetDefault("audience", "governor")
}

type (
	// ErrSigner is returned when failing to create a signer
	ErrSigner struct{}
	// ErrGenerate is returned when failing to generate a token
	ErrGenerate struct{}
)

func (e ErrSigner) Error() string {
	return "Error creating signer"
}

func (e ErrGenerate) Error() string {
	return "Error generating token"
}

func (s *service) Init(ctx context.Context, c governor.Config, r governor.ConfigReader, l governor.Logger, m governor.Router) error {
	s.logger = l
	l = s.logger.WithData(map[string]string{
		"phase": "init",
	})

	tokensecret, err := r.GetSecret("tokensecret")
	if err != nil {
		return governor.ErrWithMsg(err, "Failed to read token secret")
	}
	secret, ok := tokensecret["secret"].(string)
	if !ok {
		return governor.ErrWithKind(nil, governor.ErrInvalidConfig{}, "Invalid secret")
	}
	if secret == "" {
		return governor.ErrWithKind(nil, governor.ErrInvalidConfig{}, "Token secret is not set")
	}
	s.secret = []byte(secret)
	sig, err := jose.NewSigner(jose.SigningKey{Algorithm: jose.HS512, Key: s.secret}, (&jose.SignerOptions{}).WithType("JWT"))
	if err != nil {
		return governor.ErrWithKind(err, ErrSigner{}, "Failed to create new jwt signer")
	}
	s.signer = sig

	rsakeysecret, err := r.GetSecret("rsakey")
	if err != nil {
		return governor.ErrWithMsg(err, "Failed to read rsakey")
	}
	rsakeyPem, ok := rsakeysecret["secret"].(string)
	if !ok {
		return governor.ErrWithKind(nil, governor.ErrInvalidConfig{}, "Invalid rsakey secret")
	}
	pemBlock, _ := pem.Decode([]byte(rsakeyPem))
	if pemBlock == nil || pemBlock.Type != pemBlockType {
		return governor.ErrWithKind(nil, governor.ErrInvalidConfig{}, "Invalid rsakey pem")
	}
	rawKey, err := x509.ParsePKCS8PrivateKey(pemBlock.Bytes)
	if err != nil {
		return governor.ErrWithKind(err, governor.ErrInvalidConfig{}, "Invalid rsakey pkcs8")
	}
	key, ok := rawKey.(*rsa.PrivateKey)
	if !ok {
		return errors.New("Failed to create rsa key")
	}
	key.Precompute()
	s.privateKey = key
	s.publicKey = key.Public()
	keySig, err := jose.NewSigner(jose.SigningKey{Algorithm: jose.RS256, Key: s.privateKey}, (&jose.SignerOptions{}).WithType("JWT"))
	if err != nil {
		return governor.ErrWithKind(err, ErrSigner{}, "Failed to create new jwt RS256 signer")
	}
	s.keySigner = keySig

	jwk := &jose.JSONWebKey{
		Key:       s.publicKey,
		Algorithm: "RS256",
		Use:       "sig",
	}
	kid, err := jwk.Thumbprint(crypto.BLAKE2b_512)
	if err != nil {
		return governor.ErrWithKind(err, ErrSigner{}, "Failed to calculate jwk thumbprint")
	}
	jwk.KeyID = base64.RawURLEncoding.EncodeToString(kid)
	s.jwk = jwk

	issuer := r.GetStr("issuer")
	if issuer == "" {
		return governor.ErrWithKind(nil, governor.ErrInvalidConfig{}, "Token issuer is not set")
	}
	s.issuer = issuer

	audience := r.GetStr("audience")
	if audience == "" {
		return governor.ErrWithKind(nil, governor.ErrInvalidConfig{}, "Token audience is not set")
	}
	s.audience = audience
<commit_msg>Add kid to id token signer<commit_after>package token

import (
	"context"
	"crypto"
	"crypto/rsa"
	"crypto/x509"
	"encoding/base64"
	"encoding/pem"
	"errors"
	"strings"
	"time"

	_ "golang.org/x/crypto/blake2b" // depends on registering blake2b hash
	"gopkg.in/square/go-jose.v2"
	"gopkg.in/square/go-jose.v2/jwt"
	"xorkevin.dev/governor"
)

const (
	// ScopeAll grants all scopes to a token
	ScopeAll = "all"
)

const (
	// KindAccess is an access token kind
	KindAccess = "access"
	// KindRefresh is a refresh token kind
	KindRefresh = "refresh"
	// KindOAuthAccess is an oauth access token kind
	KindOAuthAccess = "oauth:access"
	// KindOAuthRefresh is an oauth refresh token kind
	KindOAuthRefresh = "oauth:refresh"
	// KindOAuthID is an openid id token kind
	KindOAuthID = "oauth:id"
)

const (
	pemBlockType = "PRIVATE KEY"
)

const (
	jwtHeaderKid = "kid"
	jwtHeaderJWT = "JWT"
)

type (
	// Claims is a set of fields to describe a user
	Claims struct {
		jwt.Claims
		Kind     string `json:"kind"`
		AuthTime int64  `json:"auth_time,omitempty"`
		Scope    string `json:"scope,omitempty"`
		Key      string `json:"key,omitempty"`
	}

	// Tokenizer is a token generator
	Tokenizer interface {
		GetJWKS() *jose.JSONWebKeySet
		Generate(kind string, userid string, duration int64, id string, authTime int64, scope string, key string) (string, *Claims, error)
		GenerateExt(kind string, issuer string, userid string, audience []string, duration int64, id string, authTime int64, claims interface{}) (string, error)
		Validate(kind string, tokenString string) (bool, *Claims)
		GetClaims(kind string, tokenString string) (bool, *Claims)
		GetClaimsExt(kind string, tokenString string, audience []string, claims interface{}) (bool, *Claims)
	}

	// Service is a Tokenizer and governor.Service
	Service interface {
		governor.Service
		Tokenizer
	}

	service struct {
		secret     []byte
		privateKey crypto.Signer
		publicKey  crypto.PublicKey
		issuer     string
		audience   string
		signer     jose.Signer
		keySigner  jose.Signer
		jwk        *jose.JSONWebKey
		logger     governor.Logger
	}

	ctxKeyTokenizer struct{}
)

// GetCtxTokenizer returns a Tokenizer from the context
func GetCtxTokenizer(inj governor.Injector) Tokenizer {
	v := inj.Get(ctxKeyTokenizer{})
	if v == nil {
		return nil
	}
	return v.(Tokenizer)
}

// setCtxTokenizer sets a Tokenizer in the context
func setCtxTokenizer(inj governor.Injector, t Tokenizer) {
	inj.Set(ctxKeyTokenizer{}, t)
}

// New creates a new Tokenizer
func New() Service {
	return &service{
		secret:   nil,
		issuer:   "",
		audience: "",
		signer:   nil,
	}
}

func (s *service) Register(inj governor.Injector, r governor.ConfigRegistrar, jr governor.JobRegistrar) {
	setCtxTokenizer(inj, s)

	r.SetDefault("tokensecret", "")
	r.SetDefault("issuer", "governor")
	r.SetDefault("audience", "governor")
}

type (
	// ErrSigner is returned when failing to create a signer
	ErrSigner struct{}
	// ErrGenerate is returned when failing to generate a token
	ErrGenerate struct{}
)

func (e ErrSigner) Error() string {
	return "Error creating signer"
}

func (e ErrGenerate) Error() string {
	return "Error generating token"
}

func (s *service) Init(ctx context.Context, c governor.Config, r governor.ConfigReader, l governor.Logger, m governor.Router) error {
	s.logger = l
	l = s.logger.WithData(map[string]string{
		"phase": "init",
	})

	tokensecret, err := r.GetSecret("tokensecret")
	if err != nil {
		return governor.ErrWithMsg(err, "Failed to read token secret")
	}
	secret, ok := tokensecret["secret"].(string)
	if !ok {
		return governor.ErrWithKind(nil, governor.ErrInvalidConfig{}, "Invalid secret")
	}
	if secret == "" {
		return governor.ErrWithKind(nil, governor.ErrInvalidConfig{}, "Token secret is not set")
	}
	s.secret = []byte(secret)
	sig, err := jose.NewSigner(jose.SigningKey{Algorithm: jose.HS512, Key: s.secret}, (&jose.SignerOptions{}).WithType(jwtHeaderJWT))
	if err != nil {
		return governor.ErrWithKind(err, ErrSigner{}, "Failed to create new jwt signer")
	}
	s.signer = sig

	rsakeysecret, err := r.GetSecret("rsakey")
	if err != nil {
		return governor.ErrWithMsg(err, "Failed to read rsakey")
	}
	rsakeyPem, ok := rsakeysecret["secret"].(string)
	if !ok {
		return governor.ErrWithKind(nil, governor.ErrInvalidConfig{}, "Invalid rsakey secret")
	}
	pemBlock, _ := pem.Decode([]byte(rsakeyPem))
	if pemBlock == nil || pemBlock.Type != pemBlockType {
		return governor.ErrWithKind(nil, governor.ErrInvalidConfig{}, "Invalid rsakey pem")
	}
	rawKey, err := x509.ParsePKCS8PrivateKey(pemBlock.Bytes)
	if err != nil {
		return governor.ErrWithKind(err, governor.ErrInvalidConfig{}, "Invalid rsakey pkcs8")
	}
	key, ok := rawKey.(*rsa.PrivateKey)
	if !ok {
		return errors.New("Failed to create rsa key")
	}
	key.Precompute()
	s.privateKey = key
	s.publicKey = key.Public()

	jwk := &jose.JSONWebKey{
		Key:       s.publicKey,
		Algorithm: "RS256",
		Use:       "sig",
	}
	kid, err := jwk.Thumbprint(crypto.BLAKE2b_512)
	if err != nil {
		return governor.ErrWithKind(err, ErrSigner{}, "Failed to calculate jwk thumbprint")
	}
	jwk.KeyID = base64.RawURLEncoding.EncodeToString(kid)
	s.jwk = jwk

	keySig, err := jose.NewSigner(jose.SigningKey{Algorithm: jose.RS256, Key: s.privateKey}, (&jose.SignerOptions{}).WithType(jwtHeaderJWT).WithHeader(jwtHeaderKid, jwk.KeyID))
	if err != nil {
		return governor.ErrWithKind(err, ErrSigner{}, "Failed to create new jwt RS256 signer")
	}
	s.keySigner = keySig

	issuer := r.GetStr("issuer")
	if issuer == "" {
		return governor.ErrWithKind(nil, governor.ErrInvalidConfig{}, "Token issuer is not set")
	}
	s.issuer = issuer

	audience := r.GetStr("audience")
	if audience == "" {
		return governor.ErrWithKind(nil, governor.ErrInvalidConfig{}, "Token audience is not set")
	}
	s.audience = audience

{'fraction_non_alphanumeric': 0.12577140373750542, 'fraction_numerical': 0.007301173402868318, 'mean_word_length': 4.287683823529412, 'pattern_counts': {'":': 2, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3280120', 'n_tokens_mistral': 4228, 'n_tokens_neox': 4052, 'n_words': 1063}