From 5a3ab973e602d824c0d3ef5164f815a5ca528ec5 Mon Sep 17 00:00:00 2001 From: Armon Dadgar Date: Thu, 9 Apr 2015 12:29:13 -0700 Subject: [PATCH] vault: Simplify common lease logic --- logical/lease.go | 23 ++++++++++++++++ vault/expiration.go | 64 ++++++++++++++++++--------------------------- 2 files changed, 48 insertions(+), 39 deletions(-) diff --git a/logical/lease.go b/logical/lease.go index 8999307d8..69877eca7 100644 --- a/logical/lease.go +++ b/logical/lease.go @@ -24,3 +24,26 @@ type LeaseOptions struct { // a logical backend. LeaseIssue time.Time `json:"-"` } + +// LeaseEnabled checks if leasing is enabled +func (l *LeaseOptions) LeaseEnabled() bool { + return l.Lease > 0 +} + +// LeaseTotal is the total lease time including the grace period +func (l *LeaseOptions) LeaseTotal() time.Duration { + if l.Lease == 0 { + return 0 + } + return l.Lease + l.LeaseGracePeriod +} + +// ExpirationTime computes the time until expiration including the grace period +func (l *LeaseOptions) ExpirationTime() time.Time { + var expireTime time.Time + if l.Lease > 0 { + leaseTotal := l.Lease + l.LeaseGracePeriod + expireTime = time.Now().UTC().Add(leaseTotal) + } + return expireTime +} diff --git a/vault/expiration.go b/vault/expiration.go index 4c4e91769..a378e53a5 100644 --- a/vault/expiration.go +++ b/vault/expiration.go @@ -249,7 +249,7 @@ func (m *ExpirationManager) Renew(leaseID string, increment time.Duration) (*log } // Fast-path if there is no lease - if resp == nil || resp.Secret == nil || resp.Secret.Lease == 0 { + if resp == nil || resp.Secret == nil || !resp.Secret.LeaseEnabled() { return resp, nil } @@ -262,14 +262,9 @@ func (m *ExpirationManager) Renew(leaseID string, increment time.Duration) (*log resp.Secret.LeaseID = leaseID // Update the lease entry - var expireTime time.Time - leaseTotal := resp.Secret.Lease + resp.Secret.LeaseGracePeriod - if resp.Secret.Lease > 0 { - expireTime = time.Now().UTC().Add(leaseTotal) - } le.Data = resp.Data le.Secret = resp.Secret - le.ExpireTime = expireTime + le.ExpireTime = resp.Secret.ExpirationTime() if err := m.persistEntry(le); err != nil { return nil, err } @@ -277,7 +272,12 @@ func (m *ExpirationManager) Renew(leaseID string, increment time.Duration) (*log // Update the expiration time m.pendingLock.Lock() if timer, ok := m.pending[leaseID]; ok { - timer.Reset(leaseTotal) + if le.ExpireTime.IsZero() { + timer.Stop() + delete(m.pending, leaseID) + } else { + timer.Reset(resp.Secret.LeaseTotal()) + } } m.pendingLock.Unlock() @@ -314,24 +314,22 @@ func (m *ExpirationManager) RenewToken(source string, token string) (*logical.Au } // Update the lease entry - var expireTime time.Time - leaseTotal := le.Auth.Lease + le.Auth.LeaseGracePeriod - if le.Auth.Lease > 0 { - expireTime = time.Now().UTC().Add(leaseTotal) - } - le.ExpireTime = expireTime + le.ExpireTime = le.Auth.ExpirationTime() if err := m.persistEntry(le); err != nil { return nil, err } // Update the expiration time - if !expireTime.IsZero() { - m.pendingLock.Lock() - if timer, ok := m.pending[leaseID]; ok { - timer.Reset(leaseTotal) + m.pendingLock.Lock() + if timer, ok := m.pending[leaseID]; ok { + if le.ExpireTime.IsZero() { + timer.Stop() + delete(m.pending, leaseID) + } else { + timer.Reset(le.Auth.LeaseTotal()) } - m.pendingLock.Unlock() } + m.pendingLock.Unlock() return le.Auth, nil } @@ -351,19 +349,13 @@ func (m *ExpirationManager) Register(req *logical.Request, resp *logical.Respons } // Create a lease entry - now := time.Now().UTC() - leaseTotal := resp.Secret.Lease + resp.Secret.LeaseGracePeriod - var expireTime time.Time - if resp.Secret.Lease > 0 { - expireTime = now.Add(leaseTotal) - } le := leaseEntry{ LeaseID: path.Join(req.Path, generateUUID()), Path: req.Path, Data: resp.Data, Secret: resp.Secret, - IssueTime: now, - ExpireTime: expireTime, + IssueTime: time.Now().UTC(), + ExpireTime: resp.Secret.ExpirationTime(), } // Encode the entry @@ -372,9 +364,9 @@ func (m *ExpirationManager) Register(req *logical.Request, resp *logical.Respons } // Setup revocation timer if there is a lease - if !expireTime.IsZero() { + if !le.ExpireTime.IsZero() { m.pendingLock.Lock() - m.pending[le.LeaseID] = time.AfterFunc(leaseTotal, func() { + m.pending[le.LeaseID] = time.AfterFunc(resp.Secret.LeaseTotal(), func() { m.expireID(le.LeaseID) }) m.pendingLock.Unlock() @@ -390,18 +382,12 @@ func (m *ExpirationManager) Register(req *logical.Request, resp *logical.Respons func (m *ExpirationManager) RegisterAuth(source string, auth *logical.Auth) error { defer metrics.MeasureSince([]string{"expire", "register-auth"}, time.Now()) // Create a lease entry - now := time.Now().UTC() - leaseTotal := auth.Lease + auth.LeaseGracePeriod - var expireTime time.Time - if auth.Lease > 0 { - expireTime = now.Add(leaseTotal) - } le := leaseEntry{ LeaseID: path.Join(source, m.tokenStore.SaltID(auth.ClientToken)), Auth: auth, Path: source, - IssueTime: now, - ExpireTime: expireTime, + IssueTime: time.Now().UTC(), + ExpireTime: auth.ExpirationTime(), } // Encode the entry @@ -410,9 +396,9 @@ func (m *ExpirationManager) RegisterAuth(source string, auth *logical.Auth) erro } // Setup revocation timer - if !expireTime.IsZero() { + if !le.ExpireTime.IsZero() { m.pendingLock.Lock() - m.pending[le.LeaseID] = time.AfterFunc(leaseTotal, func() { + m.pending[le.LeaseID] = time.AfterFunc(auth.LeaseTotal(), func() { m.expireID(le.LeaseID) }) m.pendingLock.Unlock()