1 files changed, 118 insertions, 0 deletions

diff --git a/jwks_fetcher.go b/jwks_fetcher.go
new file mode 100644
index 0000000..9925430
--- /dev/null
+++ b/jwks_fetcher.go
@@ -0,0 +1,118 @@
+package main
+
+import (
+        "github.com/pkg/errors"
+        "gopkg.in/square/go-jose.v2"
+        "log"
+        "net/url"
+        "time"
+)
+
+const (
+        REQUEST_BUFFER_SIZE      = 10
+        KEY_MAP_INITIAL_SIZE     = 5
+        DEFAULT_REFRESH_INTERVAL = 15 * time.Minute
+        MIN_REFRESH_INTERVAL     = 1 * time.Minute
+)
+
+type KeyRequest struct {
+        KeyId    string
+        Response chan *jose.JSONWebKey
+}
+
+type JWKSFetcher interface {
+        Run()
+        Fetch() error
+        GetKey(string) (*jose.JSONWebKey, error)
+        Done()
+}
+
+type jwksFetcher struct {
+        keyMap     map[string]jose.JSONWebKey
+        httpClient CautiousHTTPClient
+        validator  KeyValidator
+        fetchTimer *time.Timer
+        url        *url.URL
+        requests   chan *KeyRequest
+        done       chan bool
+}
+
+func NewJWKSFetcher(h CautiousHTTPClient, url *url.URL, issuer string, root string) JWKSFetcher {
+        val := NewKeyValidator(HostFromURL(issuer))
+        val.LoadRootPEM(root)
+
+        return &jwksFetcher{
+                httpClient: h,
+                validator:  val,
+                url:        url,
+                fetchTimer: time.NewTimer(DEFAULT_REFRESH_INTERVAL),
+                requests:   make(chan *KeyRequest, REQUEST_BUFFER_SIZE),
+                keyMap:     make(map[string]jose.JSONWebKey, KEY_MAP_INITIAL_SIZE),
+                done:       make(chan bool),
+        }
+}
+
+func (f *jwksFetcher) Fetch() error {
+        var jwks jose.JSONWebKeySet
+        timeout, err := f.httpClient.GetJSONExpires(f.url.String(), &jwks)
+        if err != nil {
+                return errors.WithStack(err)
+        }
+
+        for _, k := range jwks.Keys {
+                err = f.validator.Validate(k)
+                if err == nil {
+                        f.keyMap[k.KeyID] = k
+                } else {
+                        log.Printf("Rejecting key %q because %q", k.KeyID, err)
+                }
+        }
+
+        if timeout < MIN_REFRESH_INTERVAL {
+                timeout = MIN_REFRESH_INTERVAL
+        }
+
+        success := f.fetchTimer.Reset(timeout)
+        if !success {
+                f.fetchTimer = time.NewTimer(timeout)
+        }
+
+        return nil
+}
+
+func (f *jwksFetcher) Run() {
+        for {
+                select {
+                // Incoming request for a key, return key or nil in no key
+                case r := <-f.requests:
+                        if v, ok := f.keyMap[r.KeyId]; ok {
+                                r.Response <- &v
+                        } else {
+                                r.Response <- nil
+                        }
+                case <-f.fetchTimer.C:
+                        f.Fetch()
+                case <-f.done:
+                        return
+                }
+        }
+}
+
+func (f *jwksFetcher) Done() {
+        f.done <- true
+}
+
+func (f *jwksFetcher) GetKey(kid string) (*jose.JSONWebKey, error) {
+        r := &KeyRequest{
+                KeyId:    kid,
+                Response: make(chan *jose.JSONWebKey),
+        }
+
+        f.requests <- r
+
+        if res := <-r.Response; res == nil {
+                return nil, errors.Errorf("Key not found for ID")
+        } else {
+                return res, nil
+        }
+}