extend setup key logic (#91)

* feature: extend setup key logic

* fix: login test

* test: change test data to support new setup key fields

* test: add setup key tests

* test: add setup key UUID format validation

* test: add AddAccount test

* test: add more AccountManager tests

* fix: golint errors

* test: increase client interface up timout

client/cmd/login_test.go

@@ -7,6 +7,7 @@ import (
 	mgmt "github.com/wiretrustee/wiretrustee/management/server"
 	"github.com/wiretrustee/wiretrustee/util"
 	"path/filepath"
+	"strings"
 	"testing"
 )
 
@@ -38,7 +39,7 @@ func TestLogin(t *testing.T) {
 		"--config",
 		confPath,
 		"--setup-key",
-		"a2c8e62b-38f5-4553-b31e-dd66c696cebb",
+		strings.ToUpper("a2c8e62b-38f5-4553-b31e-dd66c696cebb"),
 		"--management-url",
 		mgmtURL,
 	})

client/cmd/up_test.go

@@ -9,7 +9,6 @@ import (
 	"net/url"
 	"os"
 	"path/filepath"
-	"strings"
 	"testing"
 	"time"
 )
@@ -71,7 +70,7 @@ func TestUp(t *testing.T) {
 		"--config",
 		confPath,
 		"--setup-key",
-		strings.ToUpper("a2c8e62b-38f5-4553-b31e-dd66c696cebb"),
+		"A2C8E62B-38F5-4553-B31E-DD66C696CEBB",
 		"--management-url",
 		mgmtURL.String(),
 	})
@@ -93,7 +92,7 @@ func TestUp(t *testing.T) {
 	}()
 
 	exists := false
