From 6869b48905f4170e32edecd5ab4637d58e684c7a Mon Sep 17 00:00:00 2001
From: Mikhail Bragin <bangvalo@gmail.com>
Date: Sun, 22 Aug 2021 11:29:25 +0200
Subject: [PATCH] feature: increase key usage after successful peer
 registration (#93)

---
 management/server/account.go                |  1 +
 management/server/http/handler/setupkeys.go | 37 +++++++++++++--------
 management/server/setupkey.go               | 10 ++++++
 3 files changed, 34 insertions(+), 14 deletions(-)

diff --git a/management/server/account.go b/management/server/account.go
index a89503e07..2990ef210 100644
--- a/management/server/account.go
+++ b/management/server/account.go
@@ -277,6 +277,7 @@ func (manager *AccountManager) AddPeer(setupKey string, peerKey string) (*Peer,
 	}
 
 	account.Peers[newPeer.Key] = newPeer
+	account.SetupKeys[sk.Key] = sk.IncrementUsage()
 	err = manager.Store.SaveAccount(account)
 	if err != nil {
 		return nil, status.Errorf(codes.Internal, "failed adding peer")
diff --git a/management/server/http/handler/setupkeys.go b/management/server/http/handler/setupkeys.go
index 0667cca1c..1e84efad5 100644
--- a/management/server/http/handler/setupkeys.go
+++ b/management/server/http/handler/setupkeys.go
@@ -18,13 +18,15 @@ type SetupKeys struct {
 
 // SetupKeyResponse is a response sent to the client
 type SetupKeyResponse struct {
-	Id      string
-	Key     string
-	Name    string
-	Expires time.Time
-	Type    server.SetupKeyType
-	Valid   bool
-	Revoked bool
+	Id        string
+	Key       string
+	Name      string
+	Expires   time.Time
+	Type      server.SetupKeyType
+	Valid     bool
+	Revoked   bool
+	UsedTimes int
+	LastUsed  time.Time
 }
 
 // SetupKeyRequest is a request sent by client. This object contains fields that can be modified
@@ -50,6 +52,11 @@ func (h *SetupKeys) CreateKey(w http.ResponseWriter, r *http.Request) {
 		return
 	}
 
+	if !(req.Type == server.SetupKeyReusable || req.Type == server.SetupKeyOneOff) {
+		http.Error(w, "unknown setup key type "+string(req.Type), http.StatusBadRequest)
+		return
+	}
+
 	setupKey, err := h.accountManager.AddSetupKey(accountId, req.Name, req.Type, req.ExpiresIn.Duration)
 	if err != nil {
 		errStatus, ok := status.FromError(err)
@@ -166,12 +173,14 @@ func writeSuccess(w http.ResponseWriter, key *server.SetupKey) {
 
 func toResponseBody(key *server.SetupKey) *SetupKeyResponse {
 	return &SetupKeyResponse{
-		Id:      key.Id,
-		Key:     key.Key,
-		Name:    key.Name,
-		Expires: key.ExpiresAt,
-		Type:    key.Type,
-		Valid:   key.IsValid(),
-		Revoked: key.Revoked,
+		Id:        key.Id,
+		Key:       key.Key,
+		Name:      key.Name,
+		Expires:   key.ExpiresAt,
+		Type:      key.Type,
+		Valid:     key.IsValid(),
+		Revoked:   key.Revoked,
+		UsedTimes: key.UsedTimes,
+		LastUsed:  key.LastUsed,
 	}
 }
diff --git a/management/server/setupkey.go b/management/server/setupkey.go
index 573999c2e..a17346a2d 100644
--- a/management/server/setupkey.go
+++ b/management/server/setupkey.go
@@ -35,6 +35,8 @@ type SetupKey struct {
 	Revoked bool
 	// UsedTimes indicates how many times the key was used
 	UsedTimes int
+	// LastUsed last time the key was used for peer registration
+	LastUsed time.Time
 }
 
 //Copy copies SetupKey to a new object
@@ -51,6 +53,14 @@ func (key *SetupKey) Copy() *SetupKey {
 	}
 }
 
+//IncrementUsage makes a copy of a key, increments the UsedTimes by 1 and sets LastUsed to now
+func (key *SetupKey) IncrementUsage() *SetupKey {
+	c := key.Copy()
+	c.UsedTimes = c.UsedTimes + 1
+	c.LastUsed = time.Now()
+	return c
+}
+
 // 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)