vault: Tokens can have a use count specified

vault/token_store.go

@@ -234,6 +234,7 @@ type TokenEntry struct {
 	Path        string            // Used for audit trails, this is something like "auth/user/login"
 	Meta        map[string]string // Used for auditing. This could include things like "source", "user", "ip"
 	DisplayName string            // Used for operators to be able to associate with the source
+	NumUses     int               // Used to restrict the number of uses (zero is unlimited). This is to support one-time-tokens (generalized).
 }
 
 // SetExpirationManager is used to provide the token store with
@@ -437,6 +438,13 @@ func (ts *TokenStore) handleCreate(
 		return logical.ErrorResponse("parent token lookup failed"), logical.ErrInvalidRequest
 	}
 
+	// A token with a restricted number of uses cannot create a new token
+	// otherwise it could escape the restriction count.
+	if parent.NumUses > 0 {
+		return logical.ErrorResponse("restricted use token cannot generate child tokens"),
+			logical.ErrInvalidRequest
+	}
+
 	// Check if the parent policy is root
 	isRoot := strListContains(parent.Policies, "root")
 
@@ -448,18 +456,26 @@ func (ts *TokenStore) handleCreate(
 		NoParent    bool              `mapstructure:"no_parent"`
 		Lease       string
 		DisplayName string `mapstructure:"display_name"`
+		NumUses     int    `mapstructure:"num_uses"`
 	}
 	if err := mapstructure.WeakDecode(req.Data, &data); err != nil {
 		return logical.ErrorResponse(fmt.Sprintf(
 			"Error decoding request: %s", err)), logical.ErrInvalidRequest
 	}
 
+	// Verify the number of uses is positive
+	if data.NumUses < 0 {
+		return logical.ErrorResponse("number of uses cannot be negative"),
+			logical.ErrInvalidRequest
+	}
+
 	// Setup the token entry
 	te := TokenEntry{
 		Parent:      req.ClientToken,
 		Path:        "auth/token/create",
 		Meta:        data.Metadata,
 		DisplayName: "token",
+		NumUses:     data.NumUses,
 	}
 
 	// Attach the given display name if any
@@ -615,6 +631,7 @@ func (ts *TokenStore) handleLookup(
 			"path":         out.Path,
 			"meta":         out.Meta,
 			"display_name": out.DisplayName,
+			"num_uses":     out.NumUses,
 		},
 	}
 	return resp, nil

vault/token_store_test.go

@@ -297,6 +297,68 @@ func TestTokenStore_HandleRequest_CreateToken_DisplayName(t *testing.T) {
 	}
 }
 
+func TestTokenStore_HandleRequest_CreateToken_NumUses(t *testing.T) {
+	_, ts, root := mockTokenStore(t)
+
+	req := logical.TestRequest(t, logical.WriteOperation, "create")
+	req.ClientToken = root
+	req.Data["num_uses"] = "1"
+
+	resp, err := ts.HandleRequest(req)
+	if err != nil {
+		t.Fatalf("err: %v %v", err, resp)
+	}
+
+	expected := &TokenEntry{
+		ID:          resp.Auth.ClientToken,
+		Parent:      root,
+		Policies:    []string{"root"},
+		Path:        "auth/token/create",
+		DisplayName: "token",
+		NumUses:     1,
+	}
+	out, err := ts.Lookup(resp.Auth.ClientToken)
+	if err != nil {
+		t.Fatalf("err: %v", err)
+	}
+	if !reflect.DeepEqual(out, expected) {
+		t.Fatalf("bad: %#v", out)
+	}
+}
+
+func TestTokenStore_HandleRequest_CreateToken_NumUses_Invalid(t *testing.T) {
+	_, ts, root := mockTokenStore(t)
+
+	req := logical.TestRequest(t, logical.WriteOperation, "create")
+	req.ClientToken = root
+	req.Data["num_uses"] = "-1"
+
+	resp, err := ts.HandleRequest(req)
+	if err != logical.ErrInvalidRequest {
+		t.Fatalf("err: %v %v", err, resp)
+	}
+}
+
+func TestTokenStore_HandleRequest_CreateToken_NumUses_Restricted(t *testing.T) {
+	_, ts, root := mockTokenStore(t)
+
+	req := logical.TestRequest(t, logical.WriteOperation, "create")
+	req.ClientToken = root
+	req.Data["num_uses"] = "1"
+
+	resp, err := ts.HandleRequest(req)
+	if err != nil {
+		t.Fatalf("err: %v %v", err, resp)
+	}
+
+	// We should NOT be able to use the restricted token to create a new token
+	req.ClientToken = resp.Auth.ClientToken
+	_, err = ts.HandleRequest(req)
+	if err != logical.ErrInvalidRequest {
+		t.Fatalf("err: %v %v", err, resp)
+	}
+}
+
 func TestTokenStore_HandleRequest_CreateToken_NoPolicy(t *testing.T) {
 	_, ts, root := mockTokenStore(t)
 
@@ -594,6 +656,7 @@ func TestTokenStore_HandleRequest_Lookup(t *testing.T) {
 		"path":         "auth/token/root",
 		"meta":         map[string]string(nil),
 		"display_name": "root",
+		"num_uses":     0,
 	}
 	if !reflect.DeepEqual(resp.Data, exp) {
 		t.Fatalf("bad: %#v exp: %#v", resp.Data, exp)
@@ -658,6 +721,7 @@ func TestTokenStore_HandleRequest_LookupSelf(t *testing.T) {
 		"path":         "auth/token/root",
 		"meta":         map[string]string(nil),
 		"display_name": "root",
+		"num_uses":     0,
 	}
 	if !reflect.DeepEqual(resp.Data, exp) {
 		t.Fatalf("bad: %#v exp: %#v", resp.Data, exp)