vault: integration expiration manager with core

vault/core.go

@@ -34,6 +34,10 @@ var (
 	// ErrNotInit is returned if a non-initialized barrier
 	// is attempted to be unsealed.
 	ErrNotInit = errors.New("Vault is not initialized")
+
+	// ErrInternalError is returned when we don't want to leak
+	// any information about an internal error
+	ErrInternalError = errors.New("internal error")
 )
 
 // SealConfig is used to describe the seal configuration
@@ -176,7 +180,20 @@ func (c *Core) HandleRequest(req *logical.Request) (*logical.Response, error) {
 	// TODO: Enforce ACLs
 
 	// Route the request
-	return c.router.Route(req)
+	resp, err := c.router.Route(req)
+
+	// Check if there is a lease, we must register this
+	if resp != nil && resp.IsSecret && resp.Lease != nil {
+		vaultID, err := c.expiration.Register(req, resp)
+		if err != nil {
+			c.logger.Printf("[ERR] core: failed to register lease (request: %#v, response: %#v): %v", req, resp, err)
+			return nil, ErrInternalError
+		}
+		resp.Lease.VaultID = vaultID
+	}
+
+	// Return the response and error
+	return resp, err
 }
 
 // Initialized checks if the Vault is already initialized

vault/core_test.go

@@ -3,6 +3,7 @@ package vault
 import (
 	"reflect"
 	"testing"
+	"time"
 
 	"github.com/hashicorp/vault/logical"
 	"github.com/hashicorp/vault/physical"
@@ -312,3 +313,44 @@ func TestCore_SealUnseal(t *testing.T) {
 		t.Fatalf("err: %v", err)
 	}
 }
+
+// Ensure we get a VaultID
+func TestCore_HandleRequest_Lease(t *testing.T) {
+	c, _ := TestCoreUnsealed(t)
+
+	req := &logical.Request{
+		Operation: logical.WriteOperation,
+		Path:      "secret/test",
+		Data: map[string]interface{}{
+			"foo":   "bar",
+			"lease": "1h",
+		},
+	}
+	resp, err := c.HandleRequest(req)
+	if err != nil {
+		t.Fatalf("err: %v", err)
+	}
+	if resp != nil {
+		t.Fatalf("bad: %#v", resp)
+	}
+
+	// Read the key
+	req.Operation = logical.ReadOperation
+	req.Data = nil
+	resp, err = c.HandleRequest(req)
+	if err != nil {
+		t.Fatalf("err: %v", err)
+	}
+	if resp == nil || resp.Lease == nil || resp.Data == nil {
+		t.Fatalf("bad: %#v", resp)
+	}
+	if resp.Lease.Duration != time.Hour {
+		t.Fatalf("bad: %#v", resp.Lease)
+	}
+	if resp.Lease.VaultID == "" {
+		t.Fatalf("bad: %#v", resp.Lease)
+	}
+	if resp.Data["foo"] != "bar" {
+		t.Fatalf("bad: %#v", resp.Data)
+	}
+}