logical/framework: AuthRenew callback, add LeaseExtend

/cc @armon - Going with this "standard library" of callbacks approach
to make extending leases in a customizable way easy. See the docs/tests
above.

logical/framework/backend.go

@@ -51,6 +51,11 @@ type Backend struct {
 	Rollback       RollbackFunc
 	RollbackMinAge time.Duration
 
+	// AuthRenew is the callback to call when a RenewRequest for an
+	// authentication comes in. By default, renewal won't be allowed.
+	// See the built-in AuthRenew helpers in lease.go for common callbacks.
+	AuthRenew OperationFunc
+
 	logger  *log.Logger
 	once    sync.Once
 	pathsRe []*regexp.Regexp
@@ -272,12 +277,11 @@ func (b *Backend) handleRevokeRenew(
 }
 
 func (b *Backend) handleAuthRenew(req *logical.Request) (*logical.Response, error) {
-	// TODO: make this customizable
+	if b.AuthRenew == nil {
+		return logical.ErrorResponse("this auth type doesn't support renew"), nil
+	}
 
-	// Set the lease to the requested increment
-	req.Auth.Lease = req.Auth.IncrementedLease(req.Auth.LeaseIncrement)
-	resp := &logical.Response{Auth: req.Auth}
-	return resp, nil
+	return b.AuthRenew(req, nil)
 }
 
 func (b *Backend) handleRollback(

logical/framework/backend_test.go

@@ -154,6 +154,39 @@ func TestBackendHandleRequest_helpRoot(t *testing.T) {
 	}
 }
 
+func TestBackendHandleRequest_renewAuth(t *testing.T) {
+	b := &Backend{}
+
+	resp, err := b.HandleRequest(logical.RenewAuthRequest(
+		"/foo", &logical.Auth{}, nil))
+	if err != nil {
+		t.Fatalf("err: %s", err)
+	}
+	if !resp.IsError() {
+		t.Fatalf("bad: %#v", resp)
+	}
+}
+
+func TestBackendHandleRequest_renewAuthCallback(t *testing.T) {
+	var called uint32
+	callback := func(*logical.Request, *FieldData) (*logical.Response, error) {
+		atomic.AddUint32(&called, 1)
+		return nil, nil
+	}
+
+	b := &Backend{
+		AuthRenew: callback,
+	}
+
+	_, err := b.HandleRequest(logical.RenewAuthRequest(
+		"/foo", &logical.Auth{}, nil))
+	if err != nil {
+		t.Fatalf("err: %s", err)
+	}
+	if v := atomic.LoadUint32(&called); v != 1 {
+		t.Fatalf("bad: %#v", v)
+	}
+}
 func TestBackendHandleRequest_renew(t *testing.T) {
 	var called uint32
 	callback := func(*logical.Request, *FieldData) (*logical.Response, error) {

logical/framework/lease.go

@@ -0,0 +1,48 @@
+package framework
+
+import (
+	"fmt"
+	"time"
+
+	"github.com/hashicorp/vault/logical"
+)
+
+// LeaseExtend returns an OperationFunc that can be used to simply extend
+// the lease of the auth/secret for the duration that was requested. Max
+// is the max time past the _current_ time that a lease can be extended. i.e.
+// setting it to 2 hours forces a renewal within the next 2 hours again.
+func LeaseExtend(max time.Duration) OperationFunc {
+	return func(req *logical.Request, data *FieldData) (*logical.Response, error) {
+		lease := detectLease(req)
+		if lease == nil {
+			return nil, fmt.Errorf("no lease options for request")
+		}
+
+		// Determine the requested lease
+		newLease := lease.IncrementedLease(lease.LeaseIncrement)
+
+		// Determine if the requested lease is too long
+		now := time.Now().UTC()
+		maxExpiration := now.Add(max)
+		newExpiration := now.Add(newLease)
+		if newExpiration.Sub(maxExpiration) > 0 {
+			// The new expiration is past the max expiration. In this
+			// case, admit the longest lease we can.
+			newLease = maxExpiration.Sub(lease.ExpirationTime())
+		}
+
+		// Set the lease
+		lease.Lease = newLease
+		return &logical.Response{Auth: req.Auth, Secret: req.Secret}, nil
+	}
+}
+
+func detectLease(req *logical.Request) *logical.LeaseOptions {
+	if req.Auth != nil {
+		return &req.Auth.LeaseOptions
+	} else if req.Secret != nil {
+		return &req.Secret.LeaseOptions
+	}
+
+	return nil
+}

logical/framework/lease_test.go

@@ -0,0 +1,54 @@
+package framework
+
+import (
+	"testing"
+	"time"
+
+	"github.com/hashicorp/vault/logical"
+)
+
+func TestLeaseExtend(t *testing.T) {
+	now := time.Now().UTC().Round(time.Hour)
+
+	cases := map[string]struct {
+		Max     time.Duration
+		Request time.Duration
+		Result  time.Duration
+	}{
+		"valid request, good bounds": {
+			Max:     30 * time.Hour,
+			Request: 1 * time.Hour,
+			Result:  1 * time.Hour,
+		},
+
+		"request is too long": {
+			Max:     3 * time.Hour,
+			Request: 7 * time.Hour,
+			Result:  3 * time.Hour,
+		},
+	}
+
+	for name, tc := range cases {
+		req := &logical.Request{
+			Auth: &logical.Auth{
+				LeaseOptions: logical.LeaseOptions{
+					Lease:          1 * time.Second,
+					LeaseIssue:     now,
+					LeaseIncrement: tc.Request,
+				},
+			},
+		}
+
+		callback := LeaseExtend(tc.Max)
+		resp, err := callback(req, nil)
+		if err != nil {
+			t.Fatalf("bad: %s\nerr: %s", name, err)
+		}
+
+		// Round it to the nearest hour
+		lease := now.Add(resp.Auth.Lease).Round(time.Hour).Sub(now)
+		if lease != tc.Result {
+			t.Fatalf("bad: %s\nlease: %s", name, lease)
+		}
+	}
+}

logical/request.go

@@ -73,6 +73,17 @@ func RenewRequest(
 	}
 }
 
+// RenewAuthRequest creates the structure of the renew request for an auth.
+func RenewAuthRequest(
+	path string, auth *Auth, data map[string]interface{}) *Request {
+	return &Request{
+		Operation: RenewOperation,
+		Path:      path,
+		Data:      data,
+		Auth:      auth,
+	}
+}
+
 // RevokeRequest creates the structure of the revoke request.
 func RevokeRequest(
 	path string, secret *Secret, data map[string]interface{}) *Request {

vault/expiration.go

@@ -515,9 +515,7 @@ func (m *ExpirationManager) renewAuthEntry(le *leaseEntry, increment time.Durati
 	auth.LeaseIncrement = increment
 	auth.ClientToken = ""
 
-	req := logical.RenewRequest(le.Path, nil, nil)
-	req.Auth = &auth
-
+	req := logical.RenewAuthRequest(le.Path, &auth, nil)
 	resp, err := m.router.Route(req)
 	if err != nil {
 		return nil, fmt.Errorf("failed to renew entry: %v", err)