netbird/shared/management/client/client_test.go

package client

import (
	"context"
	"fmt"
	"net"
	"os"
	"sync"
	"sync/atomic"
	"testing"
	"time"

	log "github.com/sirupsen/logrus"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"go.uber.org/mock/gomock"
	"golang.zx2c4.com/wireguard/wgctrl/wgtypes"
	"google.golang.org/grpc"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/status"

	ephemeral_manager "github.com/netbirdio/netbird/management/internals/modules/peers/ephemeral/manager"
	"github.com/netbirdio/netbird/management/server/integrations/integrated_validator/validator"

	"github.com/netbirdio/netbird/management/internals/controllers/network_map/controller"
	"github.com/netbirdio/netbird/management/internals/controllers/network_map/update_channel"
	"github.com/netbirdio/netbird/management/internals/modules/peers"
	nbgrpc "github.com/netbirdio/netbird/management/internals/shared/grpc"
	"github.com/netbirdio/netbird/management/server/job"

	"github.com/netbirdio/netbird/client/system"
	"github.com/netbirdio/netbird/encryption"
	"github.com/netbirdio/netbird/management/internals/server/config"
	mgmt "github.com/netbirdio/netbird/management/server"
	"github.com/netbirdio/netbird/management/server/activity"
	nbcache "github.com/netbirdio/netbird/management/server/cache"
	"github.com/netbirdio/netbird/management/server/groups"
	"github.com/netbirdio/netbird/management/server/mock_server"
	"github.com/netbirdio/netbird/management/server/permissions"
	"github.com/netbirdio/netbird/management/server/settings"
	"github.com/netbirdio/netbird/management/server/store"
	"github.com/netbirdio/netbird/management/server/telemetry"
	"github.com/netbirdio/netbird/management/server/types"
	mgmtProto "github.com/netbirdio/netbird/shared/management/proto"
	"github.com/netbirdio/netbird/util"
)

const ValidKey = "A2C8E62B-38F5-4553-B31E-DD66C696CEBB"

func TestMain(m *testing.M) {
	_ = util.InitLog("debug", util.LogConsole)
	code := m.Run()
	os.Exit(code)
}

func startManagement(t *testing.T) (*grpc.Server, net.Listener) {
	t.Helper()
	level, _ := log.ParseLevel("debug")
	log.SetLevel(level)

	config := &config.Config{}
	_, err := util.ReadJson("../../../management/server/testdata/management.json", config)
	if err != nil {
		t.Fatal(err)
	}

	lis, err := net.Listen("tcp", ":0")
	if err != nil {
		t.Fatal(err)
	}
	s := grpc.NewServer()
	store, cleanUp, err := store.NewTestStoreFromSQL(context.Background(), "../../../management/server/testdata/store.sql", t.TempDir())
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(cleanUp)

	eventStore := &activity.InMemoryEventStore{}

	ctrl := gomock.NewController(t)
	t.Cleanup(ctrl.Finish)

	permissionsManagerMock := permissions.NewMockManager(ctrl)
	permissionsManagerMock.
		EXPECT().
		ValidateUserPermissions(
			gomock.Any(),
			gomock.Any(),
			gomock.Any(),
			gomock.Any(),
			gomock.Any(),
		).
		Return(true, context.Background(), nil).
		AnyTimes()

	peersManger := peers.NewManager(store, permissionsManagerMock)
	settingsManagerMock := settings.NewMockManager(ctrl)
	jobManager := job.NewJobManager(nil, store, peersManger)

	ctx := context.Background()

	cacheStore, err := nbcache.NewStore(ctx, 100*time.Millisecond, 300*time.Millisecond, 100)
	if err != nil {
		t.Fatal(err)
	}

	ia, _ := validator.NewIntegratedValidator(ctx, peersManger, settingsManagerMock, eventStore, cacheStore)

	metrics, err := telemetry.NewDefaultAppMetrics(ctx)
	require.NoError(t, err)

	settingsMockManager := settings.NewMockManager(ctrl)
	settingsMockManager.
		EXPECT().
		GetSettings(
			gomock.Any(),
			gomock.Any(),
			gomock.Any(),
		).
		Return(&types.Settings{}, nil).
		AnyTimes()
	settingsMockManager.
		EXPECT().
		GetExtraSettings(gomock.Any(), gomock.Any()).
		Return(&types.ExtraSettings{}, nil).
		AnyTimes()

	updateManager := update_channel.NewPeersUpdateManager(metrics)
	requestBuffer := mgmt.NewAccountRequestBuffer(ctx, store)
	networkMapController := controller.NewController(ctx, store, metrics, updateManager, requestBuffer, mgmt.MockIntegratedValidator{}, settingsMockManager, "netbird.selfhosted", ephemeral_manager.NewEphemeralManager(store, peersManger), config, nil)
	accountManager, err := mgmt.BuildManager(context.Background(), config, store, networkMapController, jobManager, nil, "", eventStore, nil, false, ia, metrics, settingsMockManager, permissionsManagerMock, false, cacheStore)
	if err != nil {
		t.Fatal(err)
	}

	groupsManager := groups.NewManagerMock()

	secretsManager, err := nbgrpc.NewTimeBasedAuthSecretsManager(updateManager, config.TURNConfig, config.Relay, settingsMockManager, groupsManager)
	if err != nil {
		t.Fatal(err)
	}
	mgmtServer, err := nbgrpc.NewServer(config, accountManager, settingsMockManager, jobManager, secretsManager, nil, nil, mgmt.MockIntegratedValidator{}, networkMapController, nil, nil)
	if err != nil {
		t.Fatal(err)
	}
	mgmtProto.RegisterManagementServiceServer(s, mgmtServer)
	go func() {
		if err := s.Serve(lis); err != nil {
			t.Error(err)
			return
		}
	}()

	return s, lis
}

