add tag groups to enrollment key

This commit is contained in:
abhishek9686 2024-09-18 12:22:49 +04:00
parent 6640b66ac6
commit db2550b7bd
4 changed files with 29 additions and 17 deletions

View file

@ -156,6 +156,7 @@ func createEnrollmentKey(w http.ResponseWriter, r *http.Request) {
newTime, newTime,
enrollmentKeyBody.Networks, enrollmentKeyBody.Networks,
enrollmentKeyBody.Tags, enrollmentKeyBody.Tags,
enrollmentKeyBody.Groups,
enrollmentKeyBody.Unlimited, enrollmentKeyBody.Unlimited,
relayId, relayId,
) )
@ -206,7 +207,7 @@ func updateEnrollmentKey(w http.ResponseWriter, r *http.Request) {
} }
} }
newEnrollmentKey, err := logic.UpdateEnrollmentKey(keyId, relayId) newEnrollmentKey, err := logic.UpdateEnrollmentKey(keyId, relayId, enrollmentKeyBody.Groups)
if err != nil { if err != nil {
slog.Error("failed to update enrollment key", "error", err) slog.Error("failed to update enrollment key", "error", err)
logic.ReturnErrorResponse(w, r, logic.FormatError(err, "internal")) logic.ReturnErrorResponse(w, r, logic.FormatError(err, "internal"))
@ -307,6 +308,10 @@ func handleHostRegister(w http.ResponseWriter, r *http.Request) {
return return
} }
} }
newHost.Tags = make(map[models.TagID]struct{})
for _, tagI := range enrollmentKey.Groups {
newHost.Tags[tagI] = struct{}{}
}
if err = logic.CreateHost(&newHost); err != nil { if err = logic.CreateHost(&newHost); err != nil {
logger.Log( logger.Log(
0, 0,
@ -337,6 +342,10 @@ func handleHostRegister(w http.ResponseWriter, r *http.Request) {
return return
} }
logic.UpdateHostFromClient(&newHost, currHost) logic.UpdateHostFromClient(&newHost, currHost)
currHost.Tags = make(map[models.TagID]struct{})
for _, tagI := range enrollmentKey.Groups {
currHost.Tags[tagI] = struct{}{}
}
err = logic.UpsertHost(currHost) err = logic.UpsertHost(currHost)
if err != nil { if err != nil {
slog.Error("failed to update host", "id", currHost.ID, "error", err) slog.Error("failed to update host", "id", currHost.ID, "error", err)

View file

@ -37,7 +37,7 @@ var (
) )
// CreateEnrollmentKey - creates a new enrollment key in db // CreateEnrollmentKey - creates a new enrollment key in db
func CreateEnrollmentKey(uses int, expiration time.Time, networks, tags []string, unlimited bool, relay uuid.UUID) (*models.EnrollmentKey, error) { func CreateEnrollmentKey(uses int, expiration time.Time, networks, tags []string, groups []models.TagID, unlimited bool, relay uuid.UUID) (*models.EnrollmentKey, error) {
newKeyID, err := getUniqueEnrollmentID() newKeyID, err := getUniqueEnrollmentID()
if err != nil { if err != nil {
return nil, err return nil, err
@ -51,6 +51,7 @@ func CreateEnrollmentKey(uses int, expiration time.Time, networks, tags []string
Tags: []string{}, Tags: []string{},
Type: models.Undefined, Type: models.Undefined,
Relay: relay, Relay: relay,
Groups: groups,
} }
if uses > 0 { if uses > 0 {
k.UsesRemaining = uses k.UsesRemaining = uses
@ -89,7 +90,7 @@ func CreateEnrollmentKey(uses int, expiration time.Time, networks, tags []string
} }
// UpdateEnrollmentKey - updates an existing enrollment key's associated relay // UpdateEnrollmentKey - updates an existing enrollment key's associated relay
func UpdateEnrollmentKey(keyId string, relayId uuid.UUID) (*models.EnrollmentKey, error) { func UpdateEnrollmentKey(keyId string, relayId uuid.UUID, groups []models.TagID) (*models.EnrollmentKey, error) {
key, err := GetEnrollmentKey(keyId) key, err := GetEnrollmentKey(keyId)
if err != nil { if err != nil {
return nil, err return nil, err
@ -109,7 +110,7 @@ func UpdateEnrollmentKey(keyId string, relayId uuid.UUID) (*models.EnrollmentKey
} }
key.Relay = relayId key.Relay = relayId
key.Groups = groups
if err = upsertEnrollmentKey(&key); err != nil { if err = upsertEnrollmentKey(&key); err != nil {
return nil, err return nil, err
} }

View file

@ -14,35 +14,35 @@ func TestCreateEnrollmentKey(t *testing.T) {
database.InitializeDatabase() database.InitializeDatabase()
defer database.CloseDB() defer database.CloseDB()
t.Run("Can_Not_Create_Key", func(t *testing.T) { t.Run("Can_Not_Create_Key", func(t *testing.T) {
newKey, err := CreateEnrollmentKey(0, time.Time{}, nil, nil, false, uuid.Nil) newKey, err := CreateEnrollmentKey(0, time.Time{}, nil, nil, nil, false, uuid.Nil)
assert.Nil(t, newKey) assert.Nil(t, newKey)
assert.NotNil(t, err) assert.NotNil(t, err)
assert.ErrorIs(t, err, models.ErrInvalidEnrollmentKey) assert.ErrorIs(t, err, models.ErrInvalidEnrollmentKey)
}) })
t.Run("Can_Create_Key_Uses", func(t *testing.T) { t.Run("Can_Create_Key_Uses", func(t *testing.T) {
newKey, err := CreateEnrollmentKey(1, time.Time{}, nil, nil, false, uuid.Nil) newKey, err := CreateEnrollmentKey(1, time.Time{}, nil, nil, nil, false, uuid.Nil)
assert.Nil(t, err) assert.Nil(t, err)
assert.Equal(t, 1, newKey.UsesRemaining) assert.Equal(t, 1, newKey.UsesRemaining)
assert.True(t, newKey.IsValid()) assert.True(t, newKey.IsValid())
}) })
t.Run("Can_Create_Key_Time", func(t *testing.T) { t.Run("Can_Create_Key_Time", func(t *testing.T) {
newKey, err := CreateEnrollmentKey(0, time.Now().Add(time.Minute), nil, nil, false, uuid.Nil) newKey, err := CreateEnrollmentKey(0, time.Now().Add(time.Minute), nil, nil, nil, false, uuid.Nil)
assert.Nil(t, err) assert.Nil(t, err)
assert.True(t, newKey.IsValid()) assert.True(t, newKey.IsValid())
}) })
t.Run("Can_Create_Key_Unlimited", func(t *testing.T) { t.Run("Can_Create_Key_Unlimited", func(t *testing.T) {
newKey, err := CreateEnrollmentKey(0, time.Time{}, nil, nil, true, uuid.Nil) newKey, err := CreateEnrollmentKey(0, time.Time{}, nil, nil, nil, true, uuid.Nil)
assert.Nil(t, err) assert.Nil(t, err)
assert.True(t, newKey.IsValid()) assert.True(t, newKey.IsValid())
}) })
t.Run("Can_Create_Key_WithNetworks", func(t *testing.T) { t.Run("Can_Create_Key_WithNetworks", func(t *testing.T) {
newKey, err := CreateEnrollmentKey(0, time.Time{}, []string{"mynet", "skynet"}, nil, true, uuid.Nil) newKey, err := CreateEnrollmentKey(0, time.Time{}, []string{"mynet", "skynet"}, nil, nil, true, uuid.Nil)
assert.Nil(t, err) assert.Nil(t, err)
assert.True(t, newKey.IsValid()) assert.True(t, newKey.IsValid())
assert.True(t, len(newKey.Networks) == 2) assert.True(t, len(newKey.Networks) == 2)
}) })
t.Run("Can_Create_Key_WithTags", func(t *testing.T) { t.Run("Can_Create_Key_WithTags", func(t *testing.T) {
newKey, err := CreateEnrollmentKey(0, time.Time{}, nil, []string{"tag1", "tag2"}, true, uuid.Nil) newKey, err := CreateEnrollmentKey(0, time.Time{}, nil, []string{"tag1", "tag2"}, nil, true, uuid.Nil)
assert.Nil(t, err) assert.Nil(t, err)
assert.True(t, newKey.IsValid()) assert.True(t, newKey.IsValid())
assert.True(t, len(newKey.Tags) == 2) assert.True(t, len(newKey.Tags) == 2)
@ -62,7 +62,7 @@ func TestCreateEnrollmentKey(t *testing.T) {
func TestDelete_EnrollmentKey(t *testing.T) { func TestDelete_EnrollmentKey(t *testing.T) {
database.InitializeDatabase() database.InitializeDatabase()
defer database.CloseDB() defer database.CloseDB()
newKey, _ := CreateEnrollmentKey(0, time.Time{}, []string{"mynet", "skynet"}, nil, true, uuid.Nil) newKey, _ := CreateEnrollmentKey(0, time.Time{}, []string{"mynet", "skynet"}, nil, nil, true, uuid.Nil)
t.Run("Can_Delete_Key", func(t *testing.T) { t.Run("Can_Delete_Key", func(t *testing.T) {
assert.True(t, newKey.IsValid()) assert.True(t, newKey.IsValid())
err := DeleteEnrollmentKey(newKey.Value) err := DeleteEnrollmentKey(newKey.Value)
@ -83,7 +83,7 @@ func TestDelete_EnrollmentKey(t *testing.T) {
func TestDecrement_EnrollmentKey(t *testing.T) { func TestDecrement_EnrollmentKey(t *testing.T) {
database.InitializeDatabase() database.InitializeDatabase()
defer database.CloseDB() defer database.CloseDB()
newKey, _ := CreateEnrollmentKey(1, time.Time{}, nil, nil, false, uuid.Nil) newKey, _ := CreateEnrollmentKey(1, time.Time{}, nil, nil, nil, false, uuid.Nil)
t.Run("Check_initial_uses", func(t *testing.T) { t.Run("Check_initial_uses", func(t *testing.T) {
assert.True(t, newKey.IsValid()) assert.True(t, newKey.IsValid())
assert.Equal(t, newKey.UsesRemaining, 1) assert.Equal(t, newKey.UsesRemaining, 1)
@ -107,9 +107,9 @@ func TestDecrement_EnrollmentKey(t *testing.T) {
func TestUsability_EnrollmentKey(t *testing.T) { func TestUsability_EnrollmentKey(t *testing.T) {
database.InitializeDatabase() database.InitializeDatabase()
defer database.CloseDB() defer database.CloseDB()
key1, _ := CreateEnrollmentKey(1, time.Time{}, nil, nil, false, uuid.Nil) key1, _ := CreateEnrollmentKey(1, time.Time{}, nil, nil, nil, false, uuid.Nil)
key2, _ := CreateEnrollmentKey(0, time.Now().Add(time.Minute<<4), nil, nil, false, uuid.Nil) key2, _ := CreateEnrollmentKey(0, time.Now().Add(time.Minute<<4), nil, nil, nil, false, uuid.Nil)
key3, _ := CreateEnrollmentKey(0, time.Time{}, nil, nil, true, uuid.Nil) key3, _ := CreateEnrollmentKey(0, time.Time{}, nil, nil, nil, true, uuid.Nil)
t.Run("Check if valid use key can be used", func(t *testing.T) { t.Run("Check if valid use key can be used", func(t *testing.T) {
assert.Equal(t, key1.UsesRemaining, 1) assert.Equal(t, key1.UsesRemaining, 1)
ok := TryToUseEnrollmentKey(key1) ok := TryToUseEnrollmentKey(key1)
@ -145,7 +145,7 @@ func removeAllEnrollments() {
func TestTokenize_EnrollmentKeys(t *testing.T) { func TestTokenize_EnrollmentKeys(t *testing.T) {
database.InitializeDatabase() database.InitializeDatabase()
defer database.CloseDB() defer database.CloseDB()
newKey, _ := CreateEnrollmentKey(0, time.Time{}, []string{"mynet", "skynet"}, nil, true, uuid.Nil) newKey, _ := CreateEnrollmentKey(0, time.Time{}, []string{"mynet", "skynet"}, nil, nil, true, uuid.Nil)
const defaultValue = "MwE5MwE5MwE5MwE5MwE5MwE5MwE5MwE5" const defaultValue = "MwE5MwE5MwE5MwE5MwE5MwE5MwE5MwE5"
const b64value = "eyJzZXJ2ZXIiOiJhcGkubXlzZXJ2ZXIuY29tIiwidmFsdWUiOiJNd0U1TXdFNU13RTVNd0U1TXdFNU13RTVNd0U1TXdFNSJ9" const b64value = "eyJzZXJ2ZXIiOiJhcGkubXlzZXJ2ZXIuY29tIiwidmFsdWUiOiJNd0U1TXdFNU13RTVNd0U1TXdFNU13RTVNd0U1TXdFNSJ9"
const serverAddr = "api.myserver.com" const serverAddr = "api.myserver.com"
@ -178,7 +178,7 @@ func TestTokenize_EnrollmentKeys(t *testing.T) {
func TestDeTokenize_EnrollmentKeys(t *testing.T) { func TestDeTokenize_EnrollmentKeys(t *testing.T) {
database.InitializeDatabase() database.InitializeDatabase()
defer database.CloseDB() defer database.CloseDB()
newKey, _ := CreateEnrollmentKey(0, time.Time{}, []string{"mynet", "skynet"}, nil, true, uuid.Nil) newKey, _ := CreateEnrollmentKey(0, time.Time{}, []string{"mynet", "skynet"}, nil, nil, true, uuid.Nil)
const b64Value = "eyJzZXJ2ZXIiOiJhcGkubXlzZXJ2ZXIuY29tIiwidmFsdWUiOiJNd0U1TXdFNU13RTVNd0U1TXdFNU13RTVNd0U1TXdFNSJ9" const b64Value = "eyJzZXJ2ZXIiOiJhcGkubXlzZXJ2ZXIuY29tIiwidmFsdWUiOiJNd0U1TXdFNU13RTVNd0U1TXdFNU13RTVNd0U1TXdFNSJ9"
const serverAddr = "api.myserver.com" const serverAddr = "api.myserver.com"

View file

@ -52,6 +52,7 @@ type EnrollmentKey struct {
Token string `json:"token,omitempty"` // B64 value of EnrollmentToken Token string `json:"token,omitempty"` // B64 value of EnrollmentToken
Type KeyType `json:"type"` Type KeyType `json:"type"`
Relay uuid.UUID `json:"relay"` Relay uuid.UUID `json:"relay"`
Groups []TagID `json:"groups"`
} }
// APIEnrollmentKey - used to create enrollment keys via API // APIEnrollmentKey - used to create enrollment keys via API
@ -63,6 +64,7 @@ type APIEnrollmentKey struct {
Tags []string `json:"tags" validate:"required,dive,min=3,max=32"` Tags []string `json:"tags" validate:"required,dive,min=3,max=32"`
Type KeyType `json:"type"` Type KeyType `json:"type"`
Relay string `json:"relay"` Relay string `json:"relay"`
Groups []TagID `json:"groups"`
} }
// RegisterResponse - the response to a successful enrollment register // RegisterResponse - the response to a successful enrollment register