Skip to content

Commit

Permalink
Update peer methods to return posture checks
Browse files Browse the repository at this point in the history
  • Loading branch information
@bcmmbaga
bcmmbaga committed Jun 22, 2024
1 parent b45d742 commit 2b0f93e
Show file tree
Hide file tree
Showing 5 changed files with 34 additions and 43 deletions.
16 changes: 8 additions & 8 deletions management/server/account_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -85,7 +85,7 @@ func verifyCanAddPeerToAccount(t *testing.T, manager AccountManager, account *Ac
setupKey = key.Key
}

_, _, err := manager.AddPeer(setupKey, userID, peer)
_, _, _, err := manager.AddPeer(setupKey, userID, peer)
if err != nil {
t.Error("expected to add new peer successfully after creating new account, but failed", err)
}
Expand Down Expand Up @@ -997,7 +997,7 @@ func TestAccountManager_AddPeer(t *testing.T) {
expectedPeerKey := key.PublicKey().String()
expectedSetupKey := setupKey.Key

peer, _, err := manager.AddPeer(setupKey.Key, "", &nbpeer.Peer{
peer, _, _, err := manager.AddPeer(setupKey.Key, "", &nbpeer.Peer{
Key: expectedPeerKey,
Meta: nbpeer.PeerSystemMeta{Hostname: expectedPeerKey},
})
Expand Down Expand Up @@ -1065,7 +1065,7 @@ func TestAccountManager_AddPeerWithUserID(t *testing.T) {
expectedPeerKey := key.PublicKey().String()
expectedUserID := userID

peer, _, err := manager.AddPeer("", userID, &nbpeer.Peer{
peer, _, _, err := manager.AddPeer("", userID, &nbpeer.Peer{
Key: expectedPeerKey,
Meta: nbpeer.PeerSystemMeta{Hostname: expectedPeerKey},
})
Expand Down Expand Up @@ -1140,7 +1140,7 @@ func TestAccountManager_NetworkUpdates(t *testing.T) {
}
expectedPeerKey := key.PublicKey().String()

peer, _, err := manager.AddPeer(setupKey.Key, "", &nbpeer.Peer{
peer, _, _, err := manager.AddPeer(setupKey.Key, "", &nbpeer.Peer{
Key: expectedPeerKey,
Meta: nbpeer.PeerSystemMeta{Hostname: expectedPeerKey},
})
Expand Down Expand Up @@ -1315,7 +1315,7 @@ func TestAccountManager_DeletePeer(t *testing.T) {

peerKey := key.PublicKey().String()

peer, _, err := manager.AddPeer(setupKey.Key, "", &nbpeer.Peer{
peer, _, _, err := manager.AddPeer(setupKey.Key, "", &nbpeer.Peer{
Key: peerKey,
Meta: nbpeer.PeerSystemMeta{Hostname: peerKey},
})
Expand Down Expand Up @@ -1662,7 +1662,7 @@ func TestDefaultAccountManager_UpdatePeer_PeerLoginExpiration(t *testing.T) {

key, err := wgtypes.GenerateKey()
require.NoError(t, err, "unable to generate WireGuard key")
peer, _, err := manager.AddPeer("", userID, &nbpeer.Peer{
peer, _, _, err := manager.AddPeer("", userID, &nbpeer.Peer{
Key: key.PublicKey().String(),
Meta: nbpeer.PeerSystemMeta{Hostname: "test-peer"},
LoginExpirationEnabled: true,
Expand Down Expand Up @@ -1715,7 +1715,7 @@ func TestDefaultAccountManager_MarkPeerConnected_PeerLoginExpiration(t *testing.

key, err := wgtypes.GenerateKey()
require.NoError(t, err, "unable to generate WireGuard key")
_, _, err = manager.AddPeer("", userID, &nbpeer.Peer{
_, _, _, err = manager.AddPeer("", userID, &nbpeer.Peer{
Key: key.PublicKey().String(),
Meta: nbpeer.PeerSystemMeta{Hostname: "test-peer"},
LoginExpirationEnabled: true,
Expand Down Expand Up @@ -1759,7 +1759,7 @@ func TestDefaultAccountManager_UpdateAccountSettings_PeerLoginExpiration(t *test

key, err := wgtypes.GenerateKey()
require.NoError(t, err, "unable to generate WireGuard key")
_, _, err = manager.AddPeer("", userID, &nbpeer.Peer{
_, _, _, err = manager.AddPeer("", userID, &nbpeer.Peer{
Key: key.PublicKey().String(),
Meta: nbpeer.PeerSystemMeta{Hostname: "test-peer"},
LoginExpirationEnabled: true,
Expand Down
4 changes: 2 additions & 2 deletions management/server/dns_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -256,11 +256,11 @@ func initTestDNSAccount(t *testing.T, am *DefaultAccountManager) (*Account, erro
return nil, err
}

savedPeer1, _, err := am.AddPeer("", dnsAdminUserID, peer1)
savedPeer1, _, _, err := am.AddPeer("", dnsAdminUserID, peer1)
if err != nil {
return nil, err
}
_, _, err = am.AddPeer("", dnsAdminUserID, peer2)
_, _, _, err = am.AddPeer("", dnsAdminUserID, peer2)
if err != nil {
return nil, err
}
Expand Down
33 changes: 12 additions & 21 deletions management/server/mock_server/account_mock.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -30,11 +30,11 @@ type MockAccountManager struct {
ListUsersFunc func(accountID string) ([]*server.User, error)
GetPeersFunc func(accountID, userID string) ([]*nbpeer.Peer, error)
MarkPeerConnectedFunc func(peerKey string, connected bool, realIP net.IP) error
SyncAndMarkPeerFunc func(peerPubKey string, meta nbpeer.PeerSystemMeta, realIP net.IP) (*nbpeer.Peer, *server.NetworkMap, error)
SyncAndMarkPeerFunc func(peerPubKey string, meta nbpeer.PeerSystemMeta, realIP net.IP) (*nbpeer.Peer, *server.NetworkMap, []*posture.Checks, error)
DeletePeerFunc func(accountID, peerKey, userID string) error
GetNetworkMapFunc func(peerKey string) (*server.NetworkMap, error)
GetPeerNetworkFunc func(peerKey string) (*server.Network, error)
AddPeerFunc func(setupKey string, userId string, peer *nbpeer.Peer) (*nbpeer.Peer, *server.NetworkMap, error)
AddPeerFunc func(setupKey string, userId string, peer *nbpeer.Peer) (*nbpeer.Peer, *server.NetworkMap, []*posture.Checks, error)
GetGroupFunc func(accountID, groupID, userID string) (*group.Group, error)
GetAllGroupsFunc func(accountID, userID string) ([]*group.Group, error)
GetGroupByNameFunc func(accountID, groupName string) (*group.Group, error)
Expand Down Expand Up @@ -83,10 +83,9 @@ type MockAccountManager struct {
GetDNSSettingsFunc func(accountID, userID string) (*server.DNSSettings, error)
SaveDNSSettingsFunc func(accountID, userID string, dnsSettingsToSave *server.DNSSettings) error
GetPeerFunc func(accountID, peerID, userID string) (*nbpeer.Peer, error)
GetPeerAppliedPostureChecksFunc func(peerKey string) ([]posture.Checks, error)
UpdateAccountSettingsFunc func(accountID, userID string, newSettings *server.Settings) (*server.Account, error)
LoginPeerFunc func(login server.PeerLogin) (*nbpeer.Peer, *server.NetworkMap, error)
SyncPeerFunc func(sync server.PeerSync, account *server.Account) (*nbpeer.Peer, *server.NetworkMap, error)
LoginPeerFunc func(login server.PeerLogin) (*nbpeer.Peer, *server.NetworkMap, []*posture.Checks, error)
SyncPeerFunc func(sync server.PeerSync, account *server.Account) (*nbpeer.Peer, *server.NetworkMap, []*posture.Checks, error)
InviteUserFunc func(accountID string, initiatorUserID string, targetUserEmail string) error
GetAllConnectedPeersFunc func() (map[string]struct{}, error)
HasConnectedChannelFunc func(peerID string) bool
Expand All @@ -102,11 +101,11 @@ type MockAccountManager struct {
FindExistingPostureCheckFunc func(accountID string, checks *posture.ChecksDefinition) (*posture.Checks, error)
}

func (am *MockAccountManager) SyncAndMarkPeer(peerPubKey string, meta nbpeer.PeerSystemMeta, realIP net.IP) (*nbpeer.Peer, *server.NetworkMap, error) {
func (am *MockAccountManager) SyncAndMarkPeer(peerPubKey string, meta nbpeer.PeerSystemMeta, realIP net.IP) (*nbpeer.Peer, *server.NetworkMap, []*posture.Checks, error) {
if am.SyncAndMarkPeerFunc != nil {
return am.SyncAndMarkPeerFunc(peerPubKey, meta, realIP)
}
return nil, nil, status.Errorf(codes.Unimplemented, "method MarkPeerConnected is not implemented")
return nil, nil, nil, status.Errorf(codes.Unimplemented, "method MarkPeerConnected is not implemented")
}

func (am *MockAccountManager) CancelPeerRoutines(peer *nbpeer.Peer) error {
Expand Down Expand Up @@ -282,11 +281,11 @@ func (am *MockAccountManager) AddPeer(
setupKey string,
userId string,
peer *nbpeer.Peer,
) (*nbpeer.Peer, *server.NetworkMap, error) {
) (*nbpeer.Peer, *server.NetworkMap, []*posture.Checks, error) {
if am.AddPeerFunc != nil {
return am.AddPeerFunc(setupKey, userId, peer)
}
return nil, nil, status.Errorf(codes.Unimplemented, "method AddPeer is not implemented")
return nil, nil, nil, status.Errorf(codes.Unimplemented, "method AddPeer is not implemented")
}

// GetGroupByName mock implementation of GetGroupByName from server.AccountManager interface
Expand Down Expand Up @@ -627,14 +626,6 @@ func (am *MockAccountManager) GetPeer(accountID, peerID, userID string) (*nbpeer
return nil, status.Errorf(codes.Unimplemented, "method GetPeer is not implemented")
}

// GetPeerAppliedPostureChecks mocks GetPeerAppliedPostureChecks of the AccountManager interface
func (am *MockAccountManager) GetPeerAppliedPostureChecks(peerKey string) ([]posture.Checks, error) {
if am.GetPeerAppliedPostureChecksFunc != nil {
return am.GetPeerAppliedPostureChecksFunc(peerKey)
}
return nil, status.Errorf(codes.Unimplemented, "method GetPeerAppliedPostureChecks is not implemented")
}

// UpdateAccountSettings mocks UpdateAccountSettings of the AccountManager interface
func (am *MockAccountManager) UpdateAccountSettings(accountID, userID string, newSettings *server.Settings) (*server.Account, error) {
if am.UpdateAccountSettingsFunc != nil {
Expand All @@ -644,19 +635,19 @@ func (am *MockAccountManager) UpdateAccountSettings(accountID, userID string, ne
}

// LoginPeer mocks LoginPeer of the AccountManager interface
func (am *MockAccountManager) LoginPeer(login server.PeerLogin) (*nbpeer.Peer, *server.NetworkMap, error) {
func (am *MockAccountManager) LoginPeer(login server.PeerLogin) (*nbpeer.Peer, *server.NetworkMap, []*posture.Checks, error) {
if am.LoginPeerFunc != nil {
return am.LoginPeerFunc(login)
}
return nil, nil, status.Errorf(codes.Unimplemented, "method LoginPeer is not implemented")
return nil, nil, nil, status.Errorf(codes.Unimplemented, "method LoginPeer is not implemented")
}

// SyncPeer mocks SyncPeer of the AccountManager interface
func (am *MockAccountManager) SyncPeer(sync server.PeerSync, account *server.Account) (*nbpeer.Peer, *server.NetworkMap, error) {
func (am *MockAccountManager) SyncPeer(sync server.PeerSync, account *server.Account) (*nbpeer.Peer, *server.NetworkMap, []*posture.Checks, error) {
if am.SyncPeerFunc != nil {
return am.SyncPeerFunc(sync, account)
}
return nil, nil, status.Errorf(codes.Unimplemented, "method SyncPeer is not implemented")
return nil, nil, nil, status.Errorf(codes.Unimplemented, "method SyncPeer is not implemented")
}

// GetAllConnectedPeers mocks GetAllConnectedPeers of the AccountManager interface
Expand Down
4 changes: 2 additions & 2 deletions management/server/nameserver_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -851,11 +851,11 @@ func initTestNSAccount(t *testing.T, am *DefaultAccountManager) (*Account, error
return nil, err
}

_, _, err = am.AddPeer("", userID, peer1)
_, _, _, err = am.AddPeer("", userID, peer1)
if err != nil {
return nil, err
}
_, _, err = am.AddPeer("", userID, peer2)
_, _, _, err = am.AddPeer("", userID, peer2)
if err != nil {
return nil, err
}
Expand Down
20 changes: 10 additions & 10 deletions management/server/peer_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -92,7 +92,7 @@ func TestAccountManager_GetNetworkMap(t *testing.T) {
return
}

peer1, _, err := manager.AddPeer(setupKey.Key, "", &nbpeer.Peer{
peer1, _, _, err := manager.AddPeer(setupKey.Key, "", &nbpeer.Peer{
Key: peerKey1.PublicKey().String(),
Meta: nbpeer.PeerSystemMeta{Hostname: "test-peer-1"},
})
Expand All @@ -106,7 +106,7 @@ func TestAccountManager_GetNetworkMap(t *testing.T) {
t.Fatal(err)
return
}
_, _, err = manager.AddPeer(setupKey.Key, "", &nbpeer.Peer{
_, _, _, err = manager.AddPeer(setupKey.Key, "", &nbpeer.Peer{
Key: peerKey2.PublicKey().String(),
Meta: nbpeer.PeerSystemMeta{Hostname: "test-peer-2"},
})
Expand Down Expand Up @@ -165,7 +165,7 @@ func TestAccountManager_GetNetworkMapWithPolicy(t *testing.T) {
return
}

peer1, _, err := manager.AddPeer(setupKey.Key, "", &nbpeer.Peer{
peer1, _, _, err := manager.AddPeer(setupKey.Key, "", &nbpeer.Peer{
Key: peerKey1.PublicKey().String(),
Meta: nbpeer.PeerSystemMeta{Hostname: "test-peer-1"},
})
Expand All @@ -179,7 +179,7 @@ func TestAccountManager_GetNetworkMapWithPolicy(t *testing.T) {
t.Fatal(err)
return
}
peer2, _, err := manager.AddPeer(setupKey.Key, "", &nbpeer.Peer{
peer2, _, _, err := manager.AddPeer(setupKey.Key, "", &nbpeer.Peer{
Key: peerKey2.PublicKey().String(),
Meta: nbpeer.PeerSystemMeta{Hostname: "test-peer-2"},
})
Expand Down Expand Up @@ -341,7 +341,7 @@ func TestAccountManager_GetPeerNetwork(t *testing.T) {
return
}

peer1, _, err := manager.AddPeer(setupKey.Key, "", &nbpeer.Peer{
peer1, _, _, err := manager.AddPeer(setupKey.Key, "", &nbpeer.Peer{
Key: peerKey1.PublicKey().String(),
Meta: nbpeer.PeerSystemMeta{Hostname: "test-peer-1"},
})
Expand All @@ -355,7 +355,7 @@ func TestAccountManager_GetPeerNetwork(t *testing.T) {
t.Fatal(err)
return
}
_, _, err = manager.AddPeer(setupKey.Key, "", &nbpeer.Peer{
_, _, _, err = manager.AddPeer(setupKey.Key, "", &nbpeer.Peer{
Key: peerKey2.PublicKey().String(),
Meta: nbpeer.PeerSystemMeta{Hostname: "test-peer-2"},
})
Expand Down Expand Up @@ -413,7 +413,7 @@ func TestDefaultAccountManager_GetPeer(t *testing.T) {
return
}

peer1, _, err := manager.AddPeer("", someUser, &nbpeer.Peer{
peer1, _, _, err := manager.AddPeer("", someUser, &nbpeer.Peer{
Key: peerKey1.PublicKey().String(),
Meta: nbpeer.PeerSystemMeta{Hostname: "test-peer-2"},
})
Expand All @@ -429,7 +429,7 @@ func TestDefaultAccountManager_GetPeer(t *testing.T) {
}

// the second peer added with a setup key
peer2, _, err := manager.AddPeer(setupKey.Key, "", &nbpeer.Peer{
peer2, _, _, err := manager.AddPeer(setupKey.Key, "", &nbpeer.Peer{
Key: peerKey2.PublicKey().String(),
Meta: nbpeer.PeerSystemMeta{Hostname: "test-peer-2"},
})
Expand Down Expand Up @@ -601,7 +601,7 @@ func TestDefaultAccountManager_GetPeers(t *testing.T) {
return
}

_, _, err = manager.AddPeer("", someUser, &nbpeer.Peer{
_, _, _, err = manager.AddPeer("", someUser, &nbpeer.Peer{
Key: peerKey1.PublicKey().String(),
Meta: nbpeer.PeerSystemMeta{Hostname: "test-peer-1"},
})
Expand All @@ -610,7 +610,7 @@ func TestDefaultAccountManager_GetPeers(t *testing.T) {
return
}

_, _, err = manager.AddPeer("", adminUser, &nbpeer.Peer{
_, _, _, err = manager.AddPeer("", adminUser, &nbpeer.Peer{
Key: peerKey2.PublicKey().String(),
Meta: nbpeer.PeerSystemMeta{Hostname: "test-peer-2"},
})
Expand Down

0 comments on commit 2b0f93e

Please sign in to comment.