func startMockManagement(t *testing.T) (*grpc.Server, net.Listener, *mock_server.ManagementServiceServerMock, wgtypes.Key) {
	t.Helper()
	lis, err := net.Listen("tcp", ":0")
	if err != nil {
		t.Fatal(err)
	}

	s := grpc.NewServer()

	serverKey, err := wgtypes.GenerateKey()
	if err != nil {
		t.Fatal(err)
	}

	mgmtMockServer := &mock_server.ManagementServiceServerMock{
		GetServerKeyFunc: func(context.Context, *mgmtProto.Empty) (*mgmtProto.ServerKeyResponse, error) {
			response := &mgmtProto.ServerKeyResponse{
				Key: serverKey.PublicKey().String(),
			}
			return response, nil
		},
	}

	mgmtProto.RegisterManagementServiceServer(s, mgmtMockServer)
	go func() {
		if err := s.Serve(lis); err != nil {
			t.Error(err)
			return
		}
	}()

	return s, lis, mgmtMockServer, serverKey
}

func closeManagementSilently(s *grpc.Server, listener net.Listener) {
	s.GracefulStop()
	err := listener.Close()
	if err != nil {
		log.Warnf("error while closing management listener %v", err)
		return
	}
}

func TestClient_HealthCheck(t *testing.T) {
	testKey, err := wgtypes.GenerateKey()
	if err != nil {
		t.Fatal(err)
	}
	ctx := context.Background()
	s, listener := startManagement(t)
	defer closeManagementSilently(s, listener)

	client, err := NewClient(ctx, listener.Addr().String(), testKey, false)
	if err != nil {
		t.Fatal(err)
	}

	if err := client.HealthCheck(); err != nil {
		t.Errorf("health check failed: %v", err)
	}
}

func TestClient_LoginUnregistered_ShouldThrow_401(t *testing.T) {
	testKey, err := wgtypes.GenerateKey()
	if err != nil {
		t.Fatal(err)
	}
	ctx := context.Background()
	s, listener := startManagement(t)
	defer closeManagementSilently(s, listener)

	client, err := NewClient(ctx, listener.Addr().String(), testKey, false)
	if err != nil {
		t.Fatal(err)
	}
	sysInfo := system.GetInfo(context.TODO())
	_, err = client.Login(sysInfo, nil, nil)
	if err == nil {
		t.Error("expecting err on unregistered login, got nil")
	}
	if s, ok := status.FromError(err); !ok || s.Code() != codes.PermissionDenied {
		t.Errorf("expecting err code %d denied on unregistered login got %d", codes.PermissionDenied, s.Code())
	}
}

func TestClient_LoginRegistered(t *testing.T) {
	testKey, err := wgtypes.GenerateKey()
	if err != nil {
		t.Fatal(err)
	}
	ctx := context.Background()
	s, listener := startManagement(t)
	defer closeManagementSilently(s, listener)

	client, err := NewClient(ctx, listener.Addr().String(), testKey, false)
	if err != nil {
		t.Fatal(err)
	}

	info := system.GetInfo(context.TODO())
	resp, err := client.Register(ValidKey, "", info, nil, nil)
	if err != nil {
		t.Error(err)
	}

	if resp == nil {
		t.Error("expecting non nil response, got nil")
	}
}

