feature: extend setup key logic

management/server/account.go

@@ -24,19 +24,13 @@ type Account struct {
 	Peers     map[string]*Peer
 }
 
-// SetupKey represents a pre-authorized key used to register machines (peers)
-// One key might have multiple machines
-type SetupKey struct {
-	Key string
-}
-
 // Peer represents a machine connected to the network.
 // The Peer is a Wireguard peer identified by a public key
 type Peer struct {
 	// Wireguard public key
 	Key string
 	// A setup key this peer was registered with
-	SetupKey *SetupKey
+	SetupKey string
 	// IP address of the Peer
 	IP net.IP
 }
@@ -164,23 +158,39 @@ func (manager *AccountManager) createAccount(accountId string) (*Account, error)
 // Each Account has a list of pre-authorised SetupKey and if no Account has a given key err wit ha code codes.Unauthenticated
 // will be returned, meaning the key is invalid
 // Each new Peer will be assigned a new next net.IP from the Account.Network and Account.Network.LastIP will be updated (IP's are not reused).
-// If the specified setupKey is empty then a new Account will be created //todo make it more explicit?
+// If the specified setupKey is empty then a new Account will be created //todo remove this part
 func (manager *AccountManager) AddPeer(setupKey string, peerKey string) (*Peer, error) {
 	manager.mux.Lock()
 	defer manager.mux.Unlock()
 
+	upperKey := strings.ToUpper(setupKey)
+
 	var account *Account
 	var err error
 	var sk *SetupKey
-	if len(setupKey) == 0 {
+	if len(upperKey) == 0 {
 		// Empty setup key, create a new account for it.
 		account, sk = newAccount()
 	} else {
-		sk = &SetupKey{Key: strings.ToUpper(setupKey)}
+		account, err = manager.Store.GetAccountBySetupKey(upperKey)
-		account, err = manager.Store.GetAccountBySetupKey(sk.Key)
 		if err != nil {
-			return nil, status.Errorf(codes.NotFound, "unknown setupKey %s", setupKey)
+			return nil, status.Errorf(codes.NotFound, "unknown setupKey %s", upperKey)
 		}
+
+		for _, key := range account.SetupKeys {
+			if upperKey == key.Key {
+				sk = key
+			}
+		}
+
+		if sk == nil {
+			// shouldn't happen actually
+			return nil, status.Errorf(codes.NotFound, "unknown setupKey %s", upperKey)
+		}
+	}
+
+	if !sk.IsValid() {
+		return nil, status.Errorf(codes.FailedPrecondition, "setup key was expired or overused %s", upperKey)
 	}
 
 	var takenIps []net.IP
@@ -193,7 +203,7 @@ func (manager *AccountManager) AddPeer(setupKey string, peerKey string) (*Peer,
 
 	newPeer := &Peer{
 		Key:      peerKey,
-		SetupKey: sk,
+		SetupKey: sk.Key,
 		IP:       nextIp,
 	}
 
@@ -212,17 +222,16 @@ func newAccountWithId(accountId string) (*Account, *SetupKey) {
 
 	log.Debugf("creating new account")
 
-	setupKeyId := strings.ToUpper(uuid.New().String())
 	setupKeys := make(map[string]*SetupKey)
-	setupKey := &SetupKey{Key: setupKeyId}
+	setupKey := GenerateSetupKey(DefaultSetupKeyName, SetupKeyReusable, DefaultSetupKeyDuration)
-	setupKeys[setupKeyId] = setupKey
+	setupKeys[setupKey.Key] = setupKey
 	network := &Network{
 		Id:  uuid.New().String(),
 		Net: net.IPNet{IP: net.ParseIP("100.64.0.0"), Mask: net.IPMask{255, 192, 0, 0}},
 		Dns: ""}
 	peers := make(map[string]*Peer)
 
-	log.Debugf("created new account %s with setup key %s", accountId, setupKeyId)
+	log.Debugf("created new account %s with setup key %s", accountId, setupKey.Key)
 
 	return &Account{Id: accountId, SetupKeys: setupKeys, Network: network, Peers: peers}, setupKey
 }

management/server/http/handler/setupkeys.go

@@ -19,6 +19,7 @@ type SetupKeyResponse struct {
 	Name    string
 	Expires time.Time
 	Type    string
+	Valid   bool
 }
 
 func NewSetupKeysHandler(accountManager *server.AccountManager) *SetupKeys {
@@ -45,9 +46,10 @@ func (h *SetupKeys) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 		for _, key := range account.SetupKeys {
 			respBody = append(respBody, &SetupKeyResponse{
 				Key:     key.Key,
-				Name:    key.Key,
+				Name:    key.Name,
-				Expires: time.Now().Add(30 * (time.Second * 60 * 60 * 24)),
+				Expires: key.ExpiresAt,
-				Type:    "reusable",
+				Type:    string(key.Type),
+				Valid:   key.IsValid(),
 			})
 		}
 

management/server/setupkey.go

@@ -0,0 +1,56 @@
+package server
+
+import (
+	"github.com/google/uuid"
+	"strings"
+	"time"
+)
+
+const (
+	// SetupKeyReusable is a multi-use key (can be used for multiple machines)
+	SetupKeyReusable SetupKeyType = "reusable"
+	// SetupKeyOneOff is a single use key (can be used only once)
+	SetupKeyOneOff SetupKeyType = "one-off"
+
+	// DefaultSetupKeyDuration = 1 month
+	DefaultSetupKeyDuration = 24 * 30 * time.Hour
+	// DefaultSetupKeyName is a default name of the default setup key
+	DefaultSetupKeyName = "Default key"
+)
+
+// SetupKeyType is the type of setup key
+type SetupKeyType string
+
+// SetupKey represents a pre-authorized key used to register machines (peers)
+type SetupKey struct {
+	Key       string
+	Name      string
+	Type      SetupKeyType
+	CreatedAt time.Time
+	ExpiresAt time.Time
+	// Revoked indicates whether the key was revoked or not (we don't remove them for tracking purposes)
+	Revoked bool
+	// UsedTimes indicates how many times the key was used
+	UsedTimes int
+}
+
+// IsValid is true if the key was not revoked, is not expired and used not more than it was supposed to
+func (key *SetupKey) IsValid() bool {
+	expired := time.Now().After(key.ExpiresAt)
+	overUsed := key.Type == SetupKeyOneOff && key.UsedTimes >= 1
+	return !key.Revoked && !expired && !overUsed
+}
+
+// GenerateSetupKey generates a new setup key
+func GenerateSetupKey(name string, t SetupKeyType, validFor time.Duration) *SetupKey {
+	createdAt := time.Now()
+	return &SetupKey{
+		Key:       strings.ToUpper(uuid.New().String()),
+		Name:      name,
+		Type:      t,
+		CreatedAt: createdAt,
+		ExpiresAt: createdAt.Add(validFor),
+		Revoked:   false,
+		UsedTimes: 0,
+	}
+}