Alex/add max num uses to tokenhelper v2

sdk/helper/tokenhelper/tokenhelper.go

@@ -38,6 +38,9 @@ type TokenParams struct {
 
 	// The TTL to user for the token
 	TokenTTL time.Duration `json:"token_ttl" mapstructure:"token_ttl"`
+
+	// The maximum number of times a token issued from this role may be used.
+	TokenNumUses int `json:"token_num_uses,omitempty" mapstructure:"token_num_uses,omitempty"`
 }
 
 // AddTokenFields adds fields to an existing role. It panics if it would
@@ -101,6 +104,11 @@ func TokenFields() map[string]*framework.FieldSchema {
 			Type:        framework.TypeDurationSecond,
 			Description: "The initial ttl of the token to generate",
 		},
+
+		"token_num_uses": &framework.FieldSchema{
+			Type:        framework.TypeInt,
+			Description: "The maximum number of times a token may be used, a value of zero means unlimited",
+		},
 	}
 }
 
@@ -167,6 +175,13 @@ func (t *TokenParams) ParseTokenFields(req *logical.Request, d *framework.FieldD
 		return errors.New("'token_ttl' cannot be greater than 'token_max_ttl'")
 	}
 
+	if tokenNumUses, ok := d.GetOk("token_num_uses"); ok {
+		t.TokenNumUses = tokenNumUses.(int)
+	}
+	if t.TokenNumUses < 0 {
+		return errors.New("'token_num_uses' cannot be negative")
+	}
+
 	return nil
 }
 