func TestClient_Sync(t *testing.T) {
	testKey, err := wgtypes.GenerateKey()
	if err != nil {
		t.Fatal(err)
	}
	ctx := context.Background()
	s, listener := startManagement(t)
	defer closeManagementSilently(s, listener)

	client, err := NewClient(ctx, listener.Addr().String(), testKey, false)
	if err != nil {
		t.Fatal(err)
	}

	info := system.GetInfo(context.TODO())
	_, err = client.Register(ValidKey, "", info, nil, nil)
	if err != nil {
		t.Error(err)
	}

	// create and register second peer (we should receive on Sync request)
	remoteKey, err := wgtypes.GenerateKey()
	if err != nil {
		t.Error(err)
	}
	remoteClient, err := NewClient(context.TODO(), listener.Addr().String(), remoteKey, false)
	if err != nil {
		t.Fatal(err)
	}

	info = system.GetInfo(context.TODO())
	_, err = remoteClient.Register(ValidKey, "", info, nil, nil)
	if err != nil {
		t.Fatal(err)
	}

	ch := make(chan *mgmtProto.SyncResponse, 1)

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	go func() {
		err = client.Sync(ctx, func(context.Context) *system.Info { return info }, func(msg *mgmtProto.SyncResponse) error {
			ch <- msg
			return nil
		})
		if err != nil {
			return
		}
	}()

	select {
	case resp := <-ch:
		if resp.GetPeerConfig() == nil && resp.GetNetworkMap().GetPeerConfig() == nil {
			t.Error("expecting non nil PeerConfig got nil")
		}
		if resp.GetNetbirdConfig() == nil {
			t.Error("expecting non nil NetbirdConfig got nil")
		}
		// Top-level RemotePeers is deprecated and must stay empty for
		// v0.29.3+ (and dev) clients — the field rides inside NetworkMap
		// (legacy) or the NetworkMapEnvelope (components) instead.
		if len(resp.GetRemotePeers()) != 0 {
			t.Error("expecting top-level RemotePeers to be empty for v0.29.3+ clients")
		}
		// Component-capable clients receive a NetworkMapEnvelope; the
		// remote-peers list is encoded inside it. Decode it and check the
		// envelope's peers slice. Legacy peers populate NetworkMap.RemotePeers;
		// both shapes must surface exactly one remote peer.
		remotePeerKeys := remotePeerKeysFromSync(resp, testKey.PublicKey().String())
		if len(remotePeerKeys) != 1 {
			t.Errorf("expecting RemotePeers size %d got %d", 1, len(remotePeerKeys))
			return
		}
		if resp.GetNetworkMap() != nil && resp.GetNetworkMap().GetRemotePeersIsEmpty() {
			t.Error("expecting RemotePeers property to be false, got true")
		}
		if remotePeerKeys[0] != remoteKey.PublicKey().String() {
			t.Errorf("expecting RemotePeer public key %s got %s", remoteKey.PublicKey().String(), remotePeerKeys[0])
		}
	case <-time.After(3 * time.Second):
		t.Error("timeout waiting for test to finish")
	}
}

// remotePeerKeysFromSync extracts the remote-peer WG keys from either the
// legacy NetworkMap.RemotePeers list or the components NetworkMapEnvelope's
// inner peers slice (filtering out the local receiving peer identified by
// localKey, since the envelope's peers list is index-addressed and includes
// the local peer alongside remotes).
func remotePeerKeysFromSync(resp *mgmtProto.SyncResponse, localKey string) []string {
	if rp := resp.GetRemotePeers(); len(rp) > 0 {
		out := make([]string, 0, len(rp))
		for _, p := range rp {
			out = append(out, p.GetWgPubKey())
		}
		return out
	}
	if rp := resp.GetNetworkMap().GetRemotePeers(); len(rp) > 0 {
		out := make([]string, 0, len(rp))
		for _, p := range rp {
			out = append(out, p.GetWgPubKey())
		}
		return out
	}
	env := resp.GetNetworkMapEnvelope().GetFull()
	if env == nil {
		return nil
	}
	out := make([]string, 0, len(env.GetPeers()))
	for _, p := range env.GetPeers() {
		key := wgKeyFromBytes(p.GetWgPubKey())
		if key == "" || key == localKey {
			continue
		}
		out = append(out, key)
	}
	return out
}

