From 2c9ebecda770f4045c7321e9c9f411214bdeaf40 Mon Sep 17 00:00:00 2001
From: Mitchell Hashimoto <mitchell.hashimoto@gmail.com>
Date: Tue, 31 Mar 2015 21:01:12 -0700
Subject: [PATCH] vault: register zero lease entries with the expiration
 manager

/cc @armon - would appreciate a review on this one
---
 vault/core.go       |  5 +----
 vault/expiration.go | 31 +++++++++++++++++++++++--------
 2 files changed, 24 insertions(+), 12 deletions(-)

diff --git a/vault/core.go b/vault/core.go
index 2708ea4b8..bdc9a0503 100644
--- a/vault/core.go
+++ b/vault/core.go
@@ -250,10 +250,7 @@ func (c *Core) handleRequest(req *logical.Request) (*logical.Response, error) {
 	resp, err := c.router.Route(req)
 
 	// If there is a secret, we must register it with the expiration manager.
-	//
-	// TODO(mitchellh): what about secrets with a lease of 0, do we still
-	// record them so they're revoked during unmount?
-	if resp != nil && resp.Secret != nil && resp.Secret.Lease > 0 {
+	if resp != nil && resp.Secret != nil {
 		vaultID, err := c.expiration.Register(req, resp)
 		if err != nil {
 			c.logger.Printf(
diff --git a/vault/expiration.go b/vault/expiration.go
index ebd01d6a0..4c05bf60e 100644
--- a/vault/expiration.go
+++ b/vault/expiration.go
@@ -110,6 +110,11 @@ func (m *ExpirationManager) Restore() error {
 			continue
 		}
 
+		// If there is no expiry time, don't do anything
+		if le.ExpireTime.IsZero() {
+			continue
+		}
+
 		// Determine the remaining time to expiration
 		expires := le.ExpireTime.Sub(time.Now().UTC())
 		if expires <= 0 {
@@ -239,10 +244,14 @@ func (m *ExpirationManager) Renew(vaultID string, increment time.Duration) (*log
 	resp.Secret.VaultID = vaultID
 
 	// 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 = time.Now().UTC().Add(leaseTotal)
+	le.ExpireTime = expireTime
 	if err := m.persistEntry(le); err != nil {
 		return nil, err
 	}
@@ -275,13 +284,17 @@ 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{
 		VaultID:    path.Join(req.Path, generateUUID()),
 		Path:       req.Path,
 		Data:       resp.Data,
 		Secret:     resp.Secret,
 		IssueTime:  now,
-		ExpireTime: now.Add(leaseTotal),
+		ExpireTime: expireTime,
 	}
 
 	// Encode the entry
@@ -289,12 +302,14 @@ func (m *ExpirationManager) Register(req *logical.Request, resp *logical.Respons
 		return "", err
 	}
 
-	// Setup revocation timer
-	m.pendingLock.Lock()
-	m.pending[le.VaultID] = time.AfterFunc(leaseTotal, func() {
-		m.expireID(le.VaultID)
-	})
-	m.pendingLock.Unlock()
+	// Setup revocation timer if there is a lease
+	if !expireTime.IsZero() {
+		m.pendingLock.Lock()
+		m.pending[le.VaultID] = time.AfterFunc(leaseTotal, func() {
+			m.expireID(le.VaultID)
+		})
+		m.pendingLock.Unlock()
+	}
 
 	// Done
 	return le.VaultID, nil