sendnrw/netbird
2
0
Fork 0
mirror of https://github.com/netbirdio/netbird.git synced 2026-04-18 16:26:38 +00:00
Code Issues Packages Projects Releases Wiki Activity

Extend Management to support peer changes distribution (#55)

Browse Source 
* feature: add peer sync and a server public key endpoints
* test: add Management.Sync() gRpc endpoint test
* feat: implement peer sync
* docs: added some comments to the Management server
* chore: use for loop over channel when monitoring peer updates
* fix: exit infinite loop when sending updates to peers
* test: add multiple concurrent peers test for management service
* chore: remove unused test
* fix: reduce the amount peers for a concurrent peer update test

Co-authored-by: braginini <m.bragin@wiretrustee.com>
This commit is contained in:
Mikhail Bragin
2021-07-22 10:28:00 +02:00
committed by GitHub
parent d27eb317aa
commit 079d35eada
12 changed files with 974 additions and 121 deletions
Download Patch File Download Diff File Expand all files Collapse all files

294
management/management_test.go
View File

@@ -2,16 +2,20 @@ package management_test
import (
"context"
pb "github.com/golang/protobuf/proto" //nolint
log "github.com/sirupsen/logrus"
"github.com/wiretrustee/wiretrustee/signal"
"io"
"io/ioutil"
"math/rand"
"net"
"os"
"path/filepath"
"strings"
sync2 "sync"
"time"
. "github.com/onsi/ginkgo"
. "github.com/onsi/gomega"
log "github.com/sirupsen/logrus"
mgmt "github.com/wiretrustee/wiretrustee/management"
mgmtProto "github.com/wiretrustee/wiretrustee/management/proto"
"github.com/wiretrustee/wiretrustee/util"
@@ -20,16 +24,26 @@ import (
"google.golang.org/grpc/keepalive"
)
var _ = Describe("Client", func() {
const (
ValidSetupKey = "A2C8E62B-38F5-4553-B31E-DD66C696CEBB"
InvalidSetupKey = "INVALID_SETUP_KEY"
)
var _ = Describe("Management service", func() {
var (
addr string
server *grpc.Server
tmpDir string
dataDir string
addr string
server *grpc.Server
tmpDir string
dataDir string
client mgmtProto.ManagementServiceClient
serverPubKey wgtypes.Key
conn *grpc.ClientConn
)
BeforeEach(func() {
level, _ := log.ParseLevel("Debug")
log.SetLevel(level)
var err error
dataDir, err = ioutil.TempDir("", "wiretrustee_mgmt_test_tmp_*")
Expect(err).NotTo(HaveOccurred())
@@ -39,38 +53,154 @@ var _ = Describe("Client", func() {
var listener net.Listener
server, listener = startServer(dataDir)
addr = listener.Addr().String()
client, conn = createRawClient(addr)
// server public key
resp, err := client.GetServerKey(context.TODO(), &mgmtProto.Empty{})
Expect(err).NotTo(HaveOccurred())
serverPubKey, err = wgtypes.ParseKey(resp.Key)
Expect(err).NotTo(HaveOccurred())
})
AfterEach(func() {
server.Stop()
err := os.RemoveAll(tmpDir)
err := conn.Close()
Expect(err).NotTo(HaveOccurred())
err = os.RemoveAll(tmpDir)
Expect(err).NotTo(HaveOccurred())
})
Describe("Service health", func() {
Context("when it has been started", func() {
It("should be ok", func() {
client := createRawClient(addr)
healthy, err := client.IsHealthy(context.TODO(), &mgmtProto.Empty{})
Context("when calling IsHealthy endpoint", func() {
Specify("a non-error result is returned", func() {
Expect(healthy).ToNot(BeNil())
Expect(err).To(BeNil())
healthy, err := client.IsHealthy(context.TODO(), &mgmtProto.Empty{})
Expect(err).NotTo(HaveOccurred())
Expect(healthy).ToNot(BeNil())
})
})
Context("when calling Sync endpoint", func() {
Context("when there are 3 peers registered under one account", func() {
Specify("a list containing other 2 peers is returned", func() {
key, _ := wgtypes.GenerateKey()
key1, _ := wgtypes.GenerateKey()
key2, _ := wgtypes.GenerateKey()
registerPeerWithValidSetupKey(key, client)
registerPeerWithValidSetupKey(key1, client)
registerPeerWithValidSetupKey(key2, client)
messageBytes, err := pb.Marshal(&mgmtProto.SyncRequest{})
Expect(err).NotTo(HaveOccurred())
encryptedBytes, err := signal.Encrypt(messageBytes, serverPubKey, key)
Expect(err).NotTo(HaveOccurred())
sync, err := client.Sync(context.TODO(), &mgmtProto.EncryptedMessage{
WgPubKey: key.PublicKey().String(),
Body: encryptedBytes,
})
Expect(err).NotTo(HaveOccurred())
encryptedResponse := &mgmtProto.EncryptedMessage{}
err = sync.RecvMsg(encryptedResponse)
Expect(err).NotTo(HaveOccurred())
decryptedBytes, err := signal.Decrypt(encryptedResponse.Body, serverPubKey, key)
Expect(err).NotTo(HaveOccurred())
resp := &mgmtProto.SyncResponse{}
err = pb.Unmarshal(decryptedBytes, resp)
Expect(err).NotTo(HaveOccurred())
Expect(resp.Peers).To(HaveLen(2))
Expect(resp.Peers).To(ContainElements(key1.PublicKey().String(), key2.PublicKey().String()))
})
})
Context("when there is a new peer registered", func() {
Specify("an update is returned", func() {
// register only a single peer
key, _ := wgtypes.GenerateKey()
registerPeerWithValidSetupKey(key, client)
messageBytes, err := pb.Marshal(&mgmtProto.SyncRequest{})
Expect(err).NotTo(HaveOccurred())
encryptedBytes, err := signal.Encrypt(messageBytes, serverPubKey, key)
Expect(err).NotTo(HaveOccurred())
sync, err := client.Sync(context.TODO(), &mgmtProto.EncryptedMessage{
WgPubKey: key.PublicKey().String(),
Body: encryptedBytes,
})
Expect(err).NotTo(HaveOccurred())
// after the initial sync call we have 0 peer updates
encryptedResponse := &mgmtProto.EncryptedMessage{}
err = sync.RecvMsg(encryptedResponse)
Expect(err).NotTo(HaveOccurred())
decryptedBytes, err := signal.Decrypt(encryptedResponse.Body, serverPubKey, key)
Expect(err).NotTo(HaveOccurred())
resp := &mgmtProto.SyncResponse{}
err = pb.Unmarshal(decryptedBytes, resp)
Expect(resp.Peers).To(HaveLen(0))
wg := sync2.WaitGroup{}
wg.Add(1)
// continue listening on updates for a peer
go func() {
err = sync.RecvMsg(encryptedResponse)
decryptedBytes, err = signal.Decrypt(encryptedResponse.Body, serverPubKey, key)
Expect(err).NotTo(HaveOccurred())
resp = &mgmtProto.SyncResponse{}
err = pb.Unmarshal(decryptedBytes, resp)
wg.Done()
}()
// register a new peer
key1, _ := wgtypes.GenerateKey()
registerPeerWithValidSetupKey(key1, client)
wg.Wait()
Expect(err).NotTo(HaveOccurred())
Expect(resp.Peers).To(ContainElements(key1.PublicKey().String()))
Expect(resp.Peers).To(HaveLen(1))
})
})
})
Describe("Registration", func() {
Context("of a new peer without a valid setup key", func() {
It("should fail", func() {
Context("when calling GetServerKey endpoint", func() {
Specify("a public Wireguard key of the service is returned", func() {
resp, err := client.GetServerKey(context.TODO(), &mgmtProto.Empty{})
Expect(err).NotTo(HaveOccurred())
Expect(resp).ToNot(BeNil())
Expect(resp.Key).ToNot(BeNil())
Expect(resp.ExpiresAt).ToNot(BeNil())
//check if the key is a valid Wireguard key
key, err := wgtypes.ParseKey(resp.Key)
Expect(err).NotTo(HaveOccurred())
Expect(key).ToNot(BeNil())
})
})
Context("when calling RegisterPeer endpoint", func() {
Context("with an invalid setup key", func() {
Specify("an error is returned", func() {
key, _ := wgtypes.GenerateKey()
setupKey := "invalid_setup_key"
client := createRawClient(addr)
resp, err := client.RegisterPeer(context.TODO(), &mgmtProto.RegisterPeerRequest{
Key: key.PublicKey().String(),
SetupKey: setupKey,
SetupKey: InvalidSetupKey,
})
Expect(err).To(HaveOccurred())
@@ -78,57 +208,101 @@ var _ = Describe("Client", func() {
})
})
})
Describe("Registration", func() {
Context("of a new peer with a valid setup key", func() {
It("should be successful", func() {
Context("with a valid setup key", func() {
It("a non error result is returned", func() {
key, _ := wgtypes.GenerateKey()
setupKey := "A2C8E62B-38F5-4553-B31E-DD66C696CEBB" //present in the testdata/store.json file
resp := registerPeerWithValidSetupKey(key, client)
client := createRawClient(addr)
resp, err := client.RegisterPeer(context.TODO(), &mgmtProto.RegisterPeerRequest{
Key: key.PublicKey().String(),
SetupKey: setupKey,
})
Expect(err).NotTo(HaveOccurred())
Expect(resp).ToNot(BeNil())
})
})
})
Describe("Registration", func() {
Context("of a new peer with a valid setup key", func() {
It("should be persisted to a file", func() {
Context("when there are 50 peers registered under one account", func() {
Context("when there are 10 more peers registered under the same account", func() {
Specify("all of the 50 peers will get updates of 10 newly registered peers", func() {
key, _ := wgtypes.GenerateKey()
setupKey := "A2C8E62B-38F5-4553-B31E-DD66C696CEBB" //present in the testdata/store.json file
initialPeers := 20
additionalPeers := 10
client := createRawClient(addr)
_, err := client.RegisterPeer(context.TODO(), &mgmtProto.RegisterPeerRequest{
Key: key.PublicKey().String(),
SetupKey: setupKey,
})
var peers []wgtypes.Key
for i := 0; i < initialPeers; i++ {
key, _ := wgtypes.GenerateKey()
registerPeerWithValidSetupKey(key, client)
peers = append(peers, key)
}
Expect(err).NotTo(HaveOccurred())
wg := sync2.WaitGroup{}
wg.Add(initialPeers + initialPeers*additionalPeers)
for _, peer := range peers {
messageBytes, err := pb.Marshal(&mgmtProto.SyncRequest{})
Expect(err).NotTo(HaveOccurred())
encryptedBytes, err := signal.Encrypt(messageBytes, serverPubKey, peer)
Expect(err).NotTo(HaveOccurred())
store, err := util.ReadJson(filepath.Join(dataDir, "store.json"), &mgmt.FileStore{})
Expect(err).NotTo(HaveOccurred())
// receive stream
peer := peer
go func() {
Expect(store.(*mgmt.FileStore)).NotTo(BeNil())
user := store.(*mgmt.FileStore).Accounts["bf1c8084-ba50-4ce7-9439-34653001fc3b"]
Expect(user.Peers[key.PublicKey().String()]).NotTo(BeNil())
Expect(user.SetupKeys[strings.ToLower(setupKey)]).NotTo(BeNil())
// open stream
sync, err := client.Sync(context.TODO(), &mgmtProto.EncryptedMessage{
WgPubKey: peer.PublicKey().String(),
Body: encryptedBytes,
})
Expect(err).NotTo(HaveOccurred())
for {
encryptedResponse := &mgmtProto.EncryptedMessage{}
err = sync.RecvMsg(encryptedResponse)
if err == io.EOF {
break
} else if err != nil {
Expect(err).NotTo(HaveOccurred())
}
decryptedBytes, err := signal.Decrypt(encryptedResponse.Body, serverPubKey, peer)
Expect(err).NotTo(HaveOccurred())
resp := &mgmtProto.SyncResponse{}
err = pb.Unmarshal(decryptedBytes, resp)
Expect(err).NotTo(HaveOccurred())
wg.Done()
}
}()
}
time.Sleep(1 * time.Second)
for i := 0; i < additionalPeers; i++ {
key, _ := wgtypes.GenerateKey()
registerPeerWithValidSetupKey(key, client)
rand.Seed(time.Now().UnixNano())
n := rand.Intn(500)
time.Sleep(time.Duration(n) * time.Millisecond)
}
wg.Wait()
})
})
})
})
func createRawClient(addr string) mgmtProto.ManagementServiceClient {
func registerPeerWithValidSetupKey(key wgtypes.Key, client mgmtProto.ManagementServiceClient) *mgmtProto.RegisterPeerResponse {
resp, err := client.RegisterPeer(context.TODO(), &mgmtProto.RegisterPeerRequest{
Key: key.PublicKey().String(),
SetupKey: ValidSetupKey,
})
Expect(err).NotTo(HaveOccurred())
return resp
}
func createRawClient(addr string) (mgmtProto.ManagementServiceClient, *grpc.ClientConn) {
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
defer cancel()
conn, err := grpc.DialContext(ctx, addr, grpc.WithInsecure(),
@@ -137,27 +311,21 @@ func createRawClient(addr string) mgmtProto.ManagementServiceClient {
Time: 10 * time.Second,
Timeout: 2 * time.Second,
}))
if err != nil {
Fail("failed creating raw signal client")
}
Expect(err).NotTo(HaveOccurred())
return mgmtProto.NewManagementServiceClient(conn)
return mgmtProto.NewManagementServiceClient(conn), conn
}
func startServer(dataDir string) (*grpc.Server, net.Listener) {
lis, err := net.Listen("tcp", ":0")
if err != nil {
panic(err)
}
Expect(err).NotTo(HaveOccurred())
s := grpc.NewServer()
server, err := mgmt.NewServer(dataDir)
if err != nil {
panic(err)
}
Expect(err).NotTo(HaveOccurred())
mgmtProto.RegisterManagementServiceServer(s, server)
go func() {
if err := s.Serve(lis); err != nil {
log.Fatalf("failed to serve: %v", err)
Expect(err).NotTo(HaveOccurred())
}
}()