diff --git a/management/server/account.go b/management/server/account.go index 2b4e2e97d2b..2d434e2d9df 100644 --- a/management/server/account.go +++ b/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() diff --git a/management/server/http/api/openapi.yml b/management/server/http/api/openapi.yml index 377b36937aa..96a90053a3f 100644 --- a/management/server/http/api/openapi.yml +++ b/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: diff --git a/management/server/http/api/types.gen.go b/management/server/http/api/types.gen.go index 9b178f4f5d6..c42a0e6b113 100644 --- a/management/server/http/api/types.gen.go +++ b/management/server/http/api/types.gen.go @@ -449,6 +449,9 @@ type SetupKey struct { // UpdatedAt Setup key last update date UpdatedAt time.Time `json:"updated_at"` + // UsageLimit A number of times this key can be used. The value of 0 indicates the unlimited usage. + UsageLimit int `json:"usage_limit"` + // UsedTimes Usage count of setup key UsedTimes int `json:"used_times"` @@ -472,6 +475,9 @@ type SetupKeyRequest struct { // Type Setup key type, one-off for single time usage and reusable Type string `json:"type"` + + // UsageLimit A number of times this key can be used. The value of 0 indicates the unlimited usage. + UsageLimit int `json:"usage_limit"` } // User defines model for User. diff --git a/management/server/http/setupkeys.go b/management/server/http/setupkeys.go index 6af797eb0d3..a88cba29e84 100644 --- a/management/server/http/setupkeys.go +++ b/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, } } diff --git a/management/server/http/setupkeys_test.go b/management/server/http/setupkeys_test.go index fbb8a9f2c9e..2b26ebb5a77 100644 --- a/management/server/http/setupkeys_test.go +++ b/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 diff --git a/management/server/mock_server/account_mock.go b/management/server/mock_server/account_mock.go index 88c89fbff41..3030b831ddd 100644 --- a/management/server/mock_server/account_mock.go +++ b/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") } diff --git a/management/server/setupkey.go b/management/server/setupkey.go index 9d7d3637915..824d381a2ec 100644 --- a/management/server/setupkey.go +++ b/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) diff --git a/management/server/setupkey_test.go b/management/server/setupkey_test.go index 08b6d8ad084..290259116cd 100644 --- a/management/server/setupkey_test.go +++ b/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,7 +169,7 @@ 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) @@ -176,33 +177,33 @@ func TestGenerateSetupKey(t *testing.T) { } 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,