-	for start := time.Now(); time.Since(start) < 7*time.Second; {
+	for start := time.Now(); time.Since(start) < 15*time.Second; {
 		e, err := iface.Exists(iface.WgInterfaceDefault)
 		if err != nil {
 			continue

client/testdata/store.json

@@ -3,8 +3,14 @@
         "bf1c8084-ba50-4ce7-9439-34653001fc3b": {
             "Id": "bf1c8084-ba50-4ce7-9439-34653001fc3b",
             "SetupKeys": {
-                "a2c8e62b-38f5-4553-b31e-dd66c696cebb": {
+                "A2C8E62B-38F5-4553-B31E-DD66C696CEBB": {
-                    "Key": "a2c8e62b-38f5-4553-b31e-dd66c696cebb"
+                    "Key": "A2C8E62B-38F5-4553-B31E-DD66C696CEBB",
+                    "Name": "Default key",
+                    "Type": "reusable",
+                    "CreatedAt": "2021-08-19T20:46:20.005936822+02:00",
+                    "ExpiresAt": "2321-09-18T20:46:20.005936822+02:00",
+                    "Revoked": false,
+                    "UsedTimes": 0
                 }
             },
             "Network": {

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
 }
@@ -90,7 +84,7 @@ func (manager *AccountManager) GetAccount(accountId string) (*Account, error) {
 
 	account, err := manager.Store.GetAccount(accountId)
 	if err != nil {
-		return nil, status.Errorf(codes.Internal, "failed retrieving account")
+		return nil, status.Errorf(codes.NotFound, "account not found")
 	}
 
 	return account, nil
@@ -164,23 +158,40 @@ 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
+				break
+			}
+		}
+
+		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 +204,7 @@ func (manager *AccountManager) AddPeer(setupKey string, peerKey string) (*Peer,
 
 	newPeer := &Peer{
 		Key:      peerKey,
-		SetupKey: sk,
+		SetupKey: sk.Key,
 		IP:       nextIp,
 	}
 
@@ -212,17 +223,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 := GenerateDefaultSetupKey()
-	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/account_test.go

@@ -0,0 +1,211 @@
+package server
+
+import (
+	"golang.zx2c4.com/wireguard/wgctrl/wgtypes"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/status"
+	"net"
+	"testing"
+)
+
+func TestAccountManager_AddAccount(t *testing.T) {
+	manager, err := createManager(t)
+	if err != nil {
+		t.Fatal(err)
+		return
+	}
+
+	expectedId := "test_account"
+	expectedPeersSize := 0
+	expectedSetupKeysSize := 1
+	expectedNetwork := net.IPNet{
+		IP:   net.IP{100, 64, 0, 0},
+		Mask: net.IPMask{255, 192, 0, 0},
+	}
+
+	account, err := manager.AddAccount(expectedId)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	if account.Id != expectedId {
+		t.Errorf("expected account to have ID = %s, got %s", expectedId, account.Id)
+	}
+
+	if len(account.Peers) != expectedPeersSize {
+		t.Errorf("expected account to have len(Peers) = %v, got %v", expectedPeersSize, len(account.Peers))
+	}
+
+	if len(account.SetupKeys) != expectedSetupKeysSize {
+		t.Errorf("expected account to have len(SetupKeys) = %v, got %v", expectedSetupKeysSize, len(account.SetupKeys))
+	}
+
+	if account.Network.Net.String() != expectedNetwork.String() {
+		t.Errorf("expected account to have Network = %v, got %v", expectedNetwork.String(), account.Network.Net.String())
+	}
+}
+
+func TestAccountManager_GetOrCreateAccount(t *testing.T) {
+	manager, err := createManager(t)
+	if err != nil {
+		t.Fatal(err)
+		return
+	}
+
+	expectedId := "test_account"
+
+	//make sure account doesn't exist
+	account, err := manager.GetAccount(expectedId)
+	if err != nil {
+		errStatus, ok := status.FromError(err)
+		if !(ok && errStatus.Code() == codes.NotFound) {
+			t.Fatal(err)
+		}
+	}
+	if account != nil {
+		t.Fatal("expecting empty account")
+	}
+
+	account, err = manager.GetOrCreateAccount(expectedId)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	if account.Id != expectedId {
+		t.Fatalf("expected to create an account, got wrong account")
+	}
+
+	account, err = manager.GetOrCreateAccount(expectedId)
+	if err != nil {
+		t.Errorf("expected to get existing account after creation, failed")
+	}
+
+	if account.Id != expectedId {
+		t.Fatalf("expected to create an account, got wrong account")
+	}
+}
+
+func TestAccountManager_AccountExists(t *testing.T) {
+	manager, err := createManager(t)
+	if err != nil {
+		t.Fatal(err)
+		return
+	}
+
+	expectedId := "test_account"
+	_, err = manager.AddAccount(expectedId)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	exists, err := manager.AccountExists(expectedId)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	if !*exists {
+		t.Errorf("expected account to exist after creation, got false")
+	}
+
+}
+
+func TestAccountManager_GetAccount(t *testing.T) {
+	manager, err := createManager(t)
+	if err != nil {
+		t.Fatal(err)
+		return
+	}
+
+	expectedId := "test_account"
+	account, err := manager.AddAccount(expectedId)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	//AddAccount has been already tested so we can assume it is correct and compare results
+	getAccount, err := manager.GetAccount(expectedId)
+	if err != nil {
+		t.Fatal(err)
+		return
+	}
+
+	if account.Id != getAccount.Id {
+		t.Errorf("expected account.ID %s, got %s", account.Id, getAccount.Id)
+	}
+
+	for _, peer := range account.Peers {
+		if _, ok := getAccount.Peers[peer.Key]; !ok {
+			t.Errorf("expected account to have peer %s, not found", peer.Key)
+		}
+	}
+
+	for _, key := range account.SetupKeys {
+		if _, ok := getAccount.SetupKeys[key.Key]; !ok {
+			t.Errorf("expected account to have setup key %s, not found", key.Key)
+		}
+	}
+
+}
+
+func TestAccountManager_AddPeer(t *testing.T) {
+	manager, err := createManager(t)
+	if err != nil {
+		t.Fatal(err)
+		return
+	}
+
+	account, err := manager.AddAccount("test_account")
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	var setupKey *SetupKey
+	for _, key := range account.SetupKeys {
+		setupKey = key
+	}
+
+	if setupKey == nil {
+		t.Errorf("expecting account to have a default setup key")
+		return
+	}
+
+	key, err := wgtypes.GenerateKey()
+	if err != nil {
+		t.Fatal(err)
+		return
+	}
+	expectedPeerKey := key.PublicKey().String()
+	expectedPeerIP := "100.64.0.1"
+
+	peer, err := manager.AddPeer(setupKey.Key, expectedPeerKey)
+	if err != nil {
+		t.Errorf("expecting peer to be added, got failure %v", err)
+		return
+	}
+
+	if peer.Key != expectedPeerKey {
+		t.Errorf("expecting just added peer to have key = %s, got %s", expectedPeerKey, peer.Key)
+	}
+
+	if peer.Key != expectedPeerKey {
+		t.Errorf("expecting just added peer to have IP = %s, got %s", expectedPeerIP, peer.IP.String())
+	}
+
+}
+func createManager(t *testing.T) (*AccountManager, error) {
+	store, err := createStore(t)
+	if err != nil {
+		return nil, err
+	}
+	return NewManager(store), nil
+}
+
+func createStore(t *testing.T) (Store, error) {
+	dataDir := t.TempDir()
+	store, err := NewStore(dataDir)
+	if err != nil {
+		return nil, err
+	}
+
+	return store, nil
+}

management/server/file_store.go

@@ -116,7 +116,7 @@ func (s *FileStore) SaveAccount(account *Account) error {
 	// todo check that account.Id and keyId are not exist already
 	// because if keyId exists for other accounts this can be bad
 	for keyId := range account.SetupKeys {
-		s.SetupKeyId2AccountId[strings.ToLower(keyId)] = account.Id
+		s.SetupKeyId2AccountId[strings.ToUpper(keyId)] = account.Id
 	}
 
 	for _, peer := range account.Peers {

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/management_test.go

@@ -26,7 +26,6 @@ import (
 
 const (
 	ValidSetupKey = "A2C8E62B-38F5-4553-B31E-DD66C696CEBB"
-	InvalidSetupKey = "INVALID_SETUP_KEY"
 )
 
 var _ = Describe("Management service", func() {

management/server/setupkey.go

@@ -0,0 +1,61 @@
+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,
+	}
+}
+
+// GenerateDefaultSetupKey generates a default setup key
+func GenerateDefaultSetupKey() *SetupKey {
+	return GenerateSetupKey(DefaultSetupKeyName, SetupKeyReusable, DefaultSetupKeyDuration)
+}

management/server/setupkey_test.go

@@ -0,0 +1,100 @@
+package server
+
+import (
+	"github.com/google/uuid"
+	"testing"
+	"time"
+)
+
+func TestGenerateDefaultSetupKey(t *testing.T) {
+	expectedName := "Default key"
+	expectedRevoke := false
+	expectedType := "reusable"
+	expectedUsedTimes := 0
+	expectedCreatedAt := time.Now()
+	expectedExpiresAt := time.Now().Add(24 * 30 * time.Hour)
+
+	key := GenerateDefaultSetupKey()
+
+	assertKey(t, key, expectedName, expectedRevoke, expectedType, expectedUsedTimes, expectedCreatedAt, expectedExpiresAt)
+
+}
+
+func TestGenerateSetupKey(t *testing.T) {
+	expectedName := "key"
+	expectedRevoke := false
+	expectedType := "one-off"
+	expectedUsedTimes := 0
+	expectedCreatedAt := time.Now()
+	expectedExpiresAt := time.Now().Add(time.Hour)
+
+	key := GenerateSetupKey(expectedName, SetupKeyOneOff, time.Hour)
+
+	assertKey(t, key, expectedName, expectedRevoke, expectedType, expectedUsedTimes, expectedCreatedAt, expectedExpiresAt)
+
+}
+
+func TestSetupKey_IsValid(t *testing.T) {
+	validKey := GenerateSetupKey("valid key", SetupKeyOneOff, time.Hour)
+	if !validKey.IsValid() {
+		t.Errorf("expected key to be valid, got invalid %v", validKey)
+	}
+
+	// expired
+	expiredKey := GenerateSetupKey("invalid key", SetupKeyOneOff, -time.Hour)
+	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)
+	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)
+	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)
+	reusableKey.UsedTimes = 99
+	if !reusableKey.IsValid() {
+		t.Errorf("expected reusable key to be valid when used many times, got valid %v", reusableKey)
+	}
+}
+
+func assertKey(t *testing.T, key *SetupKey, expectedName string, expectedRevoke bool, expectedType string, expectedUsedTimes int, expectedCreatedAt time.Time, expectedExpiresAt time.Time) {
+	if key.Name != expectedName {
+		t.Errorf("expected setup key to have Name %v, got %v", expectedName, key.Name)
+	}
+
+	if key.Revoked != expectedRevoke {
+		t.Errorf("expected setup key to have Revoke %v, got %v", expectedRevoke, key.Revoked)
+	}
+
+	if string(key.Type) != expectedType {
+		t.Errorf("expected setup key to have Type %v, got %v", expectedType, key.Type)
+	}
+
+	if key.UsedTimes != expectedUsedTimes {
+		t.Errorf("expected setup key to have UsedTimes = %v, got %v", expectedUsedTimes, key.UsedTimes)
+	}
+
+	if key.ExpiresAt.Sub(expectedExpiresAt).Round(time.Hour) != 0 {
+		t.Errorf("expected setup key to have ExpiresAt ~ %v, got %v", expectedExpiresAt, key.ExpiresAt)
+	}
+
+	if key.CreatedAt.Sub(expectedCreatedAt).Round(time.Hour) != 0 {
+		t.Errorf("expected setup key to have CreatedAt ~ %v, got %v", expectedCreatedAt, key.CreatedAt)
+	}
+
+	_, err := uuid.Parse(key.Key)
+	if err != nil {
+		t.Errorf("expected key to be a valid UUID, got %v, %v", key.Key, err)
+	}
+}

management/server/testdata/store.json

@@ -3,8 +3,15 @@
         "bf1c8084-ba50-4ce7-9439-34653001fc3b": {
             "Id": "bf1c8084-ba50-4ce7-9439-34653001fc3b",
             "SetupKeys": {
-                "a2c8e62b-38f5-4553-b31e-dd66c696cebb": {
+                "A2C8E62B-38F5-4553-B31E-DD66C696CEBB": {
-                    "Key": "a2c8e62b-38f5-4553-b31e-dd66c696cebb"
+                    "Key": "A2C8E62B-38F5-4553-B31E-DD66C696CEBB",
+                    "Name": "Default key",
+                    "Type": "reusable",
+                    "CreatedAt": "2021-08-19T20:46:20.005936822+02:00",
+                    "ExpiresAt": "2321-09-18T20:46:20.005936822+02:00",
+                    "Revoked": false,
+                    "UsedTimes": 0
+
                 }
             },
             "Network": {