Merge pull request #2271 from gravitl/NET-83i

Net 83i

logic/enrollmentkey.go

@@ -41,12 +41,16 @@ func CreateEnrollmentKey(uses int, expiration time.Time, networks, tags []string
 		Unlimited:     unlimited,
 		Networks:      []string{},
 		Tags:          []string{},
+		Type:          models.Undefined,
 	}
 	if uses > 0 {
 		k.UsesRemaining = uses
-	}
-	if !expiration.IsZero() {
+		k.Type = models.Uses
+	} else if !expiration.IsZero() {
 		k.Expiration = expiration
+		k.Type = models.TimeExpiration
+	} else if k.Unlimited {
+		k.Type = models.Unlimited
 	}
 	if len(networks) > 0 {
 		k.Networks = networks

main.go

@@ -19,6 +19,7 @@ import (
 	"github.com/gravitl/netmaker/logger"
 	"github.com/gravitl/netmaker/logic"
 	"github.com/gravitl/netmaker/logic/pro"
+	"github.com/gravitl/netmaker/migrate"
 	"github.com/gravitl/netmaker/models"
 	"github.com/gravitl/netmaker/mq"
 	"github.com/gravitl/netmaker/netclient/ncutils"
@@ -27,7 +28,7 @@ import (
 	stunserver "github.com/gravitl/netmaker/stun-server"
 )
 
-var version = "v0.19.0"
+var version = "v0.19.1"
 
 // Start DB Connection and start API Request Handler
 func main() {
@@ -74,6 +75,7 @@ func initialize() { // Client Mode Prereq Check
 		logger.FatalLog("Error connecting to database: ", err.Error())
 	}
 	logger.Log(0, "database successfully connected")
+	migrate.Run()
 
 	logic.SetJWTSecret()
 

migrate/migrate.go

@@ -0,0 +1,50 @@
+package migrate
+
+import (
+	"encoding/json"
+
+	"github.com/gravitl/netmaker/database"
+	"github.com/gravitl/netmaker/logger"
+	"github.com/gravitl/netmaker/models"
+)
+
+// Run - runs all migrations
+func Run() {
+	updateEnrollmentKeys()
+}
+
+func updateEnrollmentKeys() {
+	rows, err := database.FetchRecords(database.ENROLLMENT_KEYS_TABLE_NAME)
+	if err != nil {
+		return
+	}
+	for _, row := range rows {
+		var key models.EnrollmentKey
+		if err = json.Unmarshal([]byte(row), &key); err != nil {
+			continue
+		}
+		if key.Type != models.Undefined {
+			logger.Log(2, "migration: enrollment key type already set")
+			continue
+		} else {
+			logger.Log(2, "migration: updating enrollment key type")
+			if key.Unlimited {
+				key.Type = models.Unlimited
+			} else if key.UsesRemaining > 0 {
+				key.Type = models.Uses
+			} else if !key.Expiration.IsZero() {
+				key.Type = models.TimeExpiration
+			}
+		}
+		data, err := json.Marshal(key)
+		if err != nil {
+			logger.Log(0, "migration: marshalling enrollment key: "+err.Error())
+			continue
+		}
+		if err = database.Insert(key.Value, string(data), database.ENROLLMENT_KEYS_TABLE_NAME); err != nil {
+			logger.Log(0, "migration: inserting enrollment key: "+err.Error())
+			continue
+		}
+
+	}
+}

models/enrollment_key.go

@@ -4,6 +4,21 @@ import (
 	"time"
 )
 
+const (
+	Undefined KeyType = iota
+	TimeExpiration
+	Uses
+	Unlimited
+)
+
+// KeyType - the type of enrollment key
+type KeyType int
+
+// String - returns the string representation of a KeyType
+func (k KeyType) String() string {
+	return [...]string{"Undefined", "TimeExpiration", "Uses", "Unlimited"}[k]
+}
+
 // EnrollmentToken - the tokenized version of an enrollmentkey;
 // to be used for host registration
 type EnrollmentToken struct {
@@ -23,6 +38,7 @@ type EnrollmentKey struct {
 	Unlimited     bool      `json:"unlimited"`
 	Tags          []string  `json:"tags"`
 	Token         string    `json:"token,omitempty"` // B64 value of EnrollmentToken
+	Type          KeyType   `json:"type"`
 }
 
 // APIEnrollmentKey - used to create enrollment keys via API
@@ -32,6 +48,7 @@ type APIEnrollmentKey struct {
 	Networks      []string `json:"networks"`
 	Unlimited     bool     `json:"unlimited"`
 	Tags          []string `json:"tags"`
+	Type          KeyType  `json:"type"`
 }
 
 // RegisterResponse - the response to a successful enrollment register
@@ -51,6 +68,9 @@ func (k *EnrollmentKey) IsValid() bool {
 	if !k.Expiration.IsZero() && time.Now().Before(k.Expiration) {
 		return true
 	}
+	if k.Type == Undefined {
+		return false
+	}
 
 	return k.Unlimited
 }