Reduced usage of NetworkInfo

Author: JoukoVirtanen

integration-tests/pkg/mock_sensor/expect_conn.go

@@ -9,15 +9,17 @@ import (
 
 	collectorAssert "github.com/stackrox/collector/integration-tests/pkg/assert"
 	"github.com/stackrox/collector/integration-tests/pkg/types"
+
+	sensorAPI "github.com/stackrox/rox/generated/internalapi/sensor"
 )
 
 // ExpectConnections waits up to the timeout for the gRPC server to receive
 // the list of expected Connections. It will first check to see if the connections
 // have been received already, and then monitor the live feed of connections
 // until timeout or until all the events have been received.
-func (s *MockSensor) ExpectConnections(t *testing.T, containerID string, timeout time.Duration, expected ...types.NetworkInfo) bool {
+func (s *MockSensor) ExpectConnections(t *testing.T, containerID string, timeout time.Duration, expected ...*sensorAPI.NetworkConnection) bool {
 
-	to_find := funk.Filter(expected, func(x types.NetworkInfo) bool {
+	to_find := funk.Filter(expected, func(x *sensorAPI.NetworkConnection) bool {
 		return !s.HasConnection(containerID, x)
 	}).([]types.NetworkInfo)
 
@@ -39,7 +41,7 @@ loop:
 				continue loop
 			}
 
-			to_find = funk.Filter(expected, func(x types.NetworkInfo) bool {
+			to_find = funk.Filter(expected, func(x *sensorAPI.NetworkConnection) bool {
 				return !s.HasConnection(containerID, x)
 			}).([]types.NetworkInfo)
 
@@ -57,7 +59,7 @@ loop:
 //
 // It does not consider the content of the events, just that a certain number
 // have been received
-func (s *MockSensor) ExpectConnectionsN(t *testing.T, containerID string, timeout time.Duration, n int) []types.NetworkInfo {
+func (s *MockSensor) ExpectConnectionsN(t *testing.T, containerID string, timeout time.Duration, n int) []*sensorAPI.NetworkConnection {
 	if len(s.Connections(containerID)) == n {
 		return s.Connections(containerID)
 	}
@@ -82,11 +84,11 @@ loop:
 // ExpectSameElementsConnections compares a list of expected connections to the observed connections. This comparison is done at the beginning, when a new
 // connection arrives, and after a timeout period. The number of connections must match and the expected and observed connections must match, but the order
 // does not matter.
-func (s *MockSensor) ExpectSameElementsConnections(t *testing.T, containerID string, timeout time.Duration, expected ...types.NetworkInfo) bool {
+func (s *MockSensor) ExpectSameElementsConnections(t *testing.T, containerID string, timeout time.Duration, expected ...*sensorAPI.NetworkConnection) bool {
 	types.SortConnections(expected)
 
-	equal := func(c1, c2 types.NetworkInfo) bool {
-		return c1.Equal(c2)
+	equal := func(c1, c2 *sensorAPI.NetworkConnection) bool {
+		return types.Equal(c1, c2)
 	}
 
 	connections := s.Connections(containerID)
@@ -123,7 +125,7 @@ func (s *MockSensor) ExpectSameElementsConnectionsScrapes(t *testing.T, containe
 		types.SortConnections(c2)
 
 		for i := range c2 {
-			if !c1[i].Equal(c2[i]) {
+			if !types.Equal(c1[i], c2[i]) {
 				return false
 			}
 		}

integration-tests/pkg/mock_sensor/server.go

@@ -167,11 +167,11 @@ func (m *MockSensor) GetConnectionsInBatches(containerID string) []types.Network
 
 // Connections returns a list of all connections that have been received for
 // a given container ID
-func (m *MockSensor) Connections(containerID string) []types.NetworkInfo {
+func (m *MockSensor) Connections(containerID string) []*sensorAPI.NetworkConnection {
 	m.networkMutex.Lock()
 	defer m.networkMutex.Unlock()
 
-	allConns := make([]types.NetworkInfo, 0)
+	allConns := make([]*sensorAPI.NetworkConnection, 0)
 	if connections, ok := m.connections[containerID]; ok {
 		conns := make([]types.NetworkInfoBatch, len(connections))
 		copy(conns, connections)
@@ -183,16 +183,16 @@ func (m *MockSensor) Connections(containerID string) []types.NetworkInfo {
 
 		return allConns
 	}
-	return make([]types.NetworkInfo, 0)
+	return make([]*sensorAPI.NetworkConnection, 0)
 }
 
 // HasConnection returns whether a given connection has been seen for a given
 // container ID
-func (m *MockSensor) HasConnection(containerID string, conn types.NetworkInfo) bool {
+func (m *MockSensor) HasConnection(containerID string, conn *sensorAPI.NetworkConnection) bool {
 	conns := m.Connections(containerID)
 	if len(conns) > 0 {
-		return slices.ContainsFunc(conns, func(c types.NetworkInfo) bool {
-			return c.Equal(conn)
+		return slices.ContainsFunc(conns, func(c *sensorAPI.NetworkConnection) bool {
+			return types.Equal(c, conn)
 		})
 	}
 
@@ -348,30 +348,15 @@ func (m *MockSensor) PushSignals(stream sensorAPI.SignalService_PushSignalsServe
 	}
 }
 
-func (m *MockSensor) convertConnection(connection *sensorAPI.NetworkConnection) types.NetworkInfo {
-	conn := types.NetworkInfo{
-		LocalAddress:   types.TranslateAddress(connection.LocalAddress),
-		RemoteAddress:  types.TranslateAddress(connection.RemoteAddress),
-		Role:           connection.GetRole().String(),
-		SocketFamily:   connection.GetSocketFamily().String(),
-		CloseTimestamp: connection.GetCloseTimestamp().String(),
-	}
-
-	m.logger.Printf("NetworkInfo: %s, %s\n", connection.GetContainerId(), conn)
-
-	return conn
-}
-
-func (m *MockSensor) convertToContainerConnsMap(connections []*sensorAPI.NetworkConnection) map[string][]types.NetworkInfo {
-	containerConnsMap := make(map[string][]types.NetworkInfo)
+func (m *MockSensor) convertToContainerConnsMap(connections []*sensorAPI.NetworkConnection) map[string][]*sensorAPI.NetworkConnection {
+	containerConnsMap := make(map[string][]*sensorAPI.NetworkConnection)
 	for _, connection := range connections {
-		conn := m.convertConnection(connection)
 		containerID := connection.GetContainerId()
 
 		if c, ok := containerConnsMap[containerID]; ok {
-			containerConnsMap[containerID] = append(c, conn)
+			containerConnsMap[containerID] = append(c, connection)
 		} else {
-			containerConnsMap[containerID] = []types.NetworkInfo{conn}
+			containerConnsMap[containerID] = []*sensorAPI.NetworkConnection{connection}
 		}
 	}
 
@@ -462,7 +447,7 @@ func (m *MockSensor) pushLineage(containerID string, process *storage.ProcessSig
 	}
 }
 
-func (m *MockSensor) pushConnections(containerConnsMap map[string][]types.NetworkInfo) {
+func (m *MockSensor) pushConnections(containerConnsMap map[string][]*sensorAPI.NetworkConnection) {
 	m.networkMutex.Lock()
 	defer m.networkMutex.Unlock()
 

integration-tests/pkg/types/network.go

@@ -1,7 +1,7 @@
 package types
 
 import (
-	"fmt"
+	"net"
 	"sort"
 
 	sensorAPI "github.com/stackrox/rox/generated/internalapi/sensor"
@@ -20,7 +20,7 @@ type NetworkInfo struct {
 	CloseTimestamp string
 }
 
-type NetworkInfoBatch []NetworkInfo
+type NetworkInfoBatch []*sensorAPI.NetworkConnection
 
 // TranslateAddress is a helper function for converting binary representations
 // of network addresses (in the signals) to usable forms for testing
@@ -52,52 +52,127 @@ func TranslateAddress(addr *sensorAPI.NetworkAddress) string {
 	return peerId.String()
 }
 
-func (n *NetworkInfo) String() string {
-	return fmt.Sprintf("%s|%s|%s|%s|%s",
-		n.LocalAddress,
-		n.RemoteAddress,
-		n.Role,
-		n.SocketFamily,
-		n.CloseTimestamp)
+func IsActive(conn *sensorAPI.NetworkConnection) bool {
+	// no close timestamp means the connection is open, and active
+	return conn.GetCloseTimestamp() == nil
 }
 
-func (n *NetworkInfo) IsActive() bool {
-	// no close timestamp means the connection is open, and active
-	return n.CloseTimestamp == NilTimestamp
+func Equal(conn1 *sensorAPI.NetworkConnection, conn2 *sensorAPI.NetworkConnection) bool {
+	return EqualNetworkAddress(conn1.LocalAddress, conn2.LocalAddress) &&
+		EqualNetworkAddress(conn1.RemoteAddress, conn2.RemoteAddress) &&
+		conn1.Role == conn2.Role &&
+		conn1.SocketFamily == conn2.SocketFamily &&
+		IsActive(conn1) == IsActive(conn2)
 }
 
-func (n *NetworkInfo) Equal(other NetworkInfo) bool {
-	return n.LocalAddress == other.LocalAddress &&
-		n.RemoteAddress == other.RemoteAddress &&
-		n.Role == other.Role &&
-		n.SocketFamily == other.SocketFamily &&
-		n.IsActive() == other.IsActive()
+func CompareBytes(b1 []byte, b2 []byte) int {
+	if len(b1) != len(b2) {
+		if len(b1) < len(b2) {
+			return -1
+		} else {
+			return 1
+		}
+	}
+
+	for i := range b1 {
+		if b1[i] != b2[i] {
+			if b1[i] < b2[i] {
+				return -1
+			} else {
+				return 1
+			}
+		}
+	}
+
+	return 0
 }
 
-func (n *NetworkInfo) Less(other NetworkInfo) bool {
-	if n.LocalAddress != other.LocalAddress {
-		return n.LocalAddress < other.LocalAddress
+func EqualNetworkAddress(addr1 *sensorAPI.NetworkAddress, addr2 *sensorAPI.NetworkAddress) bool {
+	comp := CompareBytes(addr1.GetAddressData(), addr2.GetAddressData())
+
+	if comp != 0 {
+		return false
 	}
 
-	if n.RemoteAddress != other.RemoteAddress {
-		return n.RemoteAddress < other.RemoteAddress
+	comp = CompareBytes(addr1.GetIpNetwork(), addr2.GetIpNetwork())
+
+	if comp != 0 {
+		return false
 	}
 
-	if n.Role != other.Role {
-		return n.Role < other.Role
+	return addr1.GetPort() == addr2.GetPort()
+}
+
+func LessNetworkAddress(addr1 *sensorAPI.NetworkAddress, addr2 *sensorAPI.NetworkAddress) bool {
+	comp := CompareBytes(addr1.GetAddressData(), addr2.GetAddressData())
+
+	if comp != 0 {
+		return comp < 0
 	}
 
-	if n.SocketFamily != other.SocketFamily {
-		return n.SocketFamily < other.SocketFamily
+	comp = CompareBytes(addr1.GetIpNetwork(), addr2.GetIpNetwork())
+
+	if comp != 0 {
+		return comp < 0
 	}
 
-	if n.IsActive() != other.IsActive() {
-		return n.IsActive()
+	return addr1.GetPort() < addr2.GetPort()
+}
+
+func LessNetworkConnection(conn1 *sensorAPI.NetworkConnection, conn2 *sensorAPI.NetworkConnection) bool {
+	if !EqualNetworkAddress(conn1.LocalAddress, conn2.LocalAddress) {
+		return LessNetworkAddress(conn1.GetLocalAddress(), conn2.GetLocalAddress())
+	}
+
+	if !EqualNetworkAddress(conn1.RemoteAddress, conn2.RemoteAddress) {
+		return LessNetworkAddress(conn1.GetRemoteAddress(), conn2.GetRemoteAddress())
+	}
+
+	if conn1.Role != conn2.Role {
+		return conn1.Role < conn2.Role
+	}
+
+	if conn1.SocketFamily != conn2.SocketFamily {
+		return conn1.SocketFamily < conn2.SocketFamily
+	}
+
+	if IsActive(conn1) != IsActive(conn2) {
+		return IsActive(conn1)
 	}
 
 	return false
 }
 
-func SortConnections(connections []NetworkInfo) {
-	sort.Slice(connections, func(i, j int) bool { return connections[i].Less(connections[j]) })
+func stringToIPBytes(ipStr string) []byte {
+	ip := net.ParseIP(ipStr)
+
+	if ip == nil {
+		return nil
+	}
+
+	return ip.To4()
+
+}
+
+func stringToIPNetworkBytes(ipStr string) []byte {
+	ip := net.ParseIP(ipStr)
+
+	if ip == nil {
+		return nil
+	}
+
+	return append(ip.To4(), 32)
+}
+
+func CreateNetworkAddress(ipAddress string, ipNetwork string, port uint32) *sensorAPI.NetworkAddress {
+
+	return &sensorAPI.NetworkAddress{
+		AddressData: stringToIPBytes(ipAddress),
+		IpNetwork:   stringToIPNetworkBytes(ipNetwork),
+		Port:        port,
+	}
+}
+
+func SortConnections(connections []*sensorAPI.NetworkConnection) {
+	sort.Slice(connections, func(i, j int) bool { return LessNetworkConnection(connections[i], connections[j]) })
 }

integration-tests/suites/base.go

@@ -434,8 +434,10 @@ func (s *IntegrationTestSuiteBase) getIPAddress(containerName string) (string, e
 	return s.Executor().GetContainerIP(containerName)
 }
 
-func (s *IntegrationTestSuiteBase) getPort(containerName string) (string, error) {
-	return s.Executor().GetContainerPort(containerName)
+func (s *IntegrationTestSuiteBase) getPort(containerName string) (uint32, error) {
+	portStr, err := s.Executor().GetContainerPort(containerName)
+	port, _ := strconv.ParseUint(portStr, 10, 32)
+	return uint32(port), err
 }
 
 func (s *IntegrationTestSuiteBase) StartContainerStats() {

integration-tests/suites/process_network.go

@@ -7,6 +7,8 @@ import (
 	"github.com/stackrox/collector/integration-tests/pkg/common"
 	"github.com/stackrox/collector/integration-tests/pkg/config"
 	"github.com/stackrox/collector/integration-tests/pkg/types"
+
+	sensorAPI "github.com/stackrox/rox/generated/internalapi/sensor"
 )
 
 type ProcessNetworkTestSuite struct {
@@ -15,7 +17,7 @@ type ProcessNetworkTestSuite struct {
 	clientIP        string
 	serverContainer string
 	serverIP        string
-	serverPort      string
+	serverPort      uint32
 }
 
 // Launches collector
@@ -147,22 +149,22 @@ func (s *ProcessNetworkTestSuite) TestProcessLineageInfo() {
 
 func (s *ProcessNetworkTestSuite) TestNetworkFlows() {
 	s.Sensor().ExpectConnections(s.T(), s.serverContainer, 10*time.Second,
-		types.NetworkInfo{
-			LocalAddress:   fmt.Sprintf(":%s", s.serverPort),
-			RemoteAddress:  s.clientIP,
-			Role:           "ROLE_SERVER",
-			SocketFamily:   "SOCKET_FAMILY_UNKNOWN",
-			CloseTimestamp: types.NilTimestamp,
+		&sensorAPI.NetworkConnection{
+			LocalAddress:   types.CreateNetworkAddress("", "", s.serverPort),
+			RemoteAddress:  types.CreateNetworkAddress(s.clientIP, "", s.serverPort),
+			Role:           sensorAPI.ClientServerRole_ROLE_SERVER,
+			SocketFamily:   sensorAPI.SocketFamily_SOCKET_FAMILY_UNKNOWN,
+			CloseTimestamp: nil,
 		},
 	)
 
 	s.Sensor().ExpectConnections(s.T(), s.clientContainer, 10*time.Second,
-		types.NetworkInfo{
-			LocalAddress:   "",
-			RemoteAddress:  fmt.Sprintf("%s:%s", s.serverIP, s.serverPort),
-			Role:           "ROLE_CLIENT",
-			SocketFamily:   "SOCKET_FAMILY_UNKNOWN",
-			CloseTimestamp: types.NilTimestamp,
+		&sensorAPI.NetworkConnection{
+			LocalAddress:   types.CreateNetworkAddress("", "", 0),
+			RemoteAddress:  types.CreateNetworkAddress(s.clientIP, "", s.serverPort),
+			Role:           sensorAPI.ClientServerRole_ROLE_CLIENT,
+			SocketFamily:   sensorAPI.SocketFamily_SOCKET_FAMILY_UNKNOWN,
+			CloseTimestamp: nil,
 		},
 	)
 }