mirror of
https://github.com/netbirdio/netbird.git
synced 2026-09-16 19:59:07 +02:00
Merge branch 'main' into embedded-vnc
This commit is contained in:
@@ -12,7 +12,7 @@ import (
|
||||
// Client is the interface for the management service client.
|
||||
type Client interface {
|
||||
io.Closer
|
||||
Sync(ctx context.Context, sysInfo *system.Info, msgHandler func(msg *proto.SyncResponse) error) error
|
||||
Sync(ctx context.Context, getInfo func(ctx context.Context) *system.Info, msgHandler func(msg *proto.SyncResponse) error) error
|
||||
Job(ctx context.Context, msgHandler func(msg *proto.JobRequest) *proto.JobResponse) error
|
||||
Register(setupKey string, jwtToken string, sysInfo *system.Info, sshKey []byte, dnsLabels domain.List) (*proto.LoginResponse, error)
|
||||
Login(sysInfo *system.Info, sshKey []byte, dnsLabels domain.List) (*proto.LoginResponse, error)
|
||||
|
||||
@@ -2,9 +2,11 @@ package client
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"net"
|
||||
"os"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
@@ -305,7 +307,7 @@ func TestClient_Sync(t *testing.T) {
|
||||
defer cancel()
|
||||
|
||||
go func() {
|
||||
err = client.Sync(ctx, info, func(msg *mgmtProto.SyncResponse) error {
|
||||
err = client.Sync(ctx, func(context.Context) *system.Info { return info }, func(msg *mgmtProto.SyncResponse) error {
|
||||
ch <- msg
|
||||
return nil
|
||||
})
|
||||
@@ -397,6 +399,75 @@ func wgKeyFromBytes(raw []byte) string {
|
||||
return k.String()
|
||||
}
|
||||
|
||||
func TestClient_SyncGathersInfoOnEveryConnect(t *testing.T) {
|
||||
s, lis, mgmtMockServer, serverKey := startMockManagement(t)
|
||||
defer s.GracefulStop()
|
||||
|
||||
testKey, err := wgtypes.GenerateKey()
|
||||
require.NoError(t, err)
|
||||
|
||||
hostnames := make(chan string, 2)
|
||||
mgmtMockServer.SyncFunc = func(msg *mgmtProto.EncryptedMessage, _ mgmtProto.ManagementService_SyncServer) error {
|
||||
peerKey, err := wgtypes.ParseKey(msg.GetWgPubKey())
|
||||
if err != nil {
|
||||
t.Errorf("invalid peer key: %v", err)
|
||||
return status.Error(codes.InvalidArgument, err.Error())
|
||||
}
|
||||
syncReq := &mgmtProto.SyncRequest{}
|
||||
if err := encryption.DecryptMessage(peerKey, serverKey, msg.Body, syncReq); err != nil {
|
||||
t.Errorf("decrypt sync request: %v", err)
|
||||
return status.Error(codes.InvalidArgument, err.Error())
|
||||
}
|
||||
select {
|
||||
case hostnames <- syncReq.GetMeta().GetHostname():
|
||||
default:
|
||||
}
|
||||
// Returning closes the stream, so the client reconnects and gathers again.
|
||||
return nil
|
||||
}
|
||||
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
defer cancel()
|
||||
|
||||
client, err := NewClient(ctx, lis.Addr().String(), testKey, false)
|
||||
require.NoError(t, err)
|
||||
|
||||
var gathers atomic.Int32
|
||||
done := make(chan struct{})
|
||||
go func() {
|
||||
defer close(done)
|
||||
_ = client.Sync(ctx, func(ctx context.Context) *system.Info {
|
||||
info := system.GetInfo(ctx)
|
||||
info.Hostname = fmt.Sprintf("host-%d", gathers.Add(1))
|
||||
return info
|
||||
}, func(*mgmtProto.SyncResponse) error { return nil })
|
||||
}()
|
||||
|
||||
// A connect attempt can fail before it reaches the server, so the sequence
|
||||
// numbers seen here may skip. What matters is that the reconnect carries a
|
||||
// newly gathered info instead of the one sent on the previous stream.
|
||||
var seen []int
|
||||
for len(seen) < 2 {
|
||||
select {
|
||||
case got := <-hostnames:
|
||||
var n int
|
||||
_, err := fmt.Sscanf(got, "host-%d", &n)
|
||||
require.NoError(t, err, "hostname should carry the gather sequence number")
|
||||
seen = append(seen, n)
|
||||
case <-time.After(10 * time.Second):
|
||||
t.Fatalf("timeout waiting for the second sync request, got %v", seen)
|
||||
}
|
||||
}
|
||||
assert.Greater(t, seen[1], seen[0], "the reconnect should carry a newly gathered info")
|
||||
|
||||
cancel()
|
||||
select {
|
||||
case <-done:
|
||||
case <-time.After(10 * time.Second):
|
||||
t.Fatal("timeout waiting for Sync to return after cancel")
|
||||
}
|
||||
}
|
||||
|
||||
func Test_SystemMetaDataFromClient(t *testing.T) {
|
||||
s, lis, mgmtMockServer, serverKey := startMockManagement(t)
|
||||
defer s.GracefulStop()
|
||||
|
||||
@@ -205,9 +205,9 @@ func (c *GrpcClient) ready() bool {
|
||||
|
||||
// Sync wraps the real client's Sync endpoint call and takes care of retries and encryption/decryption of messages
|
||||
// Blocking request. The result will be sent via msgHandler callback function
|
||||
func (c *GrpcClient) Sync(ctx context.Context, sysInfo *system.Info, msgHandler func(msg *proto.SyncResponse) error) error {
|
||||
func (c *GrpcClient) Sync(ctx context.Context, getInfo func(ctx context.Context) *system.Info, msgHandler func(msg *proto.SyncResponse) error) error {
|
||||
return c.withMgmtStream(ctx, func(ctx context.Context, serverPubKey wgtypes.Key, backOff backoff.BackOff) error {
|
||||
return c.handleSyncStream(ctx, serverPubKey, sysInfo, msgHandler, backOff)
|
||||
return c.handleSyncStream(ctx, serverPubKey, getInfo, msgHandler, backOff)
|
||||
})
|
||||
}
|
||||
|
||||
@@ -424,11 +424,11 @@ func (c *GrpcClient) sendJobResponse(
|
||||
return nil
|
||||
}
|
||||
|
||||
func (c *GrpcClient) handleSyncStream(ctx context.Context, serverPubKey wgtypes.Key, sysInfo *system.Info, msgHandler func(msg *proto.SyncResponse) error, backOff backoff.BackOff) error {
|
||||
func (c *GrpcClient) handleSyncStream(ctx context.Context, serverPubKey wgtypes.Key, getInfo func(ctx context.Context) *system.Info, msgHandler func(msg *proto.SyncResponse) error, backOff backoff.BackOff) error {
|
||||
ctx, cancelStream := context.WithCancel(ctx)
|
||||
defer cancelStream()
|
||||
|
||||
stream, err := c.connectToSyncStream(ctx, serverPubKey, sysInfo)
|
||||
stream, err := c.connectToSyncStream(ctx, serverPubKey, getInfo(ctx))
|
||||
if err != nil {
|
||||
log.Debugf("failed to open Management Service stream: %s", err)
|
||||
c.notifyDisconnected(err)
|
||||
|
||||
@@ -11,7 +11,7 @@ import (
|
||||
// MockClient is a mock implementation of the Client interface for testing.
|
||||
type MockClient struct {
|
||||
CloseFunc func() error
|
||||
SyncFunc func(ctx context.Context, sysInfo *system.Info, msgHandler func(msg *proto.SyncResponse) error) error
|
||||
SyncFunc func(ctx context.Context, getInfo func(ctx context.Context) *system.Info, msgHandler func(msg *proto.SyncResponse) error) error
|
||||
RegisterFunc func(setupKey string, jwtToken string, info *system.Info, sshKey []byte, dnsLabels domain.List) (*proto.LoginResponse, error)
|
||||
LoginFunc func(info *system.Info, sshKey []byte, dnsLabels domain.List) (*proto.LoginResponse, error)
|
||||
ExtendAuthSessionFunc func(info *system.Info, jwtToken string) (*proto.ExtendAuthSessionResponse, error)
|
||||
@@ -38,11 +38,11 @@ func (m *MockClient) Close() error {
|
||||
return m.CloseFunc()
|
||||
}
|
||||
|
||||
func (m *MockClient) Sync(ctx context.Context, sysInfo *system.Info, msgHandler func(msg *proto.SyncResponse) error) error {
|
||||
func (m *MockClient) Sync(ctx context.Context, getInfo func(ctx context.Context) *system.Info, msgHandler func(msg *proto.SyncResponse) error) error {
|
||||
if m.SyncFunc == nil {
|
||||
return nil
|
||||
}
|
||||
return m.SyncFunc(ctx, sysInfo, msgHandler)
|
||||
return m.SyncFunc(ctx, getInfo, msgHandler)
|
||||
}
|
||||
|
||||
func (m *MockClient) Job(ctx context.Context, msgHandler func(msg *proto.JobRequest) *proto.JobResponse) error {
|
||||
|
||||
Reference in New Issue
Block a user