netbird/management/server/api/v1alpha1/peers/peers_handler_test.go

package peers

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"net"
	"net/http"
	"net/http/httptest"
	"net/netip"
	"testing"
	"time"

	"github.com/gorilla/mux"
	"go.uber.org/mock/gomock"
	"golang.org/x/exp/maps"

	"github.com/netbirdio/netbird/management/internals/controllers/network_map"
	nbcontext "github.com/netbirdio/netbird/management/server/context"
	nbpeer "github.com/netbirdio/netbird/management/server/peer"
	"github.com/netbirdio/netbird/management/server/permissions"
	"github.com/netbirdio/netbird/management/server/permissions/modules"
	"github.com/netbirdio/netbird/management/server/permissions/operations"
	"github.com/netbirdio/netbird/management/server/types"
	"github.com/netbirdio/netbird/shared/auth"
	"github.com/netbirdio/netbird/shared/management/http/apiv1alpha1"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"

	"github.com/netbirdio/netbird/management/server/mock_server"
)

const (
	testPeerID                = "test_peer"
	noUpdateChannelTestPeerID = "no-update-channel"

	adminUser   = "admin_user"
	regularUser = "regular_user"
	serviceUser = "service_user"
)

func initTestMetaData(t *testing.T, peers ...*nbpeer.Peer) *Handler {

	peersMap := make(map[string]*nbpeer.Peer)
	for _, peer := range peers {
		peersMap[peer.ID] = peer.Copy()
	}

	policy := &types.Policy{
		ID:        "policy",
		AccountID: "test_id",
		Name:      "policy",
		Enabled:   true,
		Rules: []*types.PolicyRule{
			{
				ID:            "rule",
				Name:          "rule",
				Enabled:       true,
				Action:        "accept",
				Destinations:  []string{"group1"},
				Sources:       []string{"group1"},
				Bidirectional: true,
				Protocol:      "all",
				Ports:         []string{"80"},
			},
		},
	}

	srvUser := types.NewRegularUser(serviceUser, "", "")
	srvUser.IsServiceUser = true

	account := &types.Account{
		Id:     "test_id",
		Domain: "hotmail.com",
		Peers:  peersMap,
		Users: map[string]*types.User{
			adminUser:   types.NewAdminUser(adminUser),
			regularUser: types.NewRegularUser(regularUser, "", ""),
			serviceUser: srvUser,
		},
		Groups: map[string]*types.Group{
			"group1": {
				ID:        "group1",
				AccountID: "test_id",
				Name:      "group1",
				Issued:    "api",
				Peers:     maps.Keys(peersMap),
			},
		},
		Settings: &types.Settings{
			PeerLoginExpirationEnabled: true,
			PeerLoginExpiration:        time.Hour,
		},
		Policies: []*types.Policy{policy},
		Network: &types.Network{
			Identifier: "ciclqisab2ss43jdn8q0",
			Net: net.IPNet{
				IP:   net.ParseIP("100.67.0.0"),
				Mask: net.IPv4Mask(255, 255, 0, 0),
			},
			Serial: 51,
		},
	}

	ctrl := gomock.NewController(t)

	networkMapController := network_map.NewMockController(ctrl)
	networkMapController.EXPECT().
		GetDNSDomain(gomock.Any()).
		Return("domain").
		AnyTimes()

	ctrl2 := gomock.NewController(t)
	permissionsManager := permissions.NewMockManager(ctrl2)
	permissionsManager.EXPECT().ValidateAccountAccess(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(context.Background(), nil).AnyTimes()
	permissionsManager.EXPECT().
		ValidateUserPermissions(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Eq(modules.Peers), gomock.Eq(operations.Read)).
		DoAndReturn(func(ctx context.Context, accountID, userID string, module modules.Module, operation operations.Operation) (bool, context.Context, error) {
			user, ok := account.Users[userID]
			if !ok {
				return false, ctx, fmt.Errorf("user not found")
			}
			return user.HasAdminPower() || user.IsServiceUser, ctx, nil
		}).
		AnyTimes()

	return &Handler{
		accountManager: &mock_server.MockAccountManager{
			UpdatePeerFunc: func(_ context.Context, accountID, userID string, update *nbpeer.Peer) (*nbpeer.Peer, error) {
				var p *nbpeer.Peer
				for _, peer := range peers {
					if update.ID == peer.ID {
						p = peer.Copy()
						break
					}
				}
				p.SSHEnabled = update.SSHEnabled
				p.LoginExpirationEnabled = update.LoginExpirationEnabled
				p.Name = update.Name
				return p, nil
			},
			UpdatePeerIPFunc: func(_ context.Context, accountID, userID, peerID string, newIP netip.Addr) error {
				for _, peer := range peers {
					if peer.ID == peerID {
						peer.IP = newIP
						return nil
					}
				}
				return fmt.Errorf("peer not found")
			},
			GetPeerFunc: func(_ context.Context, accountID, peerID, userID string) (*nbpeer.Peer, error) {
				var p *nbpeer.Peer
				for _, peer := range peers {
					if peerID == peer.ID {
						p = peer.Copy()
						break
					}
				}
				return p, nil
			},
			GetUserByIDFunc: func(ctx context.Context, id string) (*types.User, error) {
				switch id {
				case adminUser:
					return account.Users[adminUser], nil
				case regularUser:
					return account.Users[regularUser], nil
				case serviceUser:
					return account.Users[serviceUser], nil
				default:
					return nil, fmt.Errorf("user not found")
				}
			},
			GetPeersFunc: func(_ context.Context, accountID, userID, nameFilter, ipFilter, macFilter string) ([]*nbpeer.Peer, error) {
				return peers, nil
			},
			GetPeerGroupsFunc: func(ctx context.Context, accountID, peerID string) ([]*types.Group, error) {
				peersID := make([]string, len(peers))
				for _, peer := range peers {
					peersID = append(peersID, peer.ID)
				}
				return []*types.Group{
					{
						ID:        "group1",
						AccountID: accountID,
						Name:      "group1",
						Issued:    "api",
						Peers:     peersID,
					},
				}, nil
			},
			GetDNSDomainFunc: func(settings *types.Settings) string {
				return "netbird.selfhosted"
			},
			GetAccountFunc: func(ctx context.Context, accountID string) (*types.Account, error) {
				return account, nil
			},
			GetAccountByIDFunc: func(ctx context.Context, accountID string, userID string) (*types.Account, error) {
				return account, nil
			},
			HasConnectedChannelFunc: func(peerID string) bool {
				statuses := make(map[string]struct{})
				for _, peer := range peers {
					if peer.ID == noUpdateChannelTestPeerID {
						break
					}
					statuses[peer.ID] = struct{}{}
				}
				_, ok := statuses[peerID]
				return ok
			},
			GetAccountSettingsFunc: func(ctx context.Context, accountID string, userID string) (*types.Settings, error) {
				return account.Settings, nil
			},
		},
		networkMapController: networkMapController,
		permissionsManager:   permissionsManager,
	}
}

// Tests the GetAllPeers endpoint reachable in the route /api/peers
// Use the metadata generated by initTestMetaData() to check for values
func TestGetPeers(t *testing.T) {

	peer := &nbpeer.Peer{
		ID:                     testPeerID,
		Key:                    "key",
		IP:                     netip.MustParseAddr("100.64.0.1"),
		IPv6:                   netip.MustParseAddr("fd00::1"),
		Status:                 &nbpeer.PeerStatus{Connected: true},
		Name:                   "PeerName",
		LoginExpirationEnabled: false,
		Meta: nbpeer.PeerSystemMeta{
			Hostname:           "hostname",
			GoOS:               "GoOS",
			Kernel:             "kernel",
			Core:               "core",
			Platform:           "platform",
			OS:                 "OS",
			WtVersion:          "development",
			SystemSerialNumber: "C02XJ0J0JGH7",
		},
	}

	peer1 := peer.Copy()
	peer1.ID = noUpdateChannelTestPeerID

	expectedUpdatedPeer := peer.Copy()
	expectedUpdatedPeer.LoginExpirationEnabled = true
	expectedUpdatedPeer.SSHEnabled = true
	expectedUpdatedPeer.Name = "New Name"

	expectedPeer1 := peer1.Copy()
	expectedPeer1.Status.Connected = false

	tt := []struct {
		name           string
		expectedStatus int
		requestType    string
		requestPath    string
		contentType    string
		requestBody    io.Reader
		expectedArray  bool
		expectedPeer   *nbpeer.Peer
	}{
		{
			name:           "GetPeersMetaData",
			requestType:    http.MethodGet,
			requestPath:    "/peers",
			expectedStatus: http.StatusOK,
			expectedArray:  true,
			expectedPeer:   peer,
		},
		{
			name:           "GetPeer with update channel",
			requestType:    http.MethodGet,
			requestPath:    "/peers/" + testPeerID,
			expectedStatus: http.StatusOK,
			expectedArray:  false,
			expectedPeer:   peer,
		},
		{
			name:           "PutPeer",
			requestType:    http.MethodPut,
			requestPath:    "/peers/" + testPeerID,
			contentType:    "application/json",
			expectedStatus: http.StatusOK,
			expectedArray:  false,
			requestBody:    bytes.NewBufferString("{\"login_expiration_enabled\":true,\"name\":\"New Name\",\"ssh_enabled\":true, \"inactivity_expiration_enabled\":true}"),
			expectedPeer:   expectedUpdatedPeer,
		},
	}

	p := initTestMetaData(t, peer, peer1)

	for _, tc := range tt {
		t.Run(tc.name, func(t *testing.T) {

			recorder := httptest.NewRecorder()
			req := httptest.NewRequest(tc.requestType, tc.requestPath, tc.requestBody)
			req = nbcontext.SetUserAuthInRequest(req, auth.UserAuth{
				UserId:    "admin_user",
				Domain:    "hotmail.com",
				AccountId: "test_id",
			})
			if tc.contentType != "" {
				req.Header.Set("Content-Type", tc.contentType)
			}

			v1validator, err := apiv1alpha1.CreateV1ApiValidatingMiddleware()
			assert.NoError(t, err)
			router := mux.NewRouter()
			router.Use(v1validator.Handler)

			router = p.WithEndpointsForRouter(router)
			router.ServeHTTP(recorder, req)

			res := recorder.Result()
			defer res.Body.Close()

			if status := recorder.Code; status != tc.expectedStatus {
				t.Fatalf("handler returned wrong status code: got %v want %v",
					status, http.StatusOK)
			}

			content, err := io.ReadAll(res.Body)
			if err != nil {
				t.Fatalf("I don't know what I expected; %v", err)
			}

			var got *apiv1alpha1.Peer
			if tc.expectedArray {
				respBody := []*apiv1alpha1.Peer{}
				err = json.Unmarshal(content, &respBody)
				if err != nil {
					t.Fatalf("Sent content is not in correct json format; %v", err)
				}

				// hardcode this check for now as we only have two peers in this suite
				assert.Equal(t, len(respBody), 2)

				for _, peer := range respBody {
					if peer.Id == testPeerID {
						got = peer
					}
				}

			} else {
				got = &apiv1alpha1.Peer{}
				err = json.Unmarshal(content, got)
				if err != nil {
					t.Fatalf("Sent content is not in correct json format; %v", err)
				}
			}

			t.Log(got)

			assert.Equal(t, tc.expectedPeer.Name, got.Name)
			assert.Equal(t, tc.expectedPeer.Meta.WtVersion, got.Version)
			assert.Equal(t, tc.expectedPeer.IP.String(), got.Ip)
			assert.Equal(t, "OS core", got.Os)
			assert.Equal(t, tc.expectedPeer.LoginExpirationEnabled, got.LoginExpirationEnabled)
			assert.Equal(t, tc.expectedPeer.SSHEnabled, got.SshEnabled)
			assert.Equal(t, tc.expectedPeer.Status.Connected, got.Connected)
			assert.Equal(t, tc.expectedPeer.Meta.SystemSerialNumber, got.SerialNumber)
		})
	}
}

func TestPeersHandlerUpdatePeerIP(t *testing.T) {
	testPeer := &nbpeer.Peer{
		ID:                     testPeerID,
		Key:                    "key",
		IP:                     netip.MustParseAddr("100.64.0.1"),
		IPv6:                   netip.MustParseAddr("fd00::1"),
		Status:                 &nbpeer.PeerStatus{Connected: false, LastSeen: time.Now()},
		Name:                   "test-host@netbird.io",
		LoginExpirationEnabled: false,
		UserID:                 regularUser,
		Meta: nbpeer.PeerSystemMeta{
			Hostname: "test-host@netbird.io",
			Core:     "22.04",
		},
	}

	p := initTestMetaData(t, testPeer)

	tt := []struct {
		name           string
		peerID         string
		requestBody    string
		callerUserID   string
		expectedStatus int
		expectedIP     string
	}{
		{
			name:           "update peer IP successfully",
			peerID:         testPeerID,
			requestBody:    `{"name":"test", "ssh_enabled":true, "login_expiration_enabled":true, "inactivity_expiration_enabled":true, "ip": "100.64.0.100"}`,
			callerUserID:   adminUser,
			expectedStatus: http.StatusOK,
			expectedIP:     "100.64.0.100",
		},
		{
			name:           "update peer IP with invalid IP",
			peerID:         testPeerID,
			requestBody:    `{"name":"test", "ssh_enabled":true, "login_expiration_enabled":true, "inactivity_expiration_enabled":true, "ip": "invalid-ip"}`,
			callerUserID:   adminUser,
			expectedStatus: http.StatusUnprocessableEntity,
		},
	}

	for _, tc := range tt {
		t.Run(tc.name, func(t *testing.T) {
			req := httptest.NewRequest(http.MethodPut, fmt.Sprintf("/peers/%s", tc.peerID), bytes.NewBuffer([]byte(tc.requestBody)))
			req.Header.Set("Content-Type", "application/json")
			req = nbcontext.SetUserAuthInRequest(req, auth.UserAuth{
				UserId:    tc.callerUserID,
				Domain:    "hotmail.com",
				AccountId: "test_id",
			})

			rr := httptest.NewRecorder()

			v1validator, err := apiv1alpha1.CreateV1ApiValidatingMiddleware()
			assert.NoError(t, err)
			router := mux.NewRouter()
			router.Use(v1validator.Handler)

			router = p.WithEndpointsForRouter(router)
			router.ServeHTTP(rr, req)

			assert.Equal(t, tc.expectedStatus, rr.Code)

			if tc.expectedStatus == http.StatusOK && tc.expectedIP != "" {
				var updatedPeer apiv1alpha1.Peer
				err := json.Unmarshal(rr.Body.Bytes(), &updatedPeer)
				require.NoError(t, err)
				assert.Equal(t, tc.expectedIP, updatedPeer.Ip)
			}
		})
	}
}