// wgKeyFromBytes mirrors the client-side decoder: the envelope ships raw 32
// bytes; reconstruct the standard base64 key the test compares against.
func wgKeyFromBytes(raw []byte) string {
	if len(raw) == 0 {
		return ""
	}
	var k wgtypes.Key
	if len(raw) != len(k) {
		return ""
	}
	copy(k[:], raw)
	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()

	testKey, err := wgtypes.GenerateKey()
	if err != nil {
		t.Fatal(err)
	}

	serverAddr := lis.Addr().String()
	ctx := context.Background()

	testClient, err := NewClient(ctx, serverAddr, testKey, false)
	if err != nil {
		t.Fatalf("error while creating testClient: %v", err)
	}

	var actualMeta *mgmtProto.PeerSystemMeta
	var actualValidKey string
	var wg sync.WaitGroup
	wg.Add(1)

	mgmtMockServer.LoginFunc = func(ctx context.Context, msg *mgmtProto.EncryptedMessage) (*mgmtProto.EncryptedMessage, error) {
		peerKey, err := wgtypes.ParseKey(msg.GetWgPubKey())
		if err != nil {
			log.Warnf("error while parsing peer's Wireguard public key %s on Sync request.", msg.WgPubKey)
			return nil, status.Errorf(codes.InvalidArgument, "provided wgPubKey %s is invalid", msg.WgPubKey)
		}

		loginReq := &mgmtProto.LoginRequest{}
		err = encryption.DecryptMessage(peerKey, serverKey, msg.Body, loginReq)
		if err != nil {
			log.Fatal(err)
		}

		actualMeta = loginReq.GetMeta()
		actualValidKey = loginReq.GetSetupKey()
		wg.Done()

		loginResp := &mgmtProto.LoginResponse{}
		encryptedResp, err := encryption.EncryptMessage(peerKey, serverKey, loginResp)
		if err != nil {
			return nil, err
		}

		return &mgmtProto.EncryptedMessage{
			WgPubKey: serverKey.PublicKey().String(),
			Body:     encryptedResp,
			Version:  0,
		}, nil
	}

	info := system.GetInfo(context.TODO())
	_, err = testClient.Register(ValidKey, "", info, nil, nil)
	if err != nil {
		t.Errorf("error while trying to register client: %v", err)
	}

	wg.Wait()

	protoNetAddr := make([]*mgmtProto.NetworkAddress, 0, len(info.NetworkAddresses))
	for _, addr := range info.NetworkAddresses {
		protoNetAddr = append(protoNetAddr, &mgmtProto.NetworkAddress{
			NetIP: addr.NetIP.String(),
			Mac:   addr.Mac,
		})

	}

	expectedMeta := &mgmtProto.PeerSystemMeta{
		Hostname:       info.Hostname,
		GoOS:           info.GoOS,
		Kernel:         info.Kernel,
		Platform:       info.Platform,
		OS:             info.OS,
		Core:           info.OSVersion,
		OSVersion:      info.OSVersion,
		NetbirdVersion: info.NetbirdVersion,
		KernelVersion:  info.KernelVersion,

		NetworkAddresses: protoNetAddr,
		SysSerialNumber:  info.SystemSerialNumber,
		SysProductName:   info.SystemProductName,
		SysManufacturer:  info.SystemManufacturer,
		Environment:      &mgmtProto.Environment{Cloud: info.Environment.Cloud, Platform: info.Environment.Platform},
	}

	assert.Equal(t, ValidKey, actualValidKey)
	if !isEqual(expectedMeta, actualMeta) {
		t.Errorf("expected and actual meta are not equal")
	}
}

func isEqual(a, b *mgmtProto.PeerSystemMeta) bool {
	if len(a.NetworkAddresses) != len(b.NetworkAddresses) {
		return false
	}

	for _, addr := range a.GetNetworkAddresses() {
		var found bool
		for _, oAddr := range b.GetNetworkAddresses() {
			if addr.GetMac() == oAddr.GetMac() && addr.GetNetIP() == oAddr.GetNetIP() {
				found = true
				continue
			}
		}
		if !found {
			return false
		}
	}

	log.Infof("------")

	return a.GetHostname() == b.GetHostname() &&
		a.GetGoOS() == b.GetGoOS() &&
		a.GetKernel() == b.GetKernel() &&
		a.GetKernelVersion() == b.GetKernelVersion() &&
		a.GetCore() == b.GetCore() &&
		a.GetPlatform() == b.GetPlatform() &&
		a.GetOS() == b.GetOS() &&
		a.GetOSVersion() == b.GetOSVersion() &&
		a.GetNetbirdVersion() == b.GetNetbirdVersion() &&
		a.GetUiVersion() == b.GetUiVersion() &&
		a.GetSysSerialNumber() == b.GetSysSerialNumber() &&
		a.GetSysProductName() == b.GetSysProductName() &&
		a.GetSysManufacturer() == b.GetSysManufacturer() &&
		a.GetEnvironment().Cloud == b.GetEnvironment().Cloud &&
		a.GetEnvironment().Platform == b.GetEnvironment().Platform
}

