Add unwrap test function and some robustness around paths for the wrap lookup function

This commit is contained in:
Jeff Mitchell 2016-05-19 11:47:18 -04:00
parent 0da8762bd5
commit 05b2d4534c
2 changed files with 89 additions and 1 deletions

View File

@ -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

74
command/unwrap_test.go Normal file
View File

@ -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)
}
}