Change Management Sync protocol to support incremental (serial) network changes (#191)

* feature: introduce NetworkMap to the management protocol with a Serial ID

* test: add Management Sync method protocol test

* test: add Management Sync method NetworkMap field check [FAILING]

* test: add Management Sync method NetworkMap field check [FAILING]

* feature: fill NetworkMap property to when Deleting peer

* feature: fill NetworkMap in the Sync protocol

* test: code review mentions - GeneratePrivateKey() in the test

* fix: wiretrustee client use wireguard GeneratePrivateKey() instead of GenerateKey()

* test: add NetworkMap test

* fix: management_proto test remove store.json on test finish

management/server/account_test.go

@@ -163,7 +163,7 @@ func TestAccountManager_AddPeer(t *testing.T) {
 		return
 	}
 
-	key, err := wgtypes.GenerateKey()
+	key, err := wgtypes.GeneratePrivateKey()
 	if err != nil {
 		t.Fatal(err)
 		return
@@ -171,7 +171,7 @@ func TestAccountManager_AddPeer(t *testing.T) {
 	expectedPeerKey := key.PublicKey().String()
 	expectedPeerIP := "100.64.0.1"
 
-	peer, err := manager.AddPeer(setupKey.Key, Peer{
+	peer, err := manager.AddPeer(setupKey.Key, &Peer{
 		Key:  expectedPeerKey,
 		Meta: PeerSystemMeta{},
 		Name: expectedPeerKey,
@@ -226,7 +226,7 @@ func TestAccountManager_DeletePeer(t *testing.T) {
 
 	peerKey := key.PublicKey().String()
 
-	_, err = manager.AddPeer(setupKey.Key, Peer{
+	_, err = manager.AddPeer(setupKey.Key, &Peer{
 		Key:  peerKey,
 		Meta: PeerSystemMeta{},
 		Name: peerKey,

management/server/grpcserver.go

@@ -141,7 +141,7 @@ func (s *Server) registerPeer(peerKey wgtypes.Key, req *proto.LoginRequest) (*Pe
 	if meta == nil {
 		return nil, status.Errorf(codes.InvalidArgument, "peer meta data was not provided")
 	}
-	peer, err := s.accountManager.AddPeer(req.GetSetupKey(), Peer{
+	peer, err := s.accountManager.AddPeer(req.GetSetupKey(), &Peer{
 		Key:  peerKey.String(),
 		Name: meta.GetHostname(),
 		Meta: PeerSystemMeta{
@@ -158,21 +158,22 @@ func (s *Server) registerPeer(peerKey wgtypes.Key, req *proto.LoginRequest) (*Pe
 		return nil, status.Errorf(codes.NotFound, "provided setup key doesn't exists")
 	}
 
-	peers, err := s.accountManager.GetPeersForAPeer(peer.Key)
+	//todo move to AccountManager the code below
+	networkMap, err := s.accountManager.GetNetworkMap(peer.Key)
 	if err != nil {
 		return nil, status.Error(codes.Internal, "internal server error")
 	}
 
 	// notify other peers of our registration
-	for _, remotePeer := range peers {
+	for _, remotePeer := range networkMap.Peers {
 		// exclude notified peer and add ourselves
 		peersToSend := []*Peer{peer}
-		for _, p := range peers {
+		for _, p := range networkMap.Peers {
 			if remotePeer.Key != p.Key {
 				peersToSend = append(peersToSend, p)
 			}
 		}
-		update := toSyncResponse(s.config, peer, peersToSend, nil)
+		update := toSyncResponse(s.config, peer, peersToSend, nil, networkMap.Network.Serial())
 		err = s.peersUpdateManager.SendUpdate(remotePeer.Key, &UpdateMessage{Update: update})
 		if err != nil {
 			// todo rethink if we should keep this return
@@ -317,7 +318,7 @@ func toRemotePeerConfig(peers []*Peer) []*proto.RemotePeerConfig {
 
 }
 
-func toSyncResponse(config *Config, peer *Peer, peers []*Peer, turnCredentials *TURNCredentials) *proto.SyncResponse {
+func toSyncResponse(config *Config, peer *Peer, peers []*Peer, turnCredentials *TURNCredentials, serial uint64) *proto.SyncResponse {
 
 	wtConfig := toWiretrusteeConfig(config, turnCredentials)
 
@@ -330,6 +331,12 @@ func toSyncResponse(config *Config, peer *Peer, peers []*Peer, turnCredentials *
 		PeerConfig:         pConfig,
 		RemotePeers:        remotePeers,
 		RemotePeersIsEmpty: len(remotePeers) == 0,
+		NetworkMap: &proto.NetworkMap{
+			Serial:             serial,
+			PeerConfig:         pConfig,
+			RemotePeers:        remotePeers,
+			RemotePeersIsEmpty: len(remotePeers) == 0,
+		},
 	}
 }
 
@@ -341,7 +348,7 @@ func (s *Server) IsHealthy(ctx context.Context, req *proto.Empty) (*proto.Empty,
 // sendInitialSync sends initial proto.SyncResponse to the peer requesting synchronization
 func (s *Server) sendInitialSync(peerKey wgtypes.Key, peer *Peer, srv proto.ManagementService_SyncServer) error {
 
-	peers, err := s.accountManager.GetPeersForAPeer(peer.Key)
+	networkMap, err := s.accountManager.GetNetworkMap(peer.Key)
 	if err != nil {
 		log.Warnf("error getting a list of peers for a peer %s", peer.Key)
 		return err
@@ -355,7 +362,7 @@ func (s *Server) sendInitialSync(peerKey wgtypes.Key, peer *Peer, srv proto.Mana
 	} else {
 		turnCredentials = nil
 	}
-	plainResp := toSyncResponse(s.config, peer, peers, turnCredentials)
+	plainResp := toSyncResponse(s.config, peer, networkMap.Peers, turnCredentials, networkMap.Network.Serial())
 
 	encryptedResp, err := encryption.EncryptMessage(peerKey, s.wgKey, plainResp)
 	if err != nil {

management/server/management_proto_test.go

@@ -0,0 +1,348 @@
+package server
+
+import (
+	"context"
+	"fmt"
+	log "github.com/sirupsen/logrus"
+	"github.com/wiretrustee/wiretrustee/encryption"
+	mgmtProto "github.com/wiretrustee/wiretrustee/management/proto"
+	"github.com/wiretrustee/wiretrustee/util"
+	"golang.zx2c4.com/wireguard/wgctrl/wgtypes"
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/credentials/insecure"
+	"google.golang.org/grpc/keepalive"
+	"net"
+	"os"
+	"path/filepath"
+	"runtime"
+	"testing"
+	"time"
+)
+
+var (
+	kaep = keepalive.EnforcementPolicy{
+		MinTime:             15 * time.Second,
+		PermitWithoutStream: true,
+	}
+
+	kasp = keepalive.ServerParameters{
+		MaxConnectionIdle:     15 * time.Second,
+		MaxConnectionAgeGrace: 5 * time.Second,
+		Time:                  5 * time.Second,
+		Timeout:               2 * time.Second,
+	}
+)
+
+const (
+	TestValidSetupKey = "A2C8E62B-38F5-4553-B31E-DD66C696CEBB"
+)
+
+// registerPeers registers peersNum peers on the management service and returns their Wireguard keys
+func registerPeers(peersNum int, client mgmtProto.ManagementServiceClient) ([]*wgtypes.Key, error) {
+
+	var peers = []*wgtypes.Key{}
+	for i := 0; i < peersNum; i++ {
+		key, err := wgtypes.GeneratePrivateKey()
+		if err != nil {
+			return nil, err
+		}
+
+		_, err = loginPeerWithValidSetupKey(key, client)
+		if err != nil {
+			return nil, err
+		}
+
+		peers = append(peers, &key)
+	}
+
+	return peers, nil
+}
+
+// getServerKey gets Management Service Wireguard public key
+func getServerKey(client mgmtProto.ManagementServiceClient) (*wgtypes.Key, error) {
+
+	keyResp, err := client.GetServerKey(context.TODO(), &mgmtProto.Empty{})
+	if err != nil {
+		return nil, err
+	}
+
+	serverKey, err := wgtypes.ParseKey(keyResp.Key)
+	if err != nil {
+		return nil, err
+	}
+
+	return &serverKey, nil
+}
+
+func Test_SyncProtocol(t *testing.T) {
+
+	dir := t.TempDir()
+	err := util.CopyFileContents("testdata/store.json", filepath.Join(dir, "store.json"))
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer func() {
+		os.Remove(filepath.Join(dir, "store.json")) //nolint
+	}()
+
+	mgmtServer, err := startManagement(33071, &Config{
+		Stuns: []*Host{{
+			Proto: "udp",
+			URI:   "stun:stun.wiretrustee.com:3468",
+		}},
+		TURNConfig: &TURNConfig{
+			TimeBasedCredentials: false,
+			CredentialsTTL:       util.Duration{},
+			Secret:               "whatever",
+			Turns: []*Host{{
+				Proto: "udp",
+				URI:   "turn:stun.wiretrustee.com:3468",
+			}},
+		},
+		Signal: &Host{
+			Proto: "http",
+			URI:   "signal.wiretrustee.com:10000",
+		},
+		Datadir:    dir,
+		HttpConfig: nil,
+	})
+	if err != nil {
+		t.Fatal(err)
+		return
+	}
+	defer mgmtServer.Stop()
+
+	client, clientConn, err := createRawClient("localhost:33071")
+	if err != nil {
+		t.Fatal(err)
+		return
+	}
+
+	defer clientConn.Close()
+
+	peers, err := registerPeers(2, client)
+	if err != nil {
+		t.Fatal(err)
+		return
+	}
+
+	serverKey, err := getServerKey(client)
+	if err != nil {
+		t.Fatal(err)
+		return
+	}
+
+	// take the first registered peer as a base for the test
+	key := *peers[0]
+
+	message, err := encryption.EncryptMessage(*serverKey, key, &mgmtProto.SyncRequest{})
+	if err != nil {
+		t.Fatal(err)
+		return
+	}
+
+	if err != nil {
+		t.Fatal(err)
+		return
+	}
+
+	sync, err := client.Sync(context.TODO(), &mgmtProto.EncryptedMessage{
+		WgPubKey: key.PublicKey().String(),
+		Body:     message,
+	})
+	if err != nil {
+		t.Fatal(err)
+		return
+	}
+
+	resp := &mgmtProto.EncryptedMessage{}
+	err = sync.RecvMsg(resp)
+	if err != nil {
+		t.Fatal(err)
+		return
+	}
+
+	syncResp := &mgmtProto.SyncResponse{}
+	err = encryption.DecryptMessage(*serverKey, key, resp.Body, syncResp)
+	if err != nil {
+		t.Fatal(err)
+		return
+	}
+
+	wiretrusteeConfig := syncResp.GetWiretrusteeConfig()
+	if wiretrusteeConfig == nil {
+		t.Fatal("expecting SyncResponse to have non-nil WiretrusteeConfig")
+	}
+
+	if wiretrusteeConfig.GetSignal() == nil {
+		t.Fatal("expecting SyncResponse to have WiretrusteeConfig with non-nil Signal config")
+	}
+
+	expectedSignalConfig := &mgmtProto.HostConfig{
+		Uri:      "signal.wiretrustee.com:10000",
+		Protocol: mgmtProto.HostConfig_HTTP,
+	}
+
+	if wiretrusteeConfig.GetSignal().GetUri() != expectedSignalConfig.GetUri() {
+		t.Fatalf("expecting SyncResponse to have WiretrusteeConfig with expected Signal URI: %v, actual: %v",
+			expectedSignalConfig.GetUri(),
+			wiretrusteeConfig.GetSignal().GetUri())
+	}
+
+	if wiretrusteeConfig.GetSignal().GetProtocol() != expectedSignalConfig.GetProtocol() {
+		t.Fatalf("expecting SyncResponse to have WiretrusteeConfig with expected Signal Protocol: %v, actual: %v",
+			expectedSignalConfig.GetProtocol().String(),
+			wiretrusteeConfig.GetSignal().GetProtocol())
+	}
+
+	expectedStunsConfig := &mgmtProto.HostConfig{
+		Uri:      "stun:stun.wiretrustee.com:3468",
+		Protocol: mgmtProto.HostConfig_UDP,
+	}
+
+	if wiretrusteeConfig.GetStuns()[0].GetUri() != expectedStunsConfig.GetUri() {
+		t.Fatalf("expecting SyncResponse to have WiretrusteeConfig with expected STUN URI: %v, actual: %v",
+			expectedStunsConfig.GetUri(),
+			wiretrusteeConfig.GetStuns()[0].GetUri())
+	}
+
+	if wiretrusteeConfig.GetStuns()[0].GetProtocol() != expectedStunsConfig.GetProtocol() {
+		t.Fatalf("expecting SyncResponse to have WiretrusteeConfig with expected STUN Protocol: %v, actual: %v",
+			expectedStunsConfig.GetProtocol(),
+			wiretrusteeConfig.GetStuns()[0].GetProtocol())
+	}
+
+	expectedTRUNHost := &mgmtProto.HostConfig{
+		Uri:      "turn:stun.wiretrustee.com:3468",
+		Protocol: mgmtProto.HostConfig_UDP,
+	}
+
+	if wiretrusteeConfig.GetTurns()[0].GetHostConfig().GetUri() != expectedTRUNHost.GetUri() {
+		t.Fatalf("expecting SyncResponse to have WiretrusteeConfig with expected TURN URI: %v, actual: %v",
+			expectedTRUNHost.GetUri(),
+			wiretrusteeConfig.GetTurns()[0].GetHostConfig().GetUri())
+	}
+
+	if wiretrusteeConfig.GetTurns()[0].GetHostConfig().GetProtocol() != expectedTRUNHost.GetProtocol() {
+		t.Fatalf("expecting SyncResponse to have WiretrusteeConfig with expected TURN Protocol: %v, actual: %v",
+			expectedTRUNHost.GetProtocol().String(),
+			wiretrusteeConfig.GetTurns()[0].GetHostConfig().GetProtocol())
+	}
+
+	// ensure backward compatibility
+
+	if syncResp.GetRemotePeers() == nil {
+		t.Fatal("expecting SyncResponse to have non-nil RemotePeers for backward compatibility")
+	}
+
+	if syncResp.GetPeerConfig() == nil {
+		t.Fatal("expecting SyncResponse to have non-nil PeerConfig for backward compatibility")
+	}
+
+	// new field - NetworkMap
+	networkMap := syncResp.GetNetworkMap()
+	if networkMap == nil {
+		t.Fatal("expecting SyncResponse to have non-nil NetworkMap")
+	}
+
+	if len(networkMap.GetRemotePeers()) != 1 {
+		t.Fatal("expecting SyncResponse to have NetworkMap with 1 remote peer")
+	}
+
+	if networkMap.GetPeerConfig() == nil {
+		t.Fatal("expecting SyncResponse to have NetworkMap with a non-nil PeerConfig")
+	}
+
+	if networkMap.GetPeerConfig().GetAddress() != "100.64.0.1/24" {
+		t.Fatal("expecting SyncResponse to have NetworkMap with a PeerConfig having valid Address")
+	}
+
+	if networkMap.GetSerial() <= 0 {
+		t.Fatalf("expecting SyncResponse to have NetworkMap with a positive Network Serial, actual %d", networkMap.GetSerial())
+	}
+}
+
+func loginPeerWithValidSetupKey(key wgtypes.Key, client mgmtProto.ManagementServiceClient) (*mgmtProto.LoginResponse, error) {
+
+	serverKey, err := getServerKey(client)
+	if err != nil {
+		return nil, err
+	}
+
+	meta := &mgmtProto.PeerSystemMeta{
+		Hostname:           key.PublicKey().String(),
+		GoOS:               runtime.GOOS,
+		OS:                 runtime.GOOS,
+		Core:               "core",
+		Platform:           "platform",
+		Kernel:             "kernel",
+		WiretrusteeVersion: "",
+	}
+	message, err := encryption.EncryptMessage(*serverKey, key, &mgmtProto.LoginRequest{SetupKey: TestValidSetupKey, Meta: meta})
+	if err != nil {
+		return nil, err
+	}
+
+	resp, err := client.Login(context.TODO(), &mgmtProto.EncryptedMessage{
+		WgPubKey: key.PublicKey().String(),
+		Body:     message,
+	})
+	if err != nil {
+		return nil, err
+	}
+
+	loginResp := &mgmtProto.LoginResponse{}
+	err = encryption.DecryptMessage(*serverKey, key, resp.Body, loginResp)
+	if err != nil {
+		return nil, err
+	}
+
+	return loginResp, nil
+
+}
+
+func startManagement(port int, config *Config) (*grpc.Server, error) {
+
+	lis, err := net.Listen("tcp", fmt.Sprintf("localhost:%d", port))
+	if err != nil {
+		return nil, err
+	}
+	s := grpc.NewServer(grpc.KeepaliveEnforcementPolicy(kaep), grpc.KeepaliveParams(kasp))
+	store, err := NewStore(config.Datadir)
+	if err != nil {
+		return nil, err
+	}
+	peersUpdateManager := NewPeersUpdateManager()
+	accountManager := NewManager(store, peersUpdateManager)
+	turnManager := NewTimeBasedAuthSecretsManager(peersUpdateManager, config.TURNConfig)
+	mgmtServer, err := NewServer(config, accountManager, peersUpdateManager, turnManager)
+	if err != nil {
+		return nil, err
+	}
+	mgmtProto.RegisterManagementServiceServer(s, mgmtServer)
+	go func() {
+		if err = s.Serve(lis); err != nil {
+			log.Fatalf("failed to serve: %v", err)
+		}
+	}()
+
+	return s, nil
+}
+
+func createRawClient(addr string) (mgmtProto.ManagementServiceClient, *grpc.ClientConn, error) {
+	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+	defer cancel()
+
+	conn, err := grpc.DialContext(ctx, addr,
+		grpc.WithTransportCredentials(insecure.NewCredentials()),
+		grpc.WithBlock(),
+		grpc.WithKeepaliveParams(keepalive.ClientParameters{
+			Time:    10 * time.Second,
+			Timeout: 2 * time.Second,
+		}))
+	if err != nil {
+		return nil, nil, err
+	}
+
+	return mgmtProto.NewManagementServiceClient(conn), conn, nil
+}

management/server/network.go

@@ -13,6 +13,11 @@ var (
 	upperIPv6 = []byte{0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff}
 )
 
+type NetworkMap struct {
+	Peers   []*Peer
+	Network *Network
+}
+
 type Network struct {
 	Id  string
 	Net net.IPNet

management/server/peer.go

@@ -137,8 +137,15 @@ func (am *AccountManager) DeletePeer(accountId string, peerKey string) (*Peer, e
 	err = am.peersUpdateManager.SendUpdate(peerKey,
 		&UpdateMessage{
 			Update: &proto.SyncResponse{
+				// fill those field for backward compatibility
 				RemotePeers:        []*proto.RemotePeerConfig{},
 				RemotePeersIsEmpty: true,
+				// new field
+				NetworkMap: &proto.NetworkMap{
+					Serial:             account.Network.Serial(),
+					RemotePeers:        []*proto.RemotePeerConfig{},
+					RemotePeersIsEmpty: true,
+				},
 			}})
 	if err != nil {
 		return nil, err
@@ -161,8 +168,15 @@ func (am *AccountManager) DeletePeer(accountId string, peerKey string) (*Peer, e
 		err = am.peersUpdateManager.SendUpdate(p.Key,
 			&UpdateMessage{
 				Update: &proto.SyncResponse{
+					// fill those field for backward compatibility
 					RemotePeers:        update,
 					RemotePeersIsEmpty: len(update) == 0,
+					// new field
+					NetworkMap: &proto.NetworkMap{
+						Serial:             account.Network.Serial(),
+						RemotePeers:        update,
+						RemotePeersIsEmpty: len(update) == 0,
+					},
 				}})
 		if err != nil {
 			return nil, err
@@ -192,9 +206,8 @@ func (am *AccountManager) GetPeerByIP(accountId string, peerIP string) (*Peer, e
 	return nil, status.Errorf(codes.NotFound, "peer with IP %s not found", peerIP)
 }
 
-// GetPeersForAPeer returns a list of peers available for a given peer (key)
-// Effectively all the peers of the original peer's account except for the peer itself
-func (am *AccountManager) GetPeersForAPeer(peerKey string) ([]*Peer, error) {
+// GetNetworkMap returns Network map for a given peer (omits original peer from the Peers result)
+func (am *AccountManager) GetNetworkMap(peerKey string) (*NetworkMap, error) {
 	am.mux.Lock()
 	defer am.mux.Unlock()
 
@@ -205,12 +218,16 @@ func (am *AccountManager) GetPeersForAPeer(peerKey string) ([]*Peer, error) {
 
 	var res []*Peer
 	for _, peer := range account.Peers {
+		// exclude original peer
 		if peer.Key != peerKey {
-			res = append(res, peer)
+			res = append(res, peer.Copy())
 		}
 	}
 
-	return res, nil
+	return &NetworkMap{
+		Peers:   res,
+		Network: account.Network.Copy(),
+	}, err
 }
 
 // AddPeer adds a new peer to the Store.
@@ -218,7 +235,7 @@ func (am *AccountManager) GetPeersForAPeer(peerKey string) ([]*Peer, error) {
 // 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).
 // The peer property is just a placeholder for the Peer properties to pass further
-func (am *AccountManager) AddPeer(setupKey string, peer Peer) (*Peer, error) {
+func (am *AccountManager) AddPeer(setupKey string, peer *Peer) (*Peer, error) {
 	am.mux.Lock()
 	defer am.mux.Unlock()
 

management/server/peer_test.go

@@ -0,0 +1,76 @@
+package server
+
+import (
+	"golang.zx2c4.com/wireguard/wgctrl/wgtypes"
+	"testing"
+)
+
+func TestAccountManager_GetNetworkMap(t *testing.T) {
+	manager, err := createManager(t)
+	if err != nil {
+		t.Fatal(err)
+		return
+	}
+
+	expectedId := "test_account"
+	userId := "account_creator"
+	account, err := manager.AddAccount(expectedId, userId)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	var setupKey *SetupKey
+	for _, key := range account.SetupKeys {
+		if key.Type == SetupKeyReusable {
+			setupKey = key
+		}
+	}
+
+	peerKey1, err := wgtypes.GeneratePrivateKey()
+	if err != nil {
+		t.Fatal(err)
+		return
+	}
+
+	_, err = manager.AddPeer(setupKey.Key, &Peer{
+		Key:  peerKey1.PublicKey().String(),
+		Meta: PeerSystemMeta{},
+		Name: "test-peer-2",
+	})
+
+	if err != nil {
+		t.Errorf("expecting peer to be added, got failure %v", err)
+		return
+	}
+
+	peerKey2, err := wgtypes.GeneratePrivateKey()
+	if err != nil {
+		t.Fatal(err)
+		return
+	}
+	_, err = manager.AddPeer(setupKey.Key, &Peer{
+		Key:  peerKey2.PublicKey().String(),
+		Meta: PeerSystemMeta{},
+		Name: "test-peer-2",
+	})
+
+	if err != nil {
+		t.Errorf("expecting peer to be added, got failure %v", err)
+		return
+	}
+
+	networkMap, err := manager.GetNetworkMap(peerKey1.PublicKey().String())
+	if err != nil {
+		t.Fatal(err)
+		return
+	}
+
+	if len(networkMap.Peers) != 1 {
+		t.Errorf("expecting Account NetworkMap to have 1 peers, got %v", len(networkMap.Peers))
+	}
+
+	if networkMap.Peers[0].Key != peerKey2.PublicKey().String() {
+		t.Errorf("expecting Account NetworkMap to have peer with a key %s, got %s", peerKey2.PublicKey().String(), networkMap.Peers[0].Key)
+	}
+
+}