Skip to content

Commit

Permalink
More documentation, and alias function for loading a KeyStore from a …
Browse files Browse the repository at this point in the history
…filepath path
  • Loading branch information
mrmelon54 committed Aug 12, 2024
1 parent 7eaf420 commit 87774ec
Show file tree
Hide file tree
Showing 4 changed files with 58 additions and 17 deletions.
2 changes: 1 addition & 1 deletion README.md
Original file line number Diff line number Diff line change
@@ -1,3 +1,3 @@
# MJWT

A simple wrapper for JWT. Contains an AccessToken and RefreshToken model.
A simple wrapper for JWT. Contains an AccessToken and RefreshToken model.
8 changes: 8 additions & 0 deletions issuer.go
Original file line number Diff line number Diff line change
Expand Up @@ -7,17 +7,21 @@ import (
"time"
)

// Issuer provides the signing for a PrivateKey identified by the KID in the
// provided KeyStore
type Issuer struct {
issuer string
kid string
signing jwt.SigningMethod
keystore *KeyStore
}

// NewIssuer creates an Issuer with an empty KeyStore
func NewIssuer(name, kid string, signing jwt.SigningMethod) (*Issuer, error) {
return NewIssuerWithKeyStore(name, kid, signing, NewKeyStore())
}

// NewIssuerWithKeyStore creates an Issuer with a provided KeyStore
func NewIssuerWithKeyStore(name, kid string, signing jwt.SigningMethod, keystore *KeyStore) (*Issuer, error) {
i := &Issuer{name, kid, signing, keystore}
if i.keystore.HasPrivateKey(kid) {
Expand All @@ -31,10 +35,12 @@ func NewIssuerWithKeyStore(name, kid string, signing jwt.SigningMethod, keystore
return i, i.keystore.SaveSingleKey(kid)
}

// GenerateJwt produces a signed JWT in string form
func (i *Issuer) GenerateJwt(sub, id string, aud jwt.ClaimStrings, dur time.Duration, claims Claims) (string, error) {
return i.SignJwt(wrapClaims[Claims](sub, id, i.issuer, aud, dur, claims))
}

// SignJwt produces a signed JWT in string form from a raw jwt.Claims structure
func (i *Issuer) SignJwt(wrapped jwt.Claims) (string, error) {
key, err := i.PrivateKey()
if err != nil {
Expand All @@ -45,10 +51,12 @@ func (i *Issuer) SignJwt(wrapped jwt.Claims) (string, error) {
return token.SignedString(key)
}

// PrivateKey outputs the rsa.PrivateKey from the KID of the Issuer
func (i *Issuer) PrivateKey() (*rsa.PrivateKey, error) {
return i.keystore.GetPrivateKey(i.kid)
}

// KeyStore outputs the underlying KeyStore used by the Issuer
func (i *Issuer) KeyStore() *KeyStore {
return i.keystore
}
1 change: 1 addition & 0 deletions jwks.go
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@ import (
"io"
)

// WriteJwkSetJson outputs the public keys used by the Issuers
func WriteJwkSetJson(w io.Writer, issuers []*Issuer) error {
enc := json.NewEncoder(w)
enc.SetIndent("", " ")
Expand Down
64 changes: 48 additions & 16 deletions keystore.go
Original file line number Diff line number Diff line change
Expand Up @@ -26,25 +26,43 @@ const PemExt = ".pem"
const PrivatePemExt = PrivateStr + PemExt
const PublicPemExt = PublicStr + PemExt

// KeyStore provides a store for a collection of private/public keypair structs
type KeyStore struct {
mu *sync.RWMutex
store map[string]*keyPair
dir afero.Fs
}

// NewKeyStore creates an empty KeyStore
func NewKeyStore() *KeyStore {
return &KeyStore{
mu: new(sync.RWMutex),
store: make(map[string]*keyPair),
}
}

// NewKeyStoreWithDir creates an empty KeyStore with an underlying afero.Fs
// filesystem for saving the internal store data
func NewKeyStoreWithDir(dir afero.Fs) *KeyStore {
keyStore := NewKeyStore()
keyStore.dir = dir
return keyStore
}

// NewKeyStoreFromPath creates an empty KeyStore. The provided path is walked to
// load the private/public keys. See implementation in NewKeyStoreFromDir.
func NewKeyStoreFromPath(dir string) (*KeyStore, error) {
abs, err := filepath.Abs(dir)
if err != nil {
return nil, err
}
return NewKeyStoreFromDir(afero.NewBasePathFs(afero.NewOsFs(), abs))
}

// NewKeyStoreFromDir creates an empty KeyStore. The provided afero.Fs is walked
// to find all private/public keys in files named `.private.pem` and
// `.public.pem` respectively. The keys are loaded into the KeyStore and any
// errors are returned immediately.
func NewKeyStoreFromDir(dir afero.Fs) (*KeyStore, error) {
keyStore := NewKeyStoreWithDir(dir)
err := afero.Walk(dir, ".", func(path string, d fs.FileInfo, err error) error {
Expand Down Expand Up @@ -94,6 +112,7 @@ type keyPair struct {
public *rsa.PublicKey
}

// LoadPrivateKey sets the rsa.PrivateKey/rsa.PublicKey for the KID
func (k *KeyStore) LoadPrivateKey(kid string, key *rsa.PrivateKey) {
k.mu.Lock()
if k.store[kid] == nil {
Expand All @@ -104,6 +123,7 @@ func (k *KeyStore) LoadPrivateKey(kid string, key *rsa.PrivateKey) {
k.mu.Unlock()
}

// LoadPublicKey sets the rsa.PublicKey for the KID
func (k *KeyStore) LoadPublicKey(kid string, key *rsa.PublicKey) {
k.mu.Lock()
if k.store[kid] == nil {
Expand All @@ -113,12 +133,14 @@ func (k *KeyStore) LoadPublicKey(kid string, key *rsa.PublicKey) {
k.mu.Unlock()
}

// RemoveKey deletes the KID keypair from the KeyStore
func (k *KeyStore) RemoveKey(kid string) {
k.mu.Lock()
delete(k.store, kid)
k.mu.Unlock()
}

// ListKeys provides a slice of the KIDs for all keys loaded in the KeyStore
func (k *KeyStore) ListKeys() []string {
k.mu.RLock()
defer k.mu.RUnlock()
Expand All @@ -129,6 +151,7 @@ func (k *KeyStore) ListKeys() []string {
return keys
}

// GetPrivateKey outputs the rsa.PrivateKey for the KID from the KeyStore
func (k *KeyStore) GetPrivateKey(kid string) (*rsa.PrivateKey, error) {
k.mu.RLock()
defer k.mu.RUnlock()
Expand All @@ -138,6 +161,7 @@ func (k *KeyStore) GetPrivateKey(kid string) (*rsa.PrivateKey, error) {
return k.store[kid].private, nil
}

// GetPublicKey outputs the rsa.PublicKey for the KID from the KeyStore
func (k *KeyStore) GetPublicKey(kid string) (*rsa.PublicKey, error) {
k.mu.RLock()
defer k.mu.RUnlock()
Expand All @@ -147,12 +171,15 @@ func (k *KeyStore) GetPublicKey(kid string) (*rsa.PublicKey, error) {
return k.store[kid].public, nil
}

// ClearKeys clears the internal map and makes a new map to release used memory
func (k *KeyStore) ClearKeys() {
k.mu.Lock()
clear(k.store)
k.store = make(map[string]*keyPair)
k.mu.Unlock()
}

// HasPrivateKey outputs true if the KID is found in the KeyStore
func (k *KeyStore) HasPrivateKey(kid string) bool {
k.mu.RLock()
defer k.mu.RUnlock()
Expand All @@ -164,6 +191,7 @@ func (k *KeyStore) internalHasPrivateKey(kid string) bool {
return v != nil && v.private != nil
}

// HasPublicKey outputs true if the KID is found in the KeyStore
func (k *KeyStore) HasPublicKey(kid string) bool {
k.mu.RLock()
defer k.mu.RUnlock()
Expand All @@ -175,6 +203,9 @@ func (k *KeyStore) internalHasPublicKey(kid string) bool {
return v != nil && v.public != nil
}

// VerifyJwt parses the provided token string and validates it against the KID
// using the KeyStore. An error is returned if the token fails to parse or if
// there is no matching KID in the KeyStore.
func (k *KeyStore) VerifyJwt(token string, claims baseTypeClaim) (*jwt.Token, error) {
withClaims, err := jwt.ParseWithClaims(token, claims, func(token *jwt.Token) (interface{}, error) {
kid, ok := token.Header["kid"].(string)
Expand All @@ -189,6 +220,8 @@ func (k *KeyStore) VerifyJwt(token string, claims baseTypeClaim) (*jwt.Token, er
return withClaims, claims.Valid()
}

// SaveSingleKey writes the rsa.PrivateKey/rsa.PublicKey for the requested KID to
// the underlying afero.Fs.
func (k *KeyStore) SaveSingleKey(kid string) error {
if k.dir == nil {
return nil
Expand All @@ -201,16 +234,11 @@ func (k *KeyStore) SaveSingleKey(kid string) error {
return ErrMissingKeyPair
}

var errs []error
if pair.private != nil {
errs = append(errs, afero.WriteFile(k.dir, kid+PrivatePemExt, rsaprivate.Encode(pair.private), 0600))
}
if pair.public != nil {
errs = append(errs, afero.WriteFile(k.dir, kid+PublicPemExt, rsapublic.Encode(pair.public), 0600))
}
return errors.Join(errs...)
return writeSingleKey(k.dir, kid, pair)
}

// SaveKeys writes the rsa.PrivateKey/rsa.PublicKey for the requested KID to the
// underlying afero.Fs.
func (k *KeyStore) SaveKeys() error {
k.mu.RLock()
defer k.mu.RUnlock()
Expand All @@ -219,15 +247,19 @@ func (k *KeyStore) SaveKeys() error {
workers.SetLimit(runtime.NumCPU())
for kid, pair := range k.store {
workers.Go(func() error {
var errs []error
if pair.private != nil {
errs = append(errs, afero.WriteFile(k.dir, kid+PrivatePemExt, rsaprivate.Encode(pair.private), 0600))
}
if pair.public != nil {
errs = append(errs, afero.WriteFile(k.dir, kid+PublicPemExt, rsapublic.Encode(pair.public), 0600))
}
return errors.Join(errs...)
return writeSingleKey(k.dir, kid, pair)
})
}
return workers.Wait()
}

func writeSingleKey(dir afero.Fs, kid string, pair *keyPair) error {
var errs []error
if pair.private != nil {
errs = append(errs, afero.WriteFile(dir, kid+PrivatePemExt, rsaprivate.Encode(pair.private), 0600))
}
if pair.public != nil {
errs = append(errs, afero.WriteFile(dir, kid+PublicPemExt, rsapublic.Encode(pair.public), 0600))
}
return errors.Join(errs...)
}

0 comments on commit 87774ec

Please sign in to comment.