Add client's interaction with management service (#71)

* Add client's interaction with management service

* Getting updates

* Fixed key and nil ptr

* Added setupKey param

* Added managment address parameter

* Fixed test

* feature: use RemotePeers from the management server instead of deprecated Peers

* merge: merge changes from main

cmd/addpeer.go

@@ -1,10 +1,13 @@
 package cmd
 
 import (
+	"os"
+
 	log "github.com/sirupsen/logrus"
 	"github.com/spf13/cobra"
 	"github.com/wiretrustee/wiretrustee/connection"
-	"os"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/status"
 )
 
 var (
@@ -16,24 +19,9 @@ var (
 		Short: "add remote peer",
 		Run: func(cmd *cobra.Command, args []string) {
 			InitLog(logLevel)
-
+			err := addPeer(key, allowedIPs)
-			if _, err := os.Stat(configPath); os.IsNotExist(err) {
-				log.Error("config doesn't exist, please run 'wiretrustee init' first")
-				os.Exit(ExitSetupFailed)
-			}
-
-			config, err := Read(configPath)
 			if err != nil {
-				log.Fatalf("Error reading config file, message: %v", err)
+				log.Errorf("Failed to add peer: %s", err)
-			}
-			config.Peers = append(config.Peers, connection.Peer{
-				WgPubKey:     key,
-				WgAllowedIps: allowedIPs,
-			})
-
-			err = config.Write(configPath)
-			if err != nil {
-				log.Errorf("failed writing config to %s: %s", config, err.Error())
 				os.Exit(ExitSetupFailed)
 			}
 		},
@@ -46,3 +34,26 @@ func init() {
 	addPeerCmd.MarkPersistentFlagRequired("key")        //nolint
 	addPeerCmd.MarkPersistentFlagRequired("allowedIPs") //nolint
 }
+
+func addPeer(peerKey string, allowedIPs string) error {
+	if _, err := os.Stat(configPath); os.IsNotExist(err) {
+		return status.Errorf(codes.FailedPrecondition, "Config doesn't exist, please run 'wiretrustee init' first")
+	}
+
+	config, err := Read(configPath)
+	if err != nil {
+		return status.Errorf(codes.FailedPrecondition, "Error reading config file, message: %s", err)
+	}
+	config.Peers = append(config.Peers, connection.Peer{
+		WgPubKey:     key,
+		WgAllowedIps: allowedIPs,
+	})
+
+	err = config.Write(configPath)
+	if err != nil {
+		log.Errorf("failed writing config to %s: %s", config, err.Error())
+		return status.Errorf(codes.Internal, "failed writing config to %s: %s", configPath, err)
+	}
+
+	return nil
+}

cmd/config.go

@@ -2,11 +2,12 @@ package cmd
 
 import (
 	"encoding/json"
-	ice "github.com/pion/ice/v2"
-	"github.com/wiretrustee/wiretrustee/connection"
 	"io/ioutil"
 	"os"
 	"path/filepath"
+
+	ice "github.com/pion/ice/v2"
+	"github.com/wiretrustee/wiretrustee/connection"
 )
 
 // Config Configuration type
@@ -17,6 +18,7 @@ type Config struct {
 	StunTurnURLs []*ice.URL
 	// host:port of the signal server
 	SignalAddr     string
+	ManagementAddr string
 	WgAddr         string
 	WgIface        string
 	IFaceBlackList []string

cmd/init.go

@@ -1,12 +1,13 @@
 package cmd
 
 import (
+	"os"
+	"strings"
+
 	ice "github.com/pion/ice/v2"
 	log "github.com/sirupsen/logrus"
 	"github.com/spf13/cobra"
 	"golang.zx2c4.com/wireguard/wgctrl/wgtypes"
-	"os"
-	"strings"
 )
 
 var (
@@ -14,6 +15,7 @@ var (
 	wgInterface    string
 	wgLocalAddr    string
 	signalAddr     string
+	managementAddr string
 	stunURLs       string
 	turnURLs       string
 
@@ -85,6 +87,7 @@ var (
 				Peers:          nil,
 				StunTurnURLs:   stunTurnURLs,
 				SignalAddr:     signalAddr,
+				ManagementAddr: managementAddr,
 				WgAddr:         wgLocalAddr,
 				WgIface:        wgInterface,
 			}
@@ -105,12 +108,14 @@ func init() {
 	initCmd.PersistentFlags().StringVar(&wgInterface, "wgInterface", "wiretrustee0", "Wireguard interface name, e.g. wiretreustee0 or wg0")
 	initCmd.PersistentFlags().StringVar(&wgLocalAddr, "wgLocalAddr", "", "Wireguard local address, e.g. 10.30.30.1/24")
 	initCmd.PersistentFlags().StringVar(&signalAddr, "signalAddr", "", "Signal server address, e.g. signal.wiretrustee.com:10000")
+	initCmd.PersistentFlags().StringVar(&managementAddr, "managementAddr", "", "Management server address, e.g. management.wiretrustee.com:10000")
 	initCmd.PersistentFlags().StringVar(&stunURLs, "stunURLs", "", "Comma separated STUN server URLs: protocol:host:port, e.g. stun:stun.l.google.com:19302,stun:stun1.l.google.com:19302")
 	//todo user:password@protocol:host:port not the best way to pass TURN credentials, do it according to https://tools.ietf.org/html/rfc7065 E.g. use oauth
 	initCmd.PersistentFlags().StringVar(&turnURLs, "turnURLs", "", "Comma separated TURN server URLs: user:password@protocol:host:port, e.g. user:password@turn:stun.wiretrustee.com:3468")
 	//initCmd.MarkPersistentFlagRequired("configPath")
 	initCmd.MarkPersistentFlagRequired("wgLocalAddr")    //nolint
 	initCmd.MarkPersistentFlagRequired("signalAddr")     //nolint
+	initCmd.MarkPersistentFlagRequired("managementAddr") //nolint
 	initCmd.MarkPersistentFlagRequired("stunURLs")       //nolint
 	initCmd.MarkPersistentFlagRequired("turnURLs")       //nolint
 }

cmd/service_test.go

@@ -2,10 +2,11 @@ package cmd
 
 import (
 	"bytes"
-	"github.com/kardianos/service"
 	"io/ioutil"
 	"os"
 	"testing"
+
+	"github.com/kardianos/service"
 )
 
 func Test_ServiceInstallCMD(t *testing.T) {
@@ -65,6 +66,8 @@ func Test_ServiceRunCMD(t *testing.T) {
 		"stun:stun.wiretrustee.com:3468",
 		"--signalAddr",
 		"signal.wiretrustee.com:10000",
+		"--managementAddr",
+		"management.wiretrustee.com:10000",
 		"--turnURLs",
 		"foo:bar@turn:stun.wiretrustee.com:3468",
 		"--wgInterface",

cmd/up.go

@@ -2,15 +2,27 @@ package cmd
 
 import (
 	"context"
+	"io"
+	"os"
+	"strings"
+	"time"
+
 	log "github.com/sirupsen/logrus"
 	"github.com/spf13/cobra"
 	"github.com/wiretrustee/wiretrustee/connection"
+	"github.com/wiretrustee/wiretrustee/encryption"
+	mgm "github.com/wiretrustee/wiretrustee/management/proto"
 	sig "github.com/wiretrustee/wiretrustee/signal"
 	"golang.zx2c4.com/wireguard/wgctrl/wgtypes"
-	"os"
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/keepalive"
+	"google.golang.org/grpc/status"
 )
 
 var (
+	setupKey string
+
 	upCmd = &cobra.Command{
 		Use:   "up",
 		Short: "start wiretrustee",
@@ -28,8 +40,9 @@ var (
 				os.Exit(ExitSetupFailed)
 			}
 
-			var sigTLSEnabled = false
+			go processManagement(config.ManagementAddr, setupKey, myKey)
 
+			var sigTLSEnabled = false
 			ctx := context.Background()
 			signalClient, err := sig.NewClient(ctx, config.SignalAddr, myKey, sigTLSEnabled)
 			if err != nil {
@@ -58,3 +71,107 @@ var (
 		},
 	}
 )
+
+func init() {
+	upCmd.PersistentFlags().StringVar(&setupKey, "setupKey", "", "Setup key to join a network, if not specified a new network will be created")
+}
+
+func processManagement(managementAddr string, setupKey string, ourPrivateKey wgtypes.Key) {
+	err := connectToManagement(managementAddr, setupKey, ourPrivateKey)
+	if err != nil {
+		log.Errorf("Failed to connect to managment server: %s", err)
+		os.Exit(ExitSetupFailed)
+	}
+
+	for {
+		_ = connectToManagement(managementAddr, setupKey, ourPrivateKey)
+	}
+}
+
+func connectToManagement(managementAddr string, setupKey string, ourPrivateKey wgtypes.Key) error {
+	log.Printf("Connecting to management server %s", managementAddr)
+	mgmCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+	defer cancel()
+	mgmConn, err := grpc.DialContext(
+		mgmCtx,
+		managementAddr,
+		grpc.WithInsecure(),
+		grpc.WithBlock(),
+		grpc.WithKeepaliveParams(keepalive.ClientParameters{
+			Time:    3 * time.Second,
+			Timeout: 2 * time.Second,
+		}))
+
+	if err != nil {
+		return status.Errorf(codes.FailedPrecondition, "Error while connecting to the Management Service %s: %s", managementAddr, err)
+	}
+	defer mgmConn.Close()
+
+	log.Printf("Connected to management server %s", managementAddr)
+
+	mgmClient := mgm.NewManagementServiceClient(mgmConn)
+	serverKeyResponse, err := mgmClient.GetServerKey(mgmCtx, &mgm.Empty{})
+	if err != nil {
+		return status.Errorf(codes.FailedPrecondition, "Error while getting server key: %s", err)
+	}
+
+	serverKey, err := wgtypes.ParseKey(serverKeyResponse.Key)
+	if err != nil {
+		return status.Errorf(codes.FailedPrecondition, "Failed parsing Wireguard public server key %s: [%s]", serverKeyResponse.Key, err)
+	}
+
+	_, err = mgmClient.RegisterPeer(mgmCtx, &mgm.RegisterPeerRequest{Key: ourPrivateKey.PublicKey().String(), SetupKey: setupKey})
+	if err != nil {
+		return status.Errorf(codes.FailedPrecondition, "Error while registering account: %s", err)
+	}
+
+	log.Println("Peer registered")
+	updatePeers(mgmClient, serverKey, ourPrivateKey)
+	return nil
+}
+
+func updatePeers(mgmClient mgm.ManagementServiceClient, remotePubKey wgtypes.Key, ourPrivateKey wgtypes.Key) {
+	log.Printf("Getting peers updates")
+	ctx := context.Background()
+	req := &mgm.SyncRequest{}
+	encryptedReq, err := encryption.EncryptMessage(remotePubKey, ourPrivateKey, req)
+	if err != nil {
+		log.Errorf("Failed to encrypt message: %s", err)
+		return
+	}
+
+	syncReq := &mgm.EncryptedMessage{WgPubKey: ourPrivateKey.PublicKey().String(), Body: encryptedReq}
+	stream, err := mgmClient.Sync(ctx, syncReq)
+	if err != nil {
+		log.Errorf("Failed to open management stream: %s", err)
+		return
+	}
+	for {
+		update, err := stream.Recv()
+		if err == io.EOF {
+			log.Errorf("Managment stream was closed: %s", err)
+			return
+		}
+		if err != nil {
+			log.Errorf("Managment stream disconnected: %s", err)
+			return
+		}
+
+		log.Infof("Got peers update")
+		resp := &mgm.SyncResponse{}
+		err = encryption.DecryptMessage(remotePubKey, ourPrivateKey, update.Body, resp)
+		if err != nil {
+			log.Errorf("Failed to decrypt message: %s", err)
+			return
+		}
+
+		for _, peer := range resp.RemotePeers {
+			log.Infof("Peer: %s", peer)
+			_ = addPeer(peer.WgPubKey, strings.Join(peer.AllowedIps, ","))
+		}
+
+		// for _, peer := range resp.RemotePeers {
+		//	log.Infof("Peer: %s", peer.WgPubKey)
+		//}
+	}
+}

go.mod

@@ -5,6 +5,7 @@ go 1.16
 require (
 	github.com/cenkalti/backoff/v4 v4.1.0
 	github.com/golang/protobuf v1.5.2
+	github.com/google/uuid v1.2.0 // indirect
 	github.com/kardianos/service v1.2.0
 	github.com/onsi/ginkgo v1.16.4
 	github.com/onsi/gomega v1.13.0

management/server/account.go

@@ -1,6 +1,8 @@
 package server
 
 import (
+	"github.com/google/uuid"
+	log "github.com/sirupsen/logrus"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/status"
 	"net"
@@ -18,7 +20,7 @@ type Account struct {
 	Id        string
 	SetupKeys map[string]*SetupKey
 	Network   *Network
-	Peers     []*Peer
+	Peers     map[string]*Peer
 }
 
 // SetupKey represents a pre-authorized key used to register machines (peers)
@@ -84,14 +86,23 @@ func (manager *AccountManager) GetPeersForAPeer(peerKey string) ([]*Peer, 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?
 func (manager *AccountManager) AddPeer(setupKey string, peerKey string) (*Peer, error) {
 	manager.mux.Lock()
 	defer manager.mux.Unlock()
 
-	account, err := manager.Store.GetAccountBySetupKey(setupKey)
+	var account *Account
+	var err error
+	var sk *SetupKey
+	if len(setupKey) == 0 {
+		// Empty setup key, create a new account for it.
+		account, sk = manager.newAccount()
+	} else {
+		sk = &SetupKey{Key: setupKey}
+		account, err = manager.Store.GetAccountBySetupKey(sk.Key)
 		if err != nil {
-		//todo
+			return nil, status.Errorf(codes.NotFound, "unknown setupKey %s", setupKey)
-		return nil, err
+		}
 	}
 
 	var takenIps []net.IP
@@ -104,11 +115,11 @@ func (manager *AccountManager) AddPeer(setupKey string, peerKey string) (*Peer,
 
 	newPeer := &Peer{
 		Key:      peerKey,
-		SetupKey: &SetupKey{Key: setupKey},
+		SetupKey: sk,
 		IP:       nextIp,
 	}
 
-	account.Peers = append(account.Peers, newPeer)
+	account.Peers[newPeer.Key] = newPeer
 	err = manager.Store.SaveAccount(account)
 	if err != nil {
 		return nil, status.Errorf(codes.Internal, "failed adding peer")
@@ -117,3 +128,21 @@ func (manager *AccountManager) AddPeer(setupKey string, peerKey string) (*Peer,
 	return newPeer, nil
 
 }
+
+// newAccount creates a new Account with a default SetupKey (doesn't store in a Store)
+func (manager *AccountManager) newAccount() (*Account, *SetupKey) {
+
+	log.Debugf("creating new account")
+
+	accountId := uuid.New().String()
+	setupKeyId := uuid.New().String()
+	setupKeys := make(map[string]*SetupKey)
+	setupKey := &SetupKey{Key: setupKeyId}
+	setupKeys[setupKeyId] = setupKey
+	network := &Network{Id: uuid.New().String(), Net: net.IPNet{}, Dns: ""}
+	peers := make(map[string]*Peer)
+
+	log.Debugf("created new account %s with setup key %s", accountId, setupKeyId)
+
+	return &Account{Id: accountId, SetupKeys: setupKeys, Network: network, Peers: peers}, setupKey
+}

management/server/file_store.go

@@ -96,11 +96,9 @@ func (s *FileStore) GetPeer(peerKey string) (*Peer, error) {
 		return nil, err
 	}
 
-	for _, peer := range account.Peers {
+	if peer, ok := account.Peers[peerKey]; ok {
-		if peer.Key == peerKey {
 		return peer, nil
 	}
-	}
 
 	return nil, status.Errorf(codes.NotFound, "peer not found")
 }

management/server/testdata/store.json

@@ -14,7 +14,8 @@
                     "Mask": "/8AAAA=="
                 },
                 "Dns": null
-            }
+            },
+            "Peers": {}
         }
     }
 }