Store updated system info on Login to Management (#323)

client/internal/connect.go

@@ -2,6 +2,7 @@ package internal
 
 import (
 	"context"
+	"github.com/netbirdio/netbird/client/system"
 	"time"
 
 	"github.com/netbirdio/netbird/iface"
@@ -193,7 +194,8 @@ func connectToManagement(ctx context.Context, managementAddr string, ourPrivateK
 		return nil, nil, status.Errorf(codes.FailedPrecondition, "failed while getting Management Service public key: %s", err)
 	}
 
-	loginResp, err := client.Login(*serverPublicKey)
+	sysInfo := system.GetInfo()
+	loginResp, err := client.Login(*serverPublicKey, sysInfo)
 	if err != nil {
 		if s, ok := status.FromError(err); ok && s.Code() == codes.PermissionDenied {
 			log.Error("peer registration required. Please run wiretrustee login command first")

client/internal/login.go

@@ -56,7 +56,8 @@ func Login(ctx context.Context, config *Config, setupKey string, jwtToken string
 
 // loginPeer attempts to login to Management Service. If peer wasn't registered, tries the registration flow.
 func loginPeer(serverPublicKey wgtypes.Key, client *mgm.GrpcClient, setupKey string, jwtToken string) (*mgmProto.LoginResponse, error) {
-	loginResp, err := client.Login(serverPublicKey)
+	sysInfo := system.GetInfo()
+	loginResp, err := client.Login(serverPublicKey, sysInfo)
 	if err != nil {
 		if s, ok := status.FromError(err); ok && s.Code() == codes.PermissionDenied {
 			log.Debugf("peer registration required")

management/client/client.go

@@ -12,7 +12,7 @@ type Client interface {
 	io.Closer
 	Sync(msgHandler func(msg *proto.SyncResponse) error) error
 	GetServerPublicKey() (*wgtypes.Key, error)
-	Register(serverKey wgtypes.Key, setupKey string, jwtToken string, info *system.Info) (*proto.LoginResponse, error)
+	Register(serverKey wgtypes.Key, setupKey string, jwtToken string, sysInfo *system.Info) (*proto.LoginResponse, error)
-	Login(serverKey wgtypes.Key) (*proto.LoginResponse, error)
+	Login(serverKey wgtypes.Key, sysInfo *system.Info) (*proto.LoginResponse, error)
 	GetDeviceAuthorizationFlow(serverKey wgtypes.Key) (*proto.DeviceAuthorizationFlow, error)
 }

management/client/client_test.go

@@ -157,7 +157,8 @@ func TestClient_LoginUnregistered_ShouldThrow_401(t *testing.T) {
 	if err != nil {
 		t.Fatal(err)
 	}
-	_, err = client.Login(*key)
+	sysInfo := system.GetInfo()
+	_, err = client.Login(*key, sysInfo)
 	if err == nil {
 		t.Error("expecting err on unregistered login, got nil")
 	}

management/client/grpc.go

@@ -228,22 +228,13 @@ func (c *GrpcClient) login(serverKey wgtypes.Key, req *proto.LoginRequest) (*pro
 // Register registers peer on Management Server. It actually calls a Login endpoint with a provided setup key
 // Takes care of encrypting and decrypting messages.
 // This method will also collect system info and send it with the request (e.g. hostname, os, etc)
-func (c *GrpcClient) Register(serverKey wgtypes.Key, setupKey string, jwtToken string, info *system.Info) (*proto.LoginResponse, error) {
+func (c *GrpcClient) Register(serverKey wgtypes.Key, setupKey string, jwtToken string, sysInfo *system.Info) (*proto.LoginResponse, error) {
-	meta := &proto.PeerSystemMeta{
+	return c.login(serverKey, &proto.LoginRequest{SetupKey: setupKey, Meta: infoToMetaData(sysInfo), JwtToken: jwtToken})
-		Hostname:           info.Hostname,
-		GoOS:               info.GoOS,
-		OS:                 info.OS,
-		Core:               info.OSVersion,
-		Platform:           info.Platform,
-		Kernel:             info.Kernel,
-		WiretrusteeVersion: info.WiretrusteeVersion,
-	}
-	return c.login(serverKey, &proto.LoginRequest{SetupKey: setupKey, Meta: meta, JwtToken: jwtToken})
 }
 
 // Login attempts login to Management Server. Takes care of encrypting and decrypting messages.
-func (c *GrpcClient) Login(serverKey wgtypes.Key) (*proto.LoginResponse, error) {
+func (c *GrpcClient) Login(serverKey wgtypes.Key, sysInfo *system.Info) (*proto.LoginResponse, error) {
-	return c.login(serverKey, &proto.LoginRequest{})
+	return c.login(serverKey, &proto.LoginRequest{Meta: infoToMetaData(sysInfo)})
 }
 
 // GetDeviceAuthorizationFlow returns a device authorization flow information.
@@ -279,3 +270,18 @@ func (c *GrpcClient) GetDeviceAuthorizationFlow(serverKey wgtypes.Key) (*proto.D
 
 	return flowInfoResp, nil
 }
+
+func infoToMetaData(info *system.Info) *proto.PeerSystemMeta {
+	if info == nil {
+		return nil
+	}
+	return &proto.PeerSystemMeta{
+		Hostname:           info.Hostname,
+		GoOS:               info.GoOS,
+		OS:                 info.OS,
+		Core:               info.OSVersion,
+		Platform:           info.Platform,
+		Kernel:             info.Kernel,
+		WiretrusteeVersion: info.WiretrusteeVersion,
+	}
+}

management/client/mock.go

@@ -11,7 +11,7 @@ type MockClient struct {
 	SyncFunc                       func(msgHandler func(msg *proto.SyncResponse) error) error
 	GetServerPublicKeyFunc         func() (*wgtypes.Key, error)
 	RegisterFunc                   func(serverKey wgtypes.Key, setupKey string, jwtToken string, info *system.Info) (*proto.LoginResponse, error)
-	LoginFunc                      func(serverKey wgtypes.Key) (*proto.LoginResponse, error)
+	LoginFunc                      func(serverKey wgtypes.Key, info *system.Info) (*proto.LoginResponse, error)
 	GetDeviceAuthorizationFlowFunc func(serverKey wgtypes.Key) (*proto.DeviceAuthorizationFlow, error)
 }
 
@@ -43,11 +43,11 @@ func (m *MockClient) Register(serverKey wgtypes.Key, setupKey string, jwtToken s
 	return m.RegisterFunc(serverKey, setupKey, jwtToken, info)
 }
 
-func (m *MockClient) Login(serverKey wgtypes.Key) (*proto.LoginResponse, error) {
+func (m *MockClient) Login(serverKey wgtypes.Key, info *system.Info) (*proto.LoginResponse, error) {
 	if m.LoginFunc == nil {
 		return nil, nil
 	}
-	return m.LoginFunc(serverKey)
+	return m.LoginFunc(serverKey, info)
 }
 
 func (m *MockClient) GetDeviceAuthorizationFlow(serverKey wgtypes.Key) (*proto.DeviceAuthorizationFlow, error) {

management/server/account.go

@@ -44,6 +44,7 @@ type AccountManager interface {
 	GetPeerByIP(accountId string, peerIP string) (*Peer, error)
 	GetNetworkMap(peerKey string) (*NetworkMap, error)
 	AddPeer(setupKey string, userId string, peer *Peer) (*Peer, error)
+	UpdatePeerMeta(peerKey string, meta PeerSystemMeta) error
 	GetUsersFromAccount(accountId string) ([]*UserInfo, error)
 	GetGroup(accountId, groupID string) (*Group, error)
 	SaveGroup(accountId string, group *Group) error

management/server/account_test.go

@@ -602,6 +602,76 @@ func TestGetUsersFromAccount(t *testing.T) {
 	}
 }
 
+func TestAccountManager_UpdatePeerMeta(t *testing.T) {
+	manager, err := createManager(t)
+	if err != nil {
+		t.Fatal(err)
+		return
+	}
+
+	account, err := manager.AddAccount("test_account", "account_creator", "")
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	var setupKey *SetupKey
+	for _, key := range account.SetupKeys {
+		setupKey = key
+	}
+
+	key, err := wgtypes.GeneratePrivateKey()
+	if err != nil {
+		t.Fatal(err)
+		return
+	}
+
+	peer, err := manager.AddPeer(setupKey.Key, "", &Peer{
+		Key: key.PublicKey().String(),
+		Meta: PeerSystemMeta{
+			Hostname:  "Hostname",
+			GoOS:      "GoOS",
+			Kernel:    "Kernel",
+			Core:      "Core",
+			Platform:  "Platform",
+			OS:        "OS",
+			WtVersion: "WtVersion",
+		},
+		Name: key.PublicKey().String(),
+	})
+	if err != nil {
+		t.Errorf("expecting peer to be added, got failure %v", err)
+		return
+	}
+
+	newMeta := PeerSystemMeta{
+		Hostname:  "new-Hostname",
+		GoOS:      "new-GoOS",
+		Kernel:    "new-Kernel",
+		Core:      "new-Core",
+		Platform:  "new-Platform",
+		OS:        "new-OS",
+		WtVersion: "new-WtVersion",
+	}
+	err = manager.UpdatePeerMeta(peer.Key, newMeta)
+	if err != nil {
+		t.Error(err)
+		return
+	}
+
+	p, err := manager.GetPeer(peer.Key)
+	if err != nil {
+		return
+	}
+
+	if err != nil {
+		t.Fatal(err)
+		return
+	}
+
+	assert.Equal(t, newMeta, p.Meta)
+
+}
+
 func createManager(t *testing.T) (*DefaultAccountManager, error) {
 	store, err := createStore(t)
 	if err != nil {

management/server/grpcserver.go

@@ -246,15 +246,16 @@ func (s *Server) Login(ctx context.Context, req *proto.EncryptedMessage) (*proto
 		return nil, status.Errorf(codes.InvalidArgument, "provided wgPubKey %s is invalid", req.WgPubKey)
 	}
 
-	peer, err := s.accountManager.GetPeer(peerKey.String())
-	if err != nil {
-		if errStatus, ok := status.FromError(err); ok && errStatus.Code() == codes.NotFound {
-			// peer doesn't exist -> check if setup key was provided
 	loginReq := &proto.LoginRequest{}
 	err = encryption.DecryptMessage(peerKey, s.wgKey, req.Body, loginReq)
 	if err != nil {
 		return nil, status.Errorf(codes.InvalidArgument, "invalid request message")
 	}
+
+	peer, err := s.accountManager.GetPeer(peerKey.String())
+	if err != nil {
+		if errStatus, ok := status.FromError(err); ok && errStatus.Code() == codes.NotFound {
+			// peer doesn't exist -> check if setup key was provided
 			if loginReq.GetJwtToken() == "" && loginReq.GetSetupKey() == "" {
 				// absent setup key -> permission denied
 				return nil, status.Errorf(codes.PermissionDenied, "provided peer with the key wgPubKey %s is not registered and no setup key or jwt was provided", peerKey.String())
@@ -269,8 +270,21 @@ func (s *Server) Login(ctx context.Context, req *proto.EncryptedMessage) (*proto
 		} else {
 			return nil, status.Error(codes.Internal, "internal server error")
 		}
+	} else if loginReq.GetMeta() != nil {
+		// update peer's system meta data on Login
+		err = s.accountManager.UpdatePeerMeta(peerKey.String(), PeerSystemMeta{
+			Hostname:  loginReq.GetMeta().GetHostname(),
+			GoOS:      loginReq.GetMeta().GetGoOS(),
+			Kernel:    loginReq.GetMeta().GetKernel(),
+			Core:      loginReq.GetMeta().GetCore(),
+			Platform:  loginReq.GetMeta().GetPlatform(),
+			OS:        loginReq.GetMeta().GetOS(),
+			WtVersion: loginReq.GetMeta().GetWiretrusteeVersion()})
+		if err != nil {
+			log.Errorf("failed updating peer system meta data %s", peerKey.String())
+			return nil, status.Error(codes.Internal, "internal server error")
+		}
 	}
-
 	// if peer has reached this point then it has logged in
 	loginResp := &proto.LoginResponse{
 		WiretrusteeConfig: toWiretrusteeConfig(s.config, nil),

management/server/mock_server/account_mock.go

@@ -38,6 +38,7 @@ type MockAccountManager struct {
 	DeleteRuleFunc                        func(accountID, ruleID string) error
 	ListRulesFunc                         func(accountID string) ([]*server.Rule, error)
 	GetUsersFromAccountFunc               func(accountID string) ([]*server.UserInfo, error)
+	UpdatePeerMetaFunc                    func(peerKey string, meta server.PeerSystemMeta) error
 }
 
 func (am *MockAccountManager) GetUsersFromAccount(accountID string) ([]*server.UserInfo, error) {
@@ -275,3 +276,10 @@ func (am *MockAccountManager) ListRules(accountID string) ([]*server.Rule, error
 	}
 	return nil, status.Errorf(codes.Unimplemented, "method ListRules not implemented")
 }
+
+func (am *MockAccountManager) UpdatePeerMeta(peerKey string, meta server.PeerSystemMeta) error {
+	if am.UpdatePeerMetaFunc != nil {
+		return am.UpdatePeerMetaFunc(peerKey, meta)
+	}
+	return status.Errorf(codes.Unimplemented, "method UpdatePeerMetaFunc not implemented")
+}

management/server/peer.go

@@ -376,3 +376,27 @@ func (am *DefaultAccountManager) AddPeer(
 
 	return newPeer, nil
 }
+
+// UpdatePeerMeta updates peer's system metadata
+func (am *DefaultAccountManager) UpdatePeerMeta(peerKey string, meta PeerSystemMeta) error {
+	am.mux.Lock()
+	defer am.mux.Unlock()
+
+	peer, err := am.Store.GetPeer(peerKey)
+	if err != nil {
+		return err
+	}
+
+	account, err := am.Store.GetPeerAccount(peerKey)
+	if err != nil {
+		return err
+	}
+
+	peerCopy := peer.Copy()
+	peerCopy.Meta = meta
+	err = am.Store.SavePeer(account.Id, peerCopy)
+	if err != nil {
+		return err
+	}
+	return nil
+}