Support new Management service protocol (NetworkMap) (#193)

* feature: support new management service protocol

* chore: add more logging to track networkmap serial

* refactor: organize peer update code in engine

* chore: fix lint issues

* refactor: extract Signal client interface

* test: add signal client mock

* refactor: introduce Management Service client interface

* chore: place management and signal clients mocks to respective packages

* test: add Serial test to the engine

* fix: lint issues

* test: unit tests for a networkMapUpdate

* test: unit tests Sync update
This commit is contained in:
Mikhail Bragin
2022-01-18 16:44:58 +01:00
committed by GitHub
parent 9a3fba3fa3
commit 5db130a12e
14 changed files with 1102 additions and 651 deletions

View File

@@ -82,7 +82,7 @@ var (
)
// loginPeer attempts to login to Management Service. If peer wasn't registered, tries the registration flow.
func loginPeer(serverPublicKey wgtypes.Key, client *mgm.Client, setupKey string) (*mgmProto.LoginResponse, error) {
func loginPeer(serverPublicKey wgtypes.Key, client *mgm.GrpcClient, setupKey string) (*mgmProto.LoginResponse, error) {
loginResp, err := client.Login(serverPublicKey)
if err != nil {
@@ -101,7 +101,7 @@ func loginPeer(serverPublicKey wgtypes.Key, client *mgm.Client, setupKey string)
// registerPeer checks whether setupKey was provided via cmd line and if not then it prompts user to enter a key.
// Otherwise tries to register with the provided setupKey via command line.
func registerPeer(serverPublicKey wgtypes.Key, client *mgm.Client, setupKey string) (*mgmProto.LoginResponse, error) {
func registerPeer(serverPublicKey wgtypes.Key, client *mgm.GrpcClient, setupKey string) (*mgmProto.LoginResponse, error) {
var err error
if setupKey == "" {

View File

@@ -83,7 +83,7 @@ func createEngineConfig(key wgtypes.Key, config *internal.Config, peerConfig *mg
}
// connectToSignal creates Signal Service client and established a connection
func connectToSignal(ctx context.Context, wtConfig *mgmProto.WiretrusteeConfig, ourPrivateKey wgtypes.Key) (*signal.Client, error) {
func connectToSignal(ctx context.Context, wtConfig *mgmProto.WiretrusteeConfig, ourPrivateKey wgtypes.Key) (*signal.GrpcClient, error) {
var sigTLSEnabled bool
if wtConfig.Signal.Protocol == mgmProto.HostConfig_HTTPS {
sigTLSEnabled = true
@@ -101,7 +101,7 @@ func connectToSignal(ctx context.Context, wtConfig *mgmProto.WiretrusteeConfig,
}
// connectToManagement creates Management Services client, establishes a connection, logs-in and gets a global Wiretrustee config (signal, turn, stun hosts, etc)
func connectToManagement(ctx context.Context, managementAddr string, ourPrivateKey wgtypes.Key, tlsEnabled bool) (*mgm.Client, *mgmProto.LoginResponse, error) {
func connectToManagement(ctx context.Context, managementAddr string, ourPrivateKey wgtypes.Key, tlsEnabled bool) (*mgm.GrpcClient, *mgmProto.LoginResponse, error) {
log.Debugf("connecting to management server %s", managementAddr)
client, err := mgm.NewClient(ctx, managementAddr, ourPrivateKey, tlsEnabled)
if err != nil {

View File

@@ -12,6 +12,7 @@ import (
mgmProto "github.com/wiretrustee/wiretrustee/management/proto"
signal "github.com/wiretrustee/wiretrustee/signal/client"
sProto "github.com/wiretrustee/wiretrustee/signal/proto"
"github.com/wiretrustee/wiretrustee/util"
"golang.zx2c4.com/wireguard/wgctrl/wgtypes"
"math/rand"
"strings"
@@ -44,9 +45,9 @@ type EngineConfig struct {
// Engine is a mechanism responsible for reacting on Signal and Management stream events and managing connections to the remote peers.
type Engine struct {
// signal is a Signal Service client
signal *signal.Client
signal signal.Client
// mgmClient is a Management Service client
mgmClient *mgm.Client
mgmClient mgm.Client
// peerConns is a map that holds all the peers that are known to this peer
peerConns map[string]*peer.Conn
@@ -64,6 +65,9 @@ type Engine struct {
ctx context.Context
wgInterface iface.WGIface
// networkSerial is the latest Serial (state ID) of the network sent by the Management service
networkSerial uint64
}
// Peer is an instance of the Connection Peer
@@ -73,17 +77,18 @@ type Peer struct {
}
// NewEngine creates a new Connection Engine
func NewEngine(signalClient *signal.Client, mgmClient *mgm.Client, config *EngineConfig, cancel context.CancelFunc, ctx context.Context) *Engine {
func NewEngine(signalClient signal.Client, mgmClient mgm.Client, config *EngineConfig, cancel context.CancelFunc, ctx context.Context) *Engine {
return &Engine{
signal: signalClient,
mgmClient: mgmClient,
peerConns: map[string]*peer.Conn{},
syncMsgMux: &sync.Mutex{},
config: config,
STUNs: []*ice.URL{},
TURNs: []*ice.URL{},
cancel: cancel,
ctx: ctx,
signal: signalClient,
mgmClient: mgmClient,
peerConns: map[string]*peer.Conn{},
syncMsgMux: &sync.Mutex{},
config: config,
STUNs: []*ice.URL{},
TURNs: []*ice.URL{},
cancel: cancel,
ctx: ctx,
networkSerial: 0,
}
}
@@ -91,7 +96,7 @@ func (e *Engine) Stop() error {
e.syncMsgMux.Lock()
defer e.syncMsgMux.Unlock()
err := e.removeAllPeerConnections()
err := e.removeAllPeers()
if err != nil {
return err
}
@@ -146,8 +151,22 @@ func (e *Engine) Start() error {
return nil
}
func (e *Engine) removePeers(peers []string) error {
for _, p := range peers {
// removePeers finds and removes peers that do not exist anymore in the network map received from the Management Service
func (e *Engine) removePeers(peersUpdate []*mgmProto.RemotePeerConfig) error {
currentPeers := make([]string, 0, len(e.peerConns))
for p := range e.peerConns {
currentPeers = append(currentPeers, p)
}
newPeers := make([]string, 0, len(peersUpdate))
for _, p := range peersUpdate {
newPeers = append(newPeers, p.GetWgPubKey())
}
toRemove := util.SliceDiff(currentPeers, newPeers)
for _, p := range toRemove {
err := e.removePeer(p)
if err != nil {
return err
@@ -157,7 +176,7 @@ func (e *Engine) removePeers(peers []string) error {
return nil
}
func (e *Engine) removeAllPeerConnections() error {
func (e *Engine) removeAllPeers() error {
log.Debugf("removing all peer connections")
for p := range e.peerConns {
err := e.removePeer(p)
@@ -189,6 +208,16 @@ func (e *Engine) GetPeerConnectionStatus(peerKey string) peer.ConnStatus {
return -1
}
func (e *Engine) GetPeers() []string {
e.syncMsgMux.Lock()
defer e.syncMsgMux.Unlock()
peers := []string{}
for s := range e.peerConns {
peers = append(peers, s)
}
return peers
}
// GetConnectedPeers returns a connection Status or nil if peer connection wasn't found
func (e *Engine) GetConnectedPeers() []string {
@@ -205,7 +234,7 @@ func (e *Engine) GetConnectedPeers() []string {
return peers
}
func signalCandidate(candidate ice.Candidate, myKey wgtypes.Key, remoteKey wgtypes.Key, s *signal.Client) error {
func signalCandidate(candidate ice.Candidate, myKey wgtypes.Key, remoteKey wgtypes.Key, s signal.Client) error {
err := s.Send(&sProto.Message{
Key: myKey.PublicKey().String(),
RemoteKey: remoteKey.String(),
@@ -223,7 +252,7 @@ func signalCandidate(candidate ice.Candidate, myKey wgtypes.Key, remoteKey wgtyp
return nil
}
func signalAuth(uFrag string, pwd string, myKey wgtypes.Key, remoteKey wgtypes.Key, s *signal.Client, isAnswer bool) error {
func signalAuth(uFrag string, pwd string, myKey wgtypes.Key, remoteKey wgtypes.Key, s signal.Client, isAnswer bool) error {
var t sProto.Body_Type
if isAnswer {
@@ -246,37 +275,42 @@ func signalAuth(uFrag string, pwd string, myKey wgtypes.Key, remoteKey wgtypes.K
return nil
}
func (e *Engine) handleSync(update *mgmProto.SyncResponse) error {
e.syncMsgMux.Lock()
defer e.syncMsgMux.Unlock()
if update.GetWiretrusteeConfig() != nil {
err := e.updateTURNs(update.GetWiretrusteeConfig().GetTurns())
if err != nil {
return err
}
err = e.updateSTUNs(update.GetWiretrusteeConfig().GetStuns())
if err != nil {
return err
}
//todo update signal
}
if update.GetNetworkMap() != nil {
// only apply new changes and ignore old ones
err := e.updateNetworkMap(update.GetNetworkMap())
if err != nil {
return err
}
}
return nil
}
// receiveManagementEvents connects to the Management Service event stream to receive updates from the management service
// E.g. when a new peer has been registered and we are allowed to connect to it.
func (e *Engine) receiveManagementEvents() {
go func() {
err := e.mgmClient.Sync(func(update *mgmProto.SyncResponse) error {
e.syncMsgMux.Lock()
defer e.syncMsgMux.Unlock()
if update.GetWiretrusteeConfig() != nil {
err := e.updateTURNs(update.GetWiretrusteeConfig().GetTurns())
if err != nil {
return err
}
err = e.updateSTUNs(update.GetWiretrusteeConfig().GetStuns())
if err != nil {
return err
}
//todo update signal
}
if update.GetRemotePeers() != nil || update.GetRemotePeersIsEmpty() {
// empty arrays are serialized by protobuf to null, but for our case empty array is a valid state.
err := e.updatePeers(update.GetRemotePeers())
if err != nil {
return err
}
}
return nil
return e.handleSync(update)
})
if err != nil {
// happens if management is unavailable for a long time.
@@ -327,27 +361,41 @@ func (e *Engine) updateTURNs(turns []*mgmProto.ProtectedHostConfig) error {
return nil
}
func (e *Engine) updatePeers(remotePeers []*mgmProto.RemotePeerConfig) error {
log.Debugf("got peers update from Management Service, total peers to connect to = %d", len(remotePeers))
remotePeerMap := make(map[string]struct{})
for _, p := range remotePeers {
remotePeerMap[p.GetWgPubKey()] = struct{}{}
func (e *Engine) updateNetworkMap(networkMap *mgmProto.NetworkMap) error {
serial := networkMap.GetSerial()
if e.networkSerial > serial {
log.Debugf("received outdated NetworkMap with serial %d, ignoring", serial)
return nil
}
//remove peers that are no longer available for us
toRemove := []string{}
for p := range e.peerConns {
if _, ok := remotePeerMap[p]; !ok {
toRemove = append(toRemove, p)
log.Debugf("got peers update from Management Service, total peers to connect to = %d", len(networkMap.GetRemotePeers()))
// cleanup request, most likely our peer has been deleted
if networkMap.GetRemotePeersIsEmpty() {
err := e.removeAllPeers()
if err != nil {
return err
}
} else {
err := e.removePeers(networkMap.GetRemotePeers())
if err != nil {
return err
}
err = e.addNewPeers(networkMap.GetRemotePeers())
if err != nil {
return err
}
}
err := e.removePeers(toRemove)
if err != nil {
return err
}
// add new peers
for _, p := range remotePeers {
e.networkSerial = serial
return nil
}
// addNewPeers finds and adds peers that were not know before but arrived from the Management service with the update
func (e *Engine) addNewPeers(peersUpdate []*mgmProto.RemotePeerConfig) error {
for _, p := range peersUpdate {
peerKey := p.GetWgPubKey()
peerIPs := p.GetAllowedIps()
if _, ok := e.peerConns[peerKey]; !ok {

View File

@@ -37,6 +37,232 @@ var (
}
)
func TestEngine_UpdateNetworkMap(t *testing.T) {
// test setup
key, err := wgtypes.GeneratePrivateKey()
if err != nil {
t.Fatal(err)
return
}
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
engine := NewEngine(&signal.MockClient{}, &mgmt.MockClient{}, &EngineConfig{
WgIfaceName: "utun100",
WgAddr: "100.64.0.1/24",
WgPrivateKey: key,
WgPort: 33100,
}, cancel, ctx)
type testCase struct {
idx int
networkMap *mgmtProto.NetworkMap
expectedLen int
expectedPeers []string
expectedSerial uint64
}
peer1 := &mgmtProto.RemotePeerConfig{
WgPubKey: "RRHf3Ma6z6mdLbriAJbqhX7+nM/B71lgw2+91q3LfhU=",
AllowedIps: []string{"100.64.0.10/24"},
}
peer2 := &mgmtProto.RemotePeerConfig{
WgPubKey: "LLHf3Ma6z6mdLbriAJbqhX7+nM/B71lgw2+91q3LfhU=",
AllowedIps: []string{"100.64.0.11/24"},
}
peer3 := &mgmtProto.RemotePeerConfig{
WgPubKey: "GGHf3Ma6z6mdLbriAJbqhX7+nM/B71lgw2+91q3LfhU=",
AllowedIps: []string{"100.64.0.12/24"},
}
// 1st case - new peer and network map has Serial grater than local => apply the update
case1 := testCase{
idx: 1,
networkMap: &mgmtProto.NetworkMap{
Serial: 1,
PeerConfig: nil,
RemotePeers: []*mgmtProto.RemotePeerConfig{
peer1,
},
RemotePeersIsEmpty: false,
},
expectedLen: 1,
expectedPeers: []string{peer1.GetWgPubKey()},
expectedSerial: 1,
}
// 2nd case - one extra peer added and network map has Serial grater than local => apply the update
case2 := testCase{
idx: 2,
networkMap: &mgmtProto.NetworkMap{
Serial: 2,
PeerConfig: nil,
RemotePeers: []*mgmtProto.RemotePeerConfig{
peer1, peer2,
},
RemotePeersIsEmpty: false,
},
expectedLen: 2,
expectedPeers: []string{peer1.GetWgPubKey(), peer2.GetWgPubKey()},
expectedSerial: 2,
}
// 3rd case - an update with 3 peers and Serial lower than the current serial of the engine => ignore the update
case3 := testCase{
idx: 3,
networkMap: &mgmtProto.NetworkMap{
Serial: 0,
PeerConfig: nil,
RemotePeers: []*mgmtProto.RemotePeerConfig{
peer1, peer2, peer3,
},
RemotePeersIsEmpty: false,
},
expectedLen: 2,
expectedPeers: []string{peer1.GetWgPubKey(), peer2.GetWgPubKey()},
expectedSerial: 2,
}
// 4th case - an update with 2 peers (1 new and 1 old) => apply the update removing old peer and adding a new one
case4 := testCase{
idx: 3,
networkMap: &mgmtProto.NetworkMap{
Serial: 4,
PeerConfig: nil,
RemotePeers: []*mgmtProto.RemotePeerConfig{
peer2, peer3,
},
RemotePeersIsEmpty: false,
},
expectedLen: 2,
expectedPeers: []string{peer2.GetWgPubKey(), peer3.GetWgPubKey()},
expectedSerial: 4,
}
// 5th case - an update with all peers to be removed
case5 := testCase{
idx: 3,
networkMap: &mgmtProto.NetworkMap{
Serial: 5,
PeerConfig: nil,
RemotePeers: []*mgmtProto.RemotePeerConfig{},
RemotePeersIsEmpty: true,
},
expectedLen: 0,
expectedPeers: nil,
expectedSerial: 5,
}
for _, c := range []testCase{case1, case2, case3, case4, case5} {
err = engine.updateNetworkMap(c.networkMap)
if err != nil {
t.Fatal(err)
return
}
if len(engine.peerConns) != c.expectedLen {
t.Errorf("case %d expecting Engine.peerConns to be of size %d, got %d", c.idx, c.expectedLen, len(engine.peerConns))
}
if engine.networkSerial != c.expectedSerial {
t.Errorf("case %d expecting Engine.networkSerial to be equal to %d, actual %d", c.idx, c.expectedSerial, engine.networkSerial)
}
for _, p := range c.expectedPeers {
if _, ok := engine.peerConns[p]; !ok {
t.Errorf("case %d expecting Engine.peerConns to contain peer %s", c.idx, p)
}
}
}
}
func TestEngine_Sync(t *testing.T) {
key, err := wgtypes.GeneratePrivateKey()
if err != nil {
t.Fatal(err)
return
}
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
// feed updates to Engine via mocked Management client
updates := make(chan *mgmtProto.SyncResponse)
defer close(updates)
syncFunc := func(msgHandler func(msg *mgmtProto.SyncResponse) error) error {
for msg := range updates {
err := msgHandler(msg)
if err != nil {
t.Fatal(err)
}
}
return nil
}
engine := NewEngine(&signal.MockClient{}, &mgmt.MockClient{SyncFunc: syncFunc}, &EngineConfig{
WgIfaceName: "utun100",
WgAddr: "100.64.0.1/24",
WgPrivateKey: key,
WgPort: 33100,
}, cancel, ctx)
defer func() {
err := engine.Stop()
if err != nil {
return
}
}()
err = engine.Start()
if err != nil {
t.Fatal(err)
return
}
peer1 := &mgmtProto.RemotePeerConfig{
WgPubKey: "RRHf3Ma6z6mdLbriAJbqhX7+nM/B71lgw2+91q3LfhU=",
AllowedIps: []string{"100.64.0.10/24"},
}
peer2 := &mgmtProto.RemotePeerConfig{
WgPubKey: "LLHf3Ma6z6mdLbriAJbqhX9+nM/B71lgw2+91q3LlhU=",
AllowedIps: []string{"100.64.0.11/24"},
}
peer3 := &mgmtProto.RemotePeerConfig{
WgPubKey: "GGHf3Ma6z6mdLbriAJbqhX9+nM/B71lgw2+91q3LlhU=",
AllowedIps: []string{"100.64.0.12/24"},
}
// 1st update with just 1 peer and serial larger than the current serial of the engine => apply update
updates <- &mgmtProto.SyncResponse{
NetworkMap: &mgmtProto.NetworkMap{
Serial: 10,
PeerConfig: nil,
RemotePeers: []*mgmtProto.RemotePeerConfig{peer1, peer2, peer3},
RemotePeersIsEmpty: false,
},
}
timeout := time.After(time.Second * 2)
for {
select {
case <-timeout:
t.Fatalf("timeout while waiting for test to finish")
default:
}
if len(engine.GetPeers()) == 3 && engine.networkSerial == 10 {
break
}
}
}
func TestEngine_MultiplePeers(t *testing.T) {
//log.SetLevel(log.DebugLevel)
@@ -58,23 +284,14 @@ func TestEngine_MultiplePeers(t *testing.T) {
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
sport := 10010
signalServer, err := startSignal(sport)
sigServer, err := startSignal(sport)
if err != nil {
t.Fatal(err)
return
}
defer signalServer.Stop()
defer sigServer.Stop()
mport := 33081
mgmtServer, err := startManagement(mport, &server.Config{
Stuns: []*server.Host{},
TURNConfig: &server.TURNConfig{},
Signal: &server.Host{
Proto: "http",
URI: "localhost:10000",
},
Datadir: dir,
HttpConfig: nil,
})
mgmtServer, err := startManagement(mport, dir)
if err != nil {
t.Fatal(err)
return
@@ -201,7 +418,18 @@ func startSignal(port int) (*grpc.Server, error) {
return s, nil
}
func startManagement(port int, config *server.Config) (*grpc.Server, error) {
func startManagement(port int, dataDir string) (*grpc.Server, error) {
config := &server.Config{
Stuns: []*server.Host{},
TURNConfig: &server.TURNConfig{},
Signal: &server.Host{
Proto: "http",
URI: "localhost:10000",
},
Datadir: dataDir,
HttpConfig: nil,
}
lis, err := net.Listen("tcp", fmt.Sprintf("localhost:%d", port))
if err != nil {

View File

@@ -419,3 +419,7 @@ func (conn *Conn) OnRemoteCandidate(candidate ice.Candidate) {
}
}()
}
func (conn *Conn) GetKey() string {
return conn.config.Key
}