Add a SetupKey usage limit

management/server/account.go

@@ -38,13 +38,8 @@ func cacheEntryExpiration() time.Duration {
 
 type AccountManager interface {
 	GetOrCreateAccountByUser(userId, domain string) (*Account, error)
-	CreateSetupKey(
-		accountId string,
-		keyName string,
-		keyType SetupKeyType,
-		expiresIn time.Duration,
-		autoGroups []string,
-	) (*SetupKey, error)
+	CreateSetupKey(accountID string, keyName string, keyType SetupKeyType, expiresIn time.Duration,
+		autoGroups []string, usageLimit int) (*SetupKey, error)
 	SaveSetupKey(accountID string, key *SetupKey) (*SetupKey, error)
 	CreateUser(accountID string, key *UserInfo) (*UserInfo, error)
 	ListSetupKeys(accountID, userID string) ([]*SetupKey, error)
@@ -945,7 +940,8 @@ func newAccountWithId(accountId, userId, domain string) *Account {
 
 	setupKeys := make(map[string]*SetupKey)
 	defaultKey := GenerateDefaultSetupKey()
-	oneOffKey := GenerateSetupKey("One-off key", SetupKeyOneOff, DefaultSetupKeyDuration, []string{})
+	oneOffKey := GenerateSetupKey("One-off key", SetupKeyOneOff, DefaultSetupKeyDuration, []string{},
+		SetupKeyUnlimitedUsage)
 	setupKeys[defaultKey.Key] = defaultKey
 	setupKeys[oneOffKey.Key] = oneOffKey
 	network := NewNetwork()

management/server/http/api/openapi.yml

@@ -193,6 +193,9 @@ components:
           description: Setup key last update date
           type: string
           format: date-time
+        usage_limit:
+          description: A number of times this key can be used. The value of 0 indicates the unlimited usage.
+          type: integer
       required:
       - id
       - key
@@ -206,6 +209,7 @@ components:
       - state
       - auto_groups
       - updated_at
+      - usage_limit
     SetupKeyRequest:
       type: object
       properties:
@@ -226,12 +230,16 @@ components:
           type: array
           items:
             type: string
+        usage_limit:
+          description: A number of times this key can be used. The value of 0 indicates the unlimited usage.
+          type: integer
       required:
         - name
         - type
         - expires_in
         - revoked
         - auto_groups
+        - usage_limit
     GroupMinimum:
       type: object
       properties:

management/server/http/setupkeys.go

@@ -50,7 +50,7 @@ func (h *SetupKeys) CreateSetupKeyHandler(w http.ResponseWriter, r *http.Request
 
 	if !(server.SetupKeyType(req.Type) == server.SetupKeyReusable ||
 		server.SetupKeyType(req.Type) == server.SetupKeyOneOff) {
-		util.WriteError(status.Errorf(status.InvalidArgument, "unknown setup key type %s", string(req.Type)), w)
+		util.WriteError(status.Errorf(status.InvalidArgument, "unknown setup key type %s", req.Type), w)
 		return
 	}
 
@@ -61,7 +61,7 @@ func (h *SetupKeys) CreateSetupKeyHandler(w http.ResponseWriter, r *http.Request
 	}
 
 	setupKey, err := h.accountManager.CreateSetupKey(account.Id, req.Name, server.SetupKeyType(req.Type), expiresIn,
-		req.AutoGroups)
+		req.AutoGroups, req.UsageLimit)
 	if err != nil {
 		util.WriteError(err, w)
 		return
@@ -201,5 +201,6 @@ func toResponseBody(key *server.SetupKey) *api.SetupKey {
 		State:      state,
 		AutoGroups: key.AutoGroups,
 		UpdatedAt:  key.UpdatedAt,
+		UsageLimit: key.UsageLimit,
 	}
 }

management/server/http/setupkeys_test.go

@@ -46,7 +46,8 @@ func initSetupKeysTestMetaData(defaultKey *server.SetupKey, newKey *server.Setup
 						"id-all":  {ID: "id-all", Name: "All"}},
 				}, user, nil
 			},
-			CreateSetupKeyFunc: func(_ string, keyName string, typ server.SetupKeyType, _ time.Duration, _ []string) (*server.SetupKey, error) {
+			CreateSetupKeyFunc: func(_ string, keyName string, typ server.SetupKeyType, _ time.Duration, _ []string,
+				_ int) (*server.SetupKey, error) {
 				if keyName == newKey.Name || typ != newKey.Type {
 					return newKey, nil
 				}
@@ -93,7 +94,8 @@ func TestSetupKeysHandlers(t *testing.T) {
 
 	adminUser := server.NewAdminUser("test_user")
 
-	newSetupKey := server.GenerateSetupKey(newSetupKeyName, server.SetupKeyReusable, 0, []string{"group-1"})
+	newSetupKey := server.GenerateSetupKey(newSetupKeyName, server.SetupKeyReusable, 0, []string{"group-1"},
+		server.SetupKeyUnlimitedUsage)
 	updatedDefaultSetupKey := defaultSetupKey.Copy()
 	updatedDefaultSetupKey.AutoGroups = []string{"group-1"}
 	updatedDefaultSetupKey.Name = updatedSetupKeyName

management/server/mock_server/account_mock.go

@@ -13,7 +13,7 @@ import (
 type MockAccountManager struct {
 	GetOrCreateAccountByUserFunc    func(userId, domain string) (*server.Account, error)
 	GetAccountByUserFunc            func(userId string) (*server.Account, error)
-	CreateSetupKeyFunc              func(accountId string, keyName string, keyType server.SetupKeyType, expiresIn time.Duration, autoGroups []string) (*server.SetupKey, error)
+	CreateSetupKeyFunc              func(accountId string, keyName string, keyType server.SetupKeyType, expiresIn time.Duration, autoGroups []string, usageLimit int) (*server.SetupKey, error)
 	GetSetupKeyFunc                 func(accountID, userID, keyID string) (*server.SetupKey, error)
 	GetAccountByUserOrAccountIdFunc func(userId, accountId, domain string) (*server.Account, error)
 	IsUserAdminFunc                 func(claims jwtclaims.AuthorizationClaims) (bool, error)
@@ -102,14 +102,15 @@ func (am *MockAccountManager) GetAccountByUser(userId string) (*server.Account,
 
 // CreateSetupKey mock implementation of CreateSetupKey from server.AccountManager interface
 func (am *MockAccountManager) CreateSetupKey(
-	accountId string,
+	accountID string,
 	keyName string,
 	keyType server.SetupKeyType,
 	expiresIn time.Duration,
 	autoGroups []string,
+	usageLimit int,
 ) (*server.SetupKey, error) {
 	if am.CreateSetupKeyFunc != nil {
-		return am.CreateSetupKeyFunc(accountId, keyName, keyType, expiresIn, autoGroups)
+		return am.CreateSetupKeyFunc(accountID, keyName, keyType, expiresIn, autoGroups, usageLimit)
 	}
 	return nil, status.Errorf(codes.Unimplemented, "method CreateSetupKey is not implemented")
 }

management/server/setupkey.go

@@ -20,7 +20,11 @@ const (
 	DefaultSetupKeyDuration = 24 * 30 * time.Hour
 	// DefaultSetupKeyName is a default name of the default setup key
 	DefaultSetupKeyName = "Default key"
+	// SetupKeyUnlimitedUsage indicates an unlimited usage of a setup key
+	SetupKeyUnlimitedUsage = 0
+)
 
+const (
 	// UpdateSetupKeyName indicates a setup key name update operation
 	UpdateSetupKeyName SetupKeyUpdateOperationType = iota
 	// UpdateSetupKeyRevoked indicates a setup key revoked filed update operation
@@ -75,6 +79,9 @@ type SetupKey struct {
 	LastUsed time.Time
 	// AutoGroups is a list of Group IDs that are auto assigned to a Peer when it uses this key to register
 	AutoGroups []string
+	// UsageLimit indicates the number of times this key can be used to enroll a machine.
+	// The value of 0 indicates the unlimited usage.
+	UsageLimit int
 }
 
 // Copy copies SetupKey to a new object
@@ -96,6 +103,7 @@ func (key *SetupKey) Copy() *SetupKey {
 		UsedTimes:  key.UsedTimes,
 		LastUsed:   key.LastUsed,
 		AutoGroups: autoGroups,
+		UsageLimit: key.UsageLimit,
 	}
 }
 
@@ -131,14 +139,23 @@ func (key *SetupKey) IsExpired() bool {
 	return time.Now().After(key.ExpiresAt)
 }
 
-// IsOverUsed if key was used too many times
+// IsOverUsed if the key was used too many times. SetupKey.UsageLimit == 0 indicates the unlimited usage.
 func (key *SetupKey) IsOverUsed() bool {
-	return key.Type == SetupKeyOneOff && key.UsedTimes >= 1
+	limit := key.UsageLimit
+	if key.Type == SetupKeyOneOff {
+		limit = 1
+	}
+	return limit > 0 && key.UsedTimes >= limit
 }
 
 // GenerateSetupKey generates a new setup key
-func GenerateSetupKey(name string, t SetupKeyType, validFor time.Duration, autoGroups []string) *SetupKey {
+func GenerateSetupKey(name string, t SetupKeyType, validFor time.Duration, autoGroups []string,
+	usageLimit int) *SetupKey {
 	key := strings.ToUpper(uuid.New().String())
+	limit := usageLimit
+	if t == SetupKeyOneOff {
+		limit = 1
+	}
 	return &SetupKey{
 		Id:         strconv.Itoa(int(Hash(key))),
 		Key:        key,
@@ -150,12 +167,14 @@ func GenerateSetupKey(name string, t SetupKeyType, validFor time.Duration, autoG
 		Revoked:    false,
 		UsedTimes:  0,
 		AutoGroups: autoGroups,
+		UsageLimit: limit,
 	}
 }
 
-// GenerateDefaultSetupKey generates a default setup key
+// GenerateDefaultSetupKey generates a default reusable setup key with an unlimited usage and 30 days expiration
 func GenerateDefaultSetupKey() *SetupKey {
-	return GenerateSetupKey(DefaultSetupKeyName, SetupKeyReusable, DefaultSetupKeyDuration, []string{})
+	return GenerateSetupKey(DefaultSetupKeyName, SetupKeyReusable, DefaultSetupKeyDuration, []string{},
+		SetupKeyUnlimitedUsage)
 }
 
 func Hash(s string) uint32 {
@@ -170,7 +189,7 @@ func Hash(s string) uint32 {
 // CreateSetupKey generates a new setup key with a given name, type, list of groups IDs to auto-assign to peers registered with this key,
 // and adds it to the specified account. A list of autoGroups IDs can be empty.
 func (am *DefaultAccountManager) CreateSetupKey(accountID string, keyName string, keyType SetupKeyType,
-	expiresIn time.Duration, autoGroups []string) (*SetupKey, error) {
+	expiresIn time.Duration, autoGroups []string, usageLimit int) (*SetupKey, error) {
 	unlock := am.Store.AcquireAccountLock(accountID)
 	defer unlock()
 
@@ -190,7 +209,7 @@ func (am *DefaultAccountManager) CreateSetupKey(accountID string, keyName string
 		}
 	}
 
-	setupKey := GenerateSetupKey(keyName, keyType, keyDuration, autoGroups)
+	setupKey := GenerateSetupKey(keyName, keyType, keyDuration, autoGroups, usageLimit)
 	account.SetupKeys[setupKey.Key] = setupKey
 
 	err = am.Store.SaveAccount(account)

management/server/setupkey_test.go

@@ -32,7 +32,8 @@ func TestDefaultAccountManager_SaveSetupKey(t *testing.T) {
 	expiresIn := time.Hour
 	keyName := "my-test-key"
 
-	key, err := manager.CreateSetupKey(account.Id, keyName, SetupKeyReusable, expiresIn, []string{})
+	key, err := manager.CreateSetupKey(account.Id, keyName, SetupKeyReusable, expiresIn, []string{},
+		SetupKeyUnlimitedUsage)
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -120,7 +121,7 @@ func TestDefaultAccountManager_CreateSetupKey(t *testing.T) {
 	for _, tCase := range []testCase{testCase1, testCase2} {
 		t.Run(tCase.name, func(t *testing.T) {
 			key, err := manager.CreateSetupKey(account.Id, tCase.expectedKeyName, SetupKeyReusable, expiresIn,
-				tCase.expectedGroups)
+				tCase.expectedGroups, SetupKeyUnlimitedUsage)
 
 			if tCase.expectedFailure {
 				if err == nil {
@@ -168,41 +169,41 @@ func TestGenerateSetupKey(t *testing.T) {
 	expectedUpdatedAt := time.Now()
 	var expectedAutoGroups []string
 
-	key := GenerateSetupKey(expectedName, SetupKeyOneOff, time.Hour, []string{})
+	key := GenerateSetupKey(expectedName, SetupKeyOneOff, time.Hour, []string{}, SetupKeyUnlimitedUsage)
 
 	assertKey(t, key, expectedName, expectedRevoke, expectedType, expectedUsedTimes, expectedCreatedAt,
 		expectedExpiresAt, strconv.Itoa(int(Hash(key.Key))), expectedUpdatedAt, expectedAutoGroups)
 
 }
 
 func TestSetupKey_IsValid(t *testing.T) {
-	validKey := GenerateSetupKey("valid key", SetupKeyOneOff, time.Hour, []string{})
+	validKey := GenerateSetupKey("valid key", SetupKeyOneOff, time.Hour, []string{}, SetupKeyUnlimitedUsage)
 	if !validKey.IsValid() {
 		t.Errorf("expected key to be valid, got invalid %v", validKey)
 	}
 
 	// expired
-	expiredKey := GenerateSetupKey("invalid key", SetupKeyOneOff, -time.Hour, []string{})
+	expiredKey := GenerateSetupKey("invalid key", SetupKeyOneOff, -time.Hour, []string{}, SetupKeyUnlimitedUsage)
 	if expiredKey.IsValid() {
 		t.Errorf("expected key to be invalid due to expiration, got valid %v", expiredKey)
 	}
 
 	// revoked
-	revokedKey := GenerateSetupKey("invalid key", SetupKeyOneOff, time.Hour, []string{})
+	revokedKey := GenerateSetupKey("invalid key", SetupKeyOneOff, time.Hour, []string{}, SetupKeyUnlimitedUsage)
 	revokedKey.Revoked = true
 	if revokedKey.IsValid() {
 		t.Errorf("expected revoked key to be invalid, got valid %v", revokedKey)
 	}
 
 	// overused
-	overUsedKey := GenerateSetupKey("invalid key", SetupKeyOneOff, time.Hour, []string{})
+	overUsedKey := GenerateSetupKey("invalid key", SetupKeyOneOff, time.Hour, []string{}, SetupKeyUnlimitedUsage)
 	overUsedKey.UsedTimes = 1
 	if overUsedKey.IsValid() {
 		t.Errorf("expected overused key to be invalid, got valid %v", overUsedKey)
 	}
 
 	// overused
-	reusableKey := GenerateSetupKey("valid key", SetupKeyReusable, time.Hour, []string{})
+	reusableKey := GenerateSetupKey("valid key", SetupKeyReusable, time.Hour, []string{}, SetupKeyUnlimitedUsage)
 	reusableKey.UsedTimes = 99
 	if !reusableKey.IsValid() {
 		t.Errorf("expected reusable key to be valid when used many times, got valid %v", reusableKey)
@@ -257,7 +258,7 @@ func assertKey(t *testing.T, key *SetupKey, expectedName string, expectedRevoke
 
 func TestSetupKey_Copy(t *testing.T) {
 
-	key := GenerateSetupKey("key name", SetupKeyOneOff, time.Hour, []string{})
+	key := GenerateSetupKey("key name", SetupKeyOneOff, time.Hour, []string{}, SetupKeyUnlimitedUsage)
 	keyCopy := key.Copy()
 
 	assertKey(t, keyCopy, key.Name, key.Revoked, string(key.Type), key.UsedTimes, key.CreatedAt, key.ExpiresAt, key.Id,