func Test_GetDeviceAuthorizationFlow(t *testing.T) {
	s, lis, mgmtMockServer, serverKey := startMockManagement(t)
	defer s.GracefulStop()

	testKey, err := wgtypes.GenerateKey()
	if err != nil {
		t.Fatal(err)
	}

	serverAddr := lis.Addr().String()
	ctx := context.Background()

	client, err := NewClient(ctx, serverAddr, testKey, false)
	if err != nil {
		t.Fatalf("error while creating testClient: %v", err)
	}

	expectedFlowInfo := &mgmtProto.DeviceAuthorizationFlow{
		Provider:       0,
		ProviderConfig: &mgmtProto.ProviderConfig{ClientID: "client"},
	}

	mgmtMockServer.GetDeviceAuthorizationFlowFunc = func(ctx context.Context, req *mgmtProto.EncryptedMessage) (*mgmtProto.EncryptedMessage, error) {
		encryptedResp, err := encryption.EncryptMessage(client.key.PublicKey(), serverKey, expectedFlowInfo)
		if err != nil {
			return nil, err
		}

		return &mgmtProto.EncryptedMessage{
			WgPubKey: serverKey.PublicKey().String(),
			Body:     encryptedResp,
			Version:  0,
		}, nil
	}

	flowInfo, err := client.GetDeviceAuthorizationFlow()
	if err != nil {
		t.Error("error while retrieving device auth flow information")
	}

	assert.Equal(t, expectedFlowInfo.Provider, flowInfo.Provider, "provider should match")
	assert.Equal(t, expectedFlowInfo.ProviderConfig.ClientID, flowInfo.ProviderConfig.ClientID, "provider configured client ID should match")
}

func Test_GetPKCEAuthorizationFlow(t *testing.T) {
	s, lis, mgmtMockServer, serverKey := startMockManagement(t)
	defer s.GracefulStop()

	testKey, err := wgtypes.GenerateKey()
	if err != nil {
		t.Fatal(err)
	}

	serverAddr := lis.Addr().String()
	ctx := context.Background()

	client, err := NewClient(ctx, serverAddr, testKey, false)
	if err != nil {
		t.Fatalf("error while creating testClient: %v", err)
	}

	expectedFlowInfo := &mgmtProto.PKCEAuthorizationFlow{
		ProviderConfig: &mgmtProto.ProviderConfig{
			ClientID:     "client",
			ClientSecret: "secret", //nolint:staticcheck
		},
	}

	var gotRequest mgmtProto.PKCEAuthorizationFlowRequest
	mgmtMockServer.GetPKCEAuthorizationFlowFunc = func(ctx context.Context, req *mgmtProto.EncryptedMessage) (*mgmtProto.EncryptedMessage, error) {
		if err := encryption.DecryptMessage(client.key.PublicKey(), serverKey, req.Body, &gotRequest); err != nil {
			return nil, err
		}

		encryptedResp, err := encryption.EncryptMessage(client.key.PublicKey(), serverKey, expectedFlowInfo)
		if err != nil {
			return nil, err
		}

		return &mgmtProto.EncryptedMessage{
			WgPubKey: serverKey.PublicKey().String(),
			Body:     encryptedResp,
			Version:  0,
		}, nil
	}

	flowInfo, err := client.GetPKCEAuthorizationFlow(true)
	require.NoError(t, err)

	assert.True(t, gotRequest.GetSessionExtend(), "session extend should reach the server")

	assert.Equal(t, expectedFlowInfo.ProviderConfig.ClientID, flowInfo.ProviderConfig.ClientID, "provider configured client ID should match")
	assert.Equal(t, expectedFlowInfo.ProviderConfig.ClientSecret, flowInfo.ProviderConfig.ClientSecret, "provider configured client secret should match") //nolint:staticcheck
}