From 810e9147306a7d2ff411b51f234ba827f1e8f6a8 Mon Sep 17 00:00:00 2001
From: Jeff Mitchell <jeffrey.mitchell@gmail.com>
Date: Thu, 19 May 2016 11:47:18 -0400
Subject: [PATCH] Add unwrap test function and some robustness around paths for
 the wrap lookup function

---
 api/client.go          | 16 ++++++++-
 command/unwrap_test.go | 74 ++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 89 insertions(+), 1 deletion(-)
 create mode 100644 command/unwrap_test.go

diff --git a/api/client.go b/api/client.go
index 376ad1df0e..70bc5e00e5 100644
--- a/api/client.go
+++ b/api/client.go
@@ -29,6 +29,11 @@ var (
 	errRedirect = errors.New("redirect")
 )
 
+// WrappingLookupFunc is a function that, given an HTTP verb and a path,
+// returns an optional string duration to be used for response wrapping (e.g.
+// "15s", or simply "15"). The path will not begin with "/v1/" or "v1/" or "/",
+// however, end-of-path forward slashes are not trimmed, so must match your
+// called path precisely.
 type WrappingLookupFunc func(operation, path string) string
 
 // Config is used to configure the creation of the client.
@@ -242,7 +247,16 @@ func (c *Client) NewRequest(method, path string) *Request {
 	}
 
 	if c.wrappingLookupFunc != nil {
-		req.WrapTTL = c.wrappingLookupFunc(method, path)
+		var lookupPath string
+		switch {
+		case strings.HasPrefix(path, "/v1/"):
+			lookupPath = strings.TrimPrefix(path, "/v1/")
+		case strings.HasPrefix(path, "v1/"):
+			lookupPath = strings.TrimPrefix(path, "v1/")
+		default:
+			lookupPath = path
+		}
+		req.WrapTTL = c.wrappingLookupFunc(method, lookupPath)
 	}
 
 	return req
diff --git a/command/unwrap_test.go b/command/unwrap_test.go
new file mode 100644
index 0000000000..ec7759edc0
--- /dev/null
+++ b/command/unwrap_test.go
@@ -0,0 +1,74 @@
+package command
+
+import (
+	"testing"
+
+	"github.com/hashicorp/vault/http"
+	"github.com/hashicorp/vault/meta"
+	"github.com/hashicorp/vault/vault"
+	"github.com/mitchellh/cli"
+)
+
+func TestUnwrap(t *testing.T) {
+	core, _, token := vault.TestCoreUnsealed(t)
+	ln, addr := http.TestServer(t, core)
+	defer ln.Close()
+
+	ui := new(cli.MockUi)
+	c := &UnwrapCommand{
+		Meta: meta.Meta{
+			ClientToken: token,
+			Ui:          ui,
+		},
+	}
+
+	args := []string{
+		"-address", addr,
+		"-field", "zip",
+	}
+
+	// Run once so the client is setup, ignore errors
+	c.Run(args)
+
+	// Get the client so we can write data
+	client, err := c.Client()
+	if err != nil {
+		t.Fatalf("err: %s", err)
+	}
+
+	wrapLookupFunc := func(method, path string) string {
+		if method == "GET" && path == "secret/foo" {
+			return "60s"
+		}
+		return ""
+	}
+	client.SetWrappingLookupFunc(wrapLookupFunc)
+
+	data := map[string]interface{}{"zip": "zap"}
+	if _, err := client.Logical().Write("secret/foo", data); err != nil {
+		t.Fatalf("err: %s", err)
+	}
+
+	outer, err := client.Logical().Read("secret/foo")
+	if err != nil {
+		t.Fatalf("err: %s", err)
+	}
+	if outer == nil {
+		t.Fatal("outer response was nil")
+	}
+	if outer.WrapInfo == nil {
+		t.Fatal("outer wrapinfo was nil, response was %#v", *outer)
+	}
+
+	args = append(args, outer.WrapInfo.Token)
+
+	// Run the read
+	if code := c.Run(args); code != 0 {
+		t.Fatalf("bad: %d\n\n%s", code, ui.ErrorWriter.String())
+	}
+
+	output := ui.OutputWriter.String()
+	if output != "zap\n" {
+		t.Fatalf("unexpectd output:\n%s", output)
+	}
+}