diff --git a/go/src/glide.lock b/go/src/glide.lock index 20c66ac529d..b63b7c2b85a 100644 --- a/go/src/glide.lock +++ b/go/src/glide.lock @@ -557,7 +557,7 @@ imports: - helpers - services - name: github.com/koding/kite - version: 1ed804424233a8f6d7ec9fdef092b16200101617 + version: 2c8b7e6356aeb55ffafa345708d3699149d6a207 subpackages: - config - dnode diff --git a/go/src/vendor/github.com/koding/kite/kontrol/handlers.go b/go/src/vendor/github.com/koding/kite/kontrol/handlers.go index 5e9d6930827..7c85ea0c306 100644 --- a/go/src/vendor/github.com/koding/kite/kontrol/handlers.go +++ b/go/src/vendor/github.com/koding/kite/kontrol/handlers.go @@ -165,6 +165,7 @@ func (k *Kontrol) HandleRegister(r *kite.Request) (interface{}, error) { func (k *Kontrol) HandleGetKites(r *kite.Request) (interface{}, error) { var args protocol.GetKitesArgs + if err := r.Args.One().Unmarshal(&args); err != nil { return nil, err } @@ -176,17 +177,21 @@ func (k *Kontrol) HandleGetKites(r *kite.Request) (interface{}, error) { } for _, kite := range kites { - // audience will go into the token as "aud" claim. - audience := getAudience(args.Query) - keyPair, err := k.getOrUpdateKeyID(kite.KeyID, r) if err != nil { return nil, err } + tok := &token{ + audience: getAudience(args.Query), + username: r.Username, + issuer: k.Kite.Kite().Username, + keyPair: keyPair, + } + // Generate token once here because we are using the same token for every // kite we return and generating many tokens is really slow. - token, err := k.generateToken(audience, r.Username, k.Kite.Kite().Username, keyPair) + token, err := k.generateToken(tok) if err != nil { return nil, err } @@ -200,14 +205,14 @@ func (k *Kontrol) HandleGetKites(r *kite.Request) (interface{}, error) { } func (k *Kontrol) HandleGetToken(r *kite.Request) (interface{}, error) { - var query *protocol.KontrolQuery + var args protocol.GetTokenArgs - if err := r.Args.One().Unmarshal(&query); err != nil { + if err := r.Args.One().Unmarshal(&args); err != nil { return nil, fmt.Errorf("invalid query: %s", err) } // check if it's exist - kites, err := k.storage.Get(query) + kites, err := k.storage.Get(&args.KontrolQuery) if err != nil { return nil, err } @@ -226,9 +231,14 @@ func (k *Kontrol) HandleGetToken(r *kite.Request) (interface{}, error) { if err != nil { return nil, err } - audience := getAudience(query) - return k.generateToken(audience, r.Username, k.Kite.Kite().Username, keyPair) + return k.generateToken(&token{ + audience: getAudience(&args.KontrolQuery), + username: r.Username, + issuer: k.Kite.Kite().Username, + keyPair: keyPair, + force: args.Force, + }) } func (k *Kontrol) HandleMachine(r *kite.Request) (interface{}, error) { diff --git a/go/src/vendor/github.com/koding/kite/kontrol/kontrol.go b/go/src/vendor/github.com/koding/kite/kontrol/kontrol.go index 5e5c0be746d..06b9115a8b0 100644 --- a/go/src/vendor/github.com/koding/kite/kontrol/kontrol.go +++ b/go/src/vendor/github.com/koding/kite/kontrol/kontrol.go @@ -89,7 +89,7 @@ type Kontrol struct { heartbeats map[string]*heartbeat heartbeatsMu sync.Mutex // protects each clients heartbeat timer - tokenCache map[string]string + tokenCache map[string]cachedToken tokenCacheMu sync.Mutex // closed notifies goroutines started by kontrol that it got closed @@ -172,7 +172,7 @@ func NewWithoutHandlers(conf *config.Config, version string) *Kontrol { clientLocks: NewIdlock(), heartbeats: make(map[string]*heartbeat), closed: make(chan struct{}), - tokenCache: make(map[string]string), + tokenCache: make(map[string]cachedToken), } // Make a copy to not modify user-provided value. @@ -506,20 +506,59 @@ func (k *Kontrol) tokenLeeway() time.Duration { return TokenLeeway } +type token struct { + audience string + username string + issuer string + keyPair *KeyPair + force bool +} + +type cachedToken struct { + signed string + timer *time.Timer +} + +func (t *token) String() string { + return t.audience + t.username + t.issuer + t.keyPair.ID +} + +// cacheToken cached the signed token under the given key. +// +// It also ensures the token is invalidated after its expiration time. +// +// If the token was already exists in the cache, it will be +// overwritten with a new value. +func (k *Kontrol) cacheToken(key, signed string) { + if ct, ok := k.tokenCache[key]; ok { + ct.timer.Stop() + } + + k.tokenCache[key] = cachedToken{ + signed: signed, + timer: time.AfterFunc(k.tokenTTL()-k.tokenLeeway(), func() { + k.tokenCacheMu.Lock() + delete(k.tokenCache, key) + k.tokenCacheMu.Unlock() + }), + } +} + // generateToken returns a JWT token string. Please see the URL for details: // http://tools.ietf.org/html/draft-ietf-oauth-json-web-token-13#section-4.1 -func (k *Kontrol) generateToken(aud, username, issuer string, kp *KeyPair) (string, error) { - uniqKey := aud + username + issuer + kp.ID +func (k *Kontrol) generateToken(tok *token) (string, error) { + uniqKey := tok.String() k.tokenCacheMu.Lock() defer k.tokenCacheMu.Unlock() - signed, ok := k.tokenCache[uniqKey] - if ok { - return signed, nil + if !tok.force { + if ct, ok := k.tokenCache[uniqKey]; ok { + return ct.signed, nil + } } - rsaPrivate, err := jwt.ParseRSAPrivateKeyFromPEM([]byte(kp.Private)) + rsaPrivate, err := jwt.ParseRSAPrivateKeyFromPEM([]byte(tok.keyPair.Private)) if err != nil { return "", err } @@ -528,9 +567,9 @@ func (k *Kontrol) generateToken(aud, username, issuer string, kp *KeyPair) (stri claims := &kitekey.KiteClaims{ StandardClaims: jwt.StandardClaims{ - Issuer: issuer, - Subject: username, - Audience: aud, + Issuer: tok.issuer, + Subject: tok.username, + Audience: tok.audience, ExpiresAt: now.Add(k.tokenTTL()).Add(k.tokenLeeway()).UTC().Unix(), IssuedAt: now.Add(-k.tokenLeeway()).UTC().Unix(), Id: uuid.NewV4().String(), @@ -541,23 +580,12 @@ func (k *Kontrol) generateToken(aud, username, issuer string, kp *KeyPair) (stri claims.NotBefore = now.Add(-k.tokenLeeway()).Unix() } - signed, err = jwt.NewWithClaims(jwt.GetSigningMethod("RS256"), claims).SignedString(rsaPrivate) + signed, err := jwt.NewWithClaims(jwt.GetSigningMethod("RS256"), claims).SignedString(rsaPrivate) if err != nil { return "", errors.New("Server error: Cannot generate a token") } - // cache our token - k.tokenCache[uniqKey] = signed - - // cache invalidation, because we cache the token in tokenCache we need to - // invalidate it expiration time. This was handled usually within JWT, but - // now we have to do it manually for our own cache. - time.AfterFunc(TokenTTL-TokenLeeway, func() { - k.tokenCacheMu.Lock() - defer k.tokenCacheMu.Unlock() - - delete(k.tokenCache, uniqKey) - }) + k.cacheToken(uniqKey, signed) return signed, nil } diff --git a/go/src/vendor/github.com/koding/kite/kontrol/kontrol_test.go b/go/src/vendor/github.com/koding/kite/kontrol/kontrol_test.go index 9537850202f..16e2deeecd5 100644 --- a/go/src/vendor/github.com/koding/kite/kontrol/kontrol_test.go +++ b/go/src/vendor/github.com/koding/kite/kontrol/kontrol_test.go @@ -572,7 +572,7 @@ func TestKontrol(t *testing.T) { // Test Kontrol.GetToken // TODO(rjeczalik): rework test to not touch Kontrol internals kon.tokenCacheMu.Lock() - kon.tokenCache = make(map[string]string) + kon.tokenCache = make(map[string]cachedToken) kon.tokenCacheMu.Unlock() _, err = exp2Kite.GetToken(&remoteMathWorker.Kite) @@ -711,7 +711,7 @@ func TestKontrolMultiKey(t *testing.T) { // Test Kontrol.GetToken // TODO(rjeczalik): rework test to not touch Kontrol internals kon.tokenCacheMu.Lock() - kon.tokenCache = make(map[string]string) // empty it + kon.tokenCache = make(map[string]cachedToken) // empty it kon.tokenCacheMu.Unlock() newToken, err := exp3Kite.GetToken(&remoteMathWorker.Kite) diff --git a/go/src/vendor/github.com/koding/kite/protocol/protocol.go b/go/src/vendor/github.com/koding/kite/protocol/protocol.go index 29a2a35788e..52ad5986b98 100644 --- a/go/src/vendor/github.com/koding/kite/protocol/protocol.go +++ b/go/src/vendor/github.com/koding/kite/protocol/protocol.go @@ -149,6 +149,13 @@ type GetKitesArgs struct { Who json.RawMessage `json:"who"` } +// GetTokenArgs is a request value for the "getToken" kontrol method. +type GetTokenArgs struct { + KontrolQuery // kite to generate a token for + + Force bool `json:"force"` // force creation of a new token +} + type WhoResult struct { Query *KontrolQuery `json:"query"` }