From 27d33ad9f792994eba30a14e672683c2e3634134 Mon Sep 17 00:00:00 2001
From: Mitchell Hashimoto <mitchell.hashimoto@gmail.com>
Date: Sat, 21 Mar 2015 16:20:30 +0100
Subject: [PATCH] logical/framework: auto-extend leases if requested

---
 logical/framework/backend.go      | 14 +-----
 logical/framework/backend_test.go | 24 ++++++++++
 logical/framework/secret.go       | 76 ++++++++++++++++++++++++++++---
 3 files changed, 96 insertions(+), 18 deletions(-)

diff --git a/logical/framework/backend.go b/logical/framework/backend.go
index d019de4c5..98c30abe3 100644
--- a/logical/framework/backend.go
+++ b/logical/framework/backend.go
@@ -181,25 +181,15 @@ func (b *Backend) handleRevokeRenew(
 		return nil, fmt.Errorf("secret is unsupported by this backend")
 	}
 
-	var fn OperationFunc
 	switch req.Operation {
 	case logical.RenewOperation:
-		fn = secret.Renew
+		return secret.HandleRenew(req)
 	case logical.RevokeOperation:
-		fn = secret.Revoke
+		return secret.HandleRevoke(req)
 	default:
 		return nil, fmt.Errorf(
 			"invalid operation for revoke/renew: %s", req.Operation)
 	}
-
-	if fn == nil {
-		return nil, logical.ErrUnsupportedOperation
-	}
-
-	return fn(req, &FieldData{
-		Raw:    req.Data,
-		Schema: secret.Fields,
-	})
 }
 
 func (b *Backend) handleRollback(
diff --git a/logical/framework/backend_test.go b/logical/framework/backend_test.go
index 7317cc737..d699858c6 100644
--- a/logical/framework/backend_test.go
+++ b/logical/framework/backend_test.go
@@ -162,6 +162,30 @@ func TestBackendHandleRequest_renew(t *testing.T) {
 	}
 }
 
+func TestBackendHandleRequest_renewExtend(t *testing.T) {
+	secret := &Secret{
+		Type:            "foo",
+		RenewExtend:     true,
+		DefaultDuration: 5 * time.Second,
+	}
+	b := &Backend{
+		Secrets: []*Secret{secret},
+	}
+
+	req := logical.RenewRequest("/foo", secret.Response(nil, nil).Secret, nil)
+	req.Secret.LeaseIncrement = 1 * time.Hour
+	resp, err := b.HandleRequest(req)
+	if err != nil {
+		t.Fatalf("err: %s", err)
+	}
+	if resp == nil || resp.Secret == nil {
+		t.Fatal("should have secret")
+	}
+	if resp.Secret.Lease != 1*time.Hour {
+		t.Fatalf("bad: %#v", resp.Secret)
+	}
+}
+
 func TestBackendHandleRequest_revoke(t *testing.T) {
 	var called uint32
 	callback := func(*logical.Request, *FieldData) (*logical.Response, error) {
diff --git a/logical/framework/secret.go b/logical/framework/secret.go
index e68ce295b..e7f3e848f 100644
--- a/logical/framework/secret.go
+++ b/logical/framework/secret.go
@@ -25,15 +25,26 @@ type Secret struct {
 	DefaultDuration    time.Duration
 	DefaultGracePeriod time.Duration
 
-	// Below are the operations that can be called on the secret.
+	// Renew is the callback called to renew this secret. If Renew is
+	// not specified and RenewExtend is false, then Renewable is set to
+	// false in the secret.
 	//
-	// Renew, if not set, will mark the secret as not renewable.
-	//
-	// Revoke is required.
-	Renew  OperationFunc
+	// RenewExtend, if true, will automatically extend the lease of this
+	// secret type. You can specify RenewExtendMax to specify the max
+	// duration it can be extended, otherwise it will be extended potentially
+	// indefinitely.
+	Renew          OperationFunc
+	RenewExtend    bool
+	RenewExtendMax time.Duration
+
+	// Revoke is the callback called to revoke this secret. This is required.
 	Revoke OperationFunc
 }
 
+func (s *Secret) Renewable() bool {
+	return s.Renew != nil || s.RenewExtend
+}
+
 func (s *Secret) Response(
 	data, internal map[string]interface{}) *logical.Response {
 	internalData := make(map[string]interface{})
@@ -44,7 +55,7 @@ func (s *Secret) Response(
 
 	return &logical.Response{
 		Secret: &logical.Secret{
-			Renewable:        s.Renew != nil,
+			Renewable:        s.Renewable(),
 			Lease:            s.DefaultDuration,
 			LeaseGracePeriod: s.DefaultGracePeriod,
 			InternalData:     internalData,
@@ -53,3 +64,56 @@ func (s *Secret) Response(
 		Data: data,
 	}
 }
+
+// HandleRenew is the request handler for renewing this secret.
+func (s *Secret) HandleRenew(req *logical.Request) (*logical.Response, error) {
+	if !s.Renewable() {
+		return nil, logical.ErrUnsupportedOperation
+	}
+
+	data := &FieldData{
+		Raw:    req.Data,
+		Schema: s.Fields,
+	}
+
+	// If we have a callback, we just call that and that does all the logic.
+	if s.Renew != nil {
+		return s.Renew(req, data)
+	}
+
+	// If we're using RenewExtend, then just automaticaly extend.
+	if s.RenewExtend {
+		return s.HandleRenewExtend(req, data)
+	}
+
+	return nil, logical.ErrUnsupportedOperation
+}
+
+// HandleRenewExtend is the OperationFunc that just extends the lease
+// of the secret.
+func (s *Secret) HandleRenewExtend(
+	req *logical.Request, data *FieldData) (*logical.Response, error) {
+	// First copy the original secret/data
+	var resp logical.Response
+	resp.Secret = req.Secret
+	resp.Data = req.Data
+
+	// Now extend the lease by the amount specified.
+	resp.Secret.Lease = req.Secret.LeaseIncrement
+
+	return &resp, nil
+}
+
+// HandleRevoke is the request handler for renewing this secret.
+func (s *Secret) HandleRevoke(req *logical.Request) (*logical.Response, error) {
+	data := &FieldData{
+		Raw:    req.Data,
+		Schema: s.Fields,
+	}
+
+	if s.Revoke != nil {
+		return s.Revoke(req, data)
+	}
+
+	return nil, logical.ErrUnsupportedOperation
+}