@@ -179,6 +194,7 @@ func (t *TokenParams) PopulateTokenData(m map[string]interface{}) {
 	m["token_policies"] = t.TokenPolicies
 	m["token_type"] = t.TokenType.String()
 	m["token_ttl"] = t.TokenTTL.Seconds()
+	m["token_num_uses"] = t.TokenNumUses
 }
 
 func (t *TokenParams) PopulateTokenAuth(auth *logical.Auth) {
@@ -190,6 +206,7 @@ func (t *TokenParams) PopulateTokenAuth(auth *logical.Auth) {
 	auth.Policies = t.TokenPolicies
 	auth.TokenType = t.TokenType
 	auth.TTL = t.TokenTTL
+	auth.NumUses = t.TokenNumUses
 }
 
 const (

vault/token_store.go

@@ -409,7 +409,7 @@ func (ts *TokenStore) paths() []*framework.Path {
 		ExistenceCheck: ts.tokenStoreRoleExistenceCheck,
 	}
 
-	tokenhelper.AddTokenFieldsWithAllowList(rolesPath.Fields, []string{"token_bound_cidrs", "token_explicit_max_ttl", "token_period", "token_type", "token_no_default_policy"})
+	tokenhelper.AddTokenFieldsWithAllowList(rolesPath.Fields, []string{"token_bound_cidrs", "token_explicit_max_ttl", "token_period", "token_type", "token_no_default_policy", "token_num_uses"})
 	p = append(p, rolesPath)
 
 	return p
@@ -2234,6 +2234,16 @@ func (ts *TokenStore) handleCreateCommon(ctx context.Context, req *logical.Reque
 			renewable = false
 		}
 
+		// Update num_uses which is equal to req.Data["num_uses"] at this point
+		// 0 means unlimited so 1 is actually less than 0
+		switch {
+		case role.TokenNumUses > 0 && te.NumUses == 0:
+			te.NumUses = role.TokenNumUses
+		case role.TokenNumUses < te.NumUses:
+			te.NumUses = role.TokenNumUses
+		default:
+		}
+
 		if role.PathSuffix != "" {
 			te.Path = fmt.Sprintf("%s/%s", te.Path, role.PathSuffix)
 		}
@@ -2978,6 +2988,9 @@ func (ts *TokenStore) tokenStoreRoleRead(ctx context.Context, req *logical.Reque
 	if len(role.BoundCIDRs) > 0 {
 		resp.Data["bound_cidrs"] = role.BoundCIDRs
 	}
+	if role.TokenNumUses > 0 {
+		resp.Data["token_num_uses"] = role.TokenNumUses
+	}
 
 	return resp, nil
 }
@@ -3149,6 +3162,12 @@ func (ts *TokenStore) tokenStoreRoleCreateUpdate(ctx context.Context, req *logic
 		}
 	}
 
+	// no legacy version without the token_ prefix to check for
+	tokenNumUses, ok := data.GetOk("token_num_uses")
+	if ok {
+		entry.TokenNumUses = tokenNumUses.(int)
+	}
+
 	// Run validity checks on token type
 	if entry.TokenType == logical.TokenTypeBatch {
 		if !entry.Orphan {

vault/token_store_test.go

@@ -2682,6 +2682,7 @@ func TestTokenStore_RoleCRUD(t *testing.T) {
 		"path_suffix":      "happenin",
 		"bound_cidrs":      []string{"0.0.0.0/0"},
 		"explicit_max_ttl": "2h",
+		"token_num_uses":   123,
 	}
 
 	resp, err = core.HandleRequest(namespace.RootContext(nil), req)
@@ -2717,6 +2718,7 @@ func TestTokenStore_RoleCRUD(t *testing.T) {
 		"token_explicit_max_ttl": int64(7200),
 		"renewable":              true,
 		"token_type":             "default-service",
+		"token_num_uses":         123,
 	}
 
 	if resp.Data["bound_cidrs"].([]*sockaddr.SockAddrMarshaler)[0].String() != "0.0.0.0/0" {
@@ -2741,6 +2743,7 @@ func TestTokenStore_RoleCRUD(t *testing.T) {
 		"path_suffix":      "happenin",
 		"renewable":        false,
 		"explicit_max_ttl": "80h",
+		"token_num_uses":   0,
 	}
 
 	resp, err = core.HandleRequest(namespace.RootContext(nil), req)
@@ -3992,6 +3995,99 @@ func TestTokenStore_Periodic(t *testing.T) {
 	}
 }
 
+func testTokenStore_NumUses_ErrorCheckHelper(t *testing.T, resp *logical.Response, err error) {
+	if err != nil {
+		t.Fatal(err)
+	}
+	if resp == nil {
+		t.Fatal("response was nil")
+	}
+	if resp.Auth == nil {
+		t.Fatalf(fmt.Sprintf("response auth was nil, resp is %#v", *resp))
+	}
+	if resp.Auth.ClientToken == "" {
+		t.Fatalf("bad: %#v", resp)
+	}
+}
+
+func testTokenStore_NumUses_SelfLookupHelper(t *testing.T, core *Core, clientToken string, expectedNumUses int) {
+	req := logical.TestRequest(t, logical.ReadOperation, "auth/token/lookup-self")
+	req.ClientToken = clientToken
+	resp, err := core.HandleRequest(namespace.RootContext(nil), req)
+	if err != nil {
+		t.Fatalf("err: %v", err)
+	}
+	// Just used the token, this should decrement the num_uses counter
+	expectedNumUses = expectedNumUses - 1
+	actualNumUses := resp.Data["num_uses"].(int)
+
+	if actualNumUses != expectedNumUses {
+		t.Fatalf("num_uses mismatch (expected %d, got %d)", expectedNumUses, actualNumUses)
+	}
+}
+func TestTokenStore_NumUses(t *testing.T) {
+	core, _, root := TestCoreUnsealed(t)
+	roleNumUses := 10
+	lesserNumUses := 5
+	greaterNumUses := 15
+
+	// Create a test role setting token_num_uses
+	req := logical.TestRequest(t, logical.UpdateOperation, "auth/token/roles/test")
+	req.ClientToken = root
+	req.Data = map[string]interface{}{
+		"token_num_uses": roleNumUses,
+	}
+	resp, err := core.HandleRequest(namespace.RootContext(nil), req)
+	if err != nil || (resp != nil && resp.IsError()) {
+		t.Fatalf("err: %v\nresp: %#v", err, resp)
+	}
+	if resp != nil {
+		t.Fatalf("expected a nil response")
+	}
+
+	// Generate some tokens from the test role
+	req.ClientToken = root
+	req.Path = "auth/token/create/test"
+
+	// first token, num_uses is expected to come from the role
+	resp, err = core.HandleRequest(namespace.RootContext(nil), req)
+	testTokenStore_NumUses_ErrorCheckHelper(t, resp, err)
+	noOverrideToken := resp.Auth.ClientToken
+
+	// second token, override num_uses with a lesser value, this should become the value
+	// applied to the token
+	req.Data = map[string]interface{}{
+		"num_uses": lesserNumUses,
+	}
+	resp, err = core.HandleRequest(namespace.RootContext(nil), req)
+	testTokenStore_NumUses_ErrorCheckHelper(t, resp, err)
+	lesserOverrideToken := resp.Auth.ClientToken
+
+	// third token, override num_uses with a greater value, the value
+	// applied to the token should come from the role
+	req.Data = map[string]interface{}{
+		"num_uses": greaterNumUses,
+	}
+	resp, err = core.HandleRequest(namespace.RootContext(nil), req)
+	testTokenStore_NumUses_ErrorCheckHelper(t, resp, err)
+	greaterOverrideToken := resp.Auth.ClientToken
+
+	// fourth token, override num_uses with a zero value, a num_uses value of zero
+	// has an internal meaning of unlimited so num_uses == 1 is actually less than
+	// num_uses == 0. In this case, the lesser value of the role should be applied.
+	req.Data = map[string]interface{}{
+		"num_uses": 0,
+	}
+	resp, err = core.HandleRequest(namespace.RootContext(nil), req)
+	testTokenStore_NumUses_ErrorCheckHelper(t, resp, err)
+	zeroOverrideToken := resp.Auth.ClientToken
+
+	testTokenStore_NumUses_SelfLookupHelper(t, core, noOverrideToken, roleNumUses)
+	testTokenStore_NumUses_SelfLookupHelper(t, core, lesserOverrideToken, lesserNumUses)
+	testTokenStore_NumUses_SelfLookupHelper(t, core, greaterOverrideToken, roleNumUses)
+	testTokenStore_NumUses_SelfLookupHelper(t, core, zeroOverrideToken, roleNumUses)
+}
+
 func TestTokenStore_Periodic_ExplicitMax(t *testing.T) {
 	core, _, root := TestCoreUnsealed(t)