From d2d628b67653fff4339d90ec72f88ef3cf693e6e Mon Sep 17 00:00:00 2001 From: Chris McCarthy Date: Fri, 5 Jan 2024 22:11:37 +0000 Subject: [PATCH] #2139 - Fixed unicast and broadcast functionality properly --- CHANGELOG.md | 8 + .../simulator/network/hardware/base.py | 27 ++- .../network/hardware/nodes/router.py | 45 ++++- .../simulator/system/core/session_manager.py | 83 +++++--- .../simulator/system/core/software_manager.py | 22 ++- src/primaite/simulator/system/software.py | 21 +- .../network/test_broadcast.py | 180 ++++++++++++++++++ 7 files changed, 341 insertions(+), 45 deletions(-) create mode 100644 tests/integration_tests/network/test_broadcast.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 96634b28..60961802 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -38,10 +38,18 @@ SessionManager. - HTTP Services: `WebBrowser` to simulate a web client and `WebServer` - Fixed an issue where the services were still able to run even though the node the service is installed on is turned off - NTP Services: `NTPClient` and `NTPServer` +- **RouterNIC Class**: Introduced a new class `RouterNIC`, extending the standard `NIC` functionality. This class is specifically designed for router operations, optimizing the processing and routing of network traffic. + - **Custom Layer-3 Processing**: The `RouterNIC` class includes custom handling for network frames, bypassing standard Node NIC's Layer 3 broadcast/unicast checks. This allows for more efficient routing behavior in network scenarios where router-specific frame processing is required. + - **Enhanced Frame Reception**: The `receive_frame` method in `RouterNIC` is tailored to handle frames based on Layer 2 (Ethernet) checks, focusing on MAC address-based routing and broadcast frame acceptance. + ### Changed - Integrated the RouteTable into the Routers frame processing. - Frames are now dropped when their TTL reaches 0 +- **NIC Functionality Update**: Updated the Network Interface Card (`NIC`) functionality to support Layer 3 (L3) broadcasts. + - **Layer 3 Broadcast Handling**: Enhanced the existing `NIC` classes to correctly process and handle Layer 3 broadcasts. This update allows devices using standard NICs to effectively participate in network activities that involve L3 broadcasting. + - **Improved Frame Reception Logic**: The `receive_frame` method of the `NIC` class has been updated to include additional checks and handling for L3 broadcasts, ensuring proper frame processing in a wider range of network scenarios. + ### Removed - Removed legacy simulation modules: `acl`, `common`, `environment`, `links`, `nodes`, `pol` diff --git a/src/primaite/simulator/network/hardware/base.py b/src/primaite/simulator/network/hardware/base.py index c27378a8..7e6e0a3b 100644 --- a/src/primaite/simulator/network/hardware/base.py +++ b/src/primaite/simulator/network/hardware/base.py @@ -274,11 +274,20 @@ class NIC(SimComponent): def receive_frame(self, frame: Frame) -> bool: """ - Receive a network frame from the connected link if the NIC is enabled. + Receive a network frame from the connected link, processing it if the NIC is enabled. - The Frame is passed to the Node. + This method decrements the Time To Live (TTL) of the frame, captures it using PCAP (Packet Capture), and checks + if the frame is either a broadcast or destined for this NIC. If the frame is acceptable, it is passed to the + connected node. The method also handles the discarding of frames with TTL expired and logs this event. - :param frame: The network frame being received. + The frame's reception is based on various conditions: + - If the NIC is disabled, the frame is not processed. + - If the TTL of the frame reaches zero after decrement, it is discarded and logged. + - If the frame is a broadcast or its destination MAC/IP address matches this NIC's, it is accepted. + - All other frames are dropped and logged or printed to the console. + + :param frame: The network frame being received. This should be an instance of the Frame class. + :return: Returns True if the frame is processed and passed to the node, False otherwise. """ if self.enabled: frame.decrement_ttl() @@ -288,7 +297,17 @@ class NIC(SimComponent): frame.set_received_timestamp() self.pcap.capture(frame) # If this destination or is broadcast - if frame.ethernet.dst_mac_addr == self.mac_address or frame.ethernet.dst_mac_addr == "ff:ff:ff:ff:ff:ff": + accept_frame = False + + # Check if it's a broadcast: + if frame.ethernet.dst_mac_addr == "ff:ff:ff:ff:ff:ff": + if frame.ip.dst_ip_address in {self.ip_address, self.ip_network.broadcast_address}: + accept_frame = True + else: + if frame.ethernet.dst_mac_addr == self.mac_address: + accept_frame = True + + if accept_frame: self._connected_node.receive_frame(frame=frame, from_nic=self) return True return False diff --git a/src/primaite/simulator/network/hardware/nodes/router.py b/src/primaite/simulator/network/hardware/nodes/router.py index 172cc711..473712ea 100644 --- a/src/primaite/simulator/network/hardware/nodes/router.py +++ b/src/primaite/simulator/network/hardware/nodes/router.py @@ -690,6 +690,47 @@ class RouterICMP(ICMP): self.router.process_frame(frame, from_nic) +class RouterNIC(NIC): + """ + A Router-specific Network Interface Card (NIC) that extends the standard NIC functionality. + + This class overrides the standard Node NIC's Layer 3 (L3) broadcast/unicast checks. It is designed + to handle network frames in a manner specific to routers, allowing them to efficiently process + and route network traffic. + """ + + def receive_frame(self, frame: Frame) -> bool: + """ + Receive and process a network frame from the connected link, provided the NIC is enabled. + + This method is tailored for router behavior. It decrements the frame's Time To Live (TTL), checks for TTL + expiration, and captures the frame using PCAP (Packet Capture). The frame is accepted if it is destined for + this NIC's MAC address or is a broadcast frame. + + Key Differences from Standard NIC: + - Does not perform Layer 3 (IP-based) broadcast checks. + - Only checks for Layer 2 (Ethernet) destination MAC address and broadcast frames. + + :param frame: The network frame being received. This should be an instance of the Frame class. + :return: Returns True if the frame is processed and passed to the connected node, False otherwise. + """ + if self.enabled: + frame.decrement_ttl() + if frame.ip and frame.ip.ttl < 1: + self._connected_node.sys_log.info("Frame discarded as TTL limit reached") + return False + frame.set_received_timestamp() + self.pcap.capture(frame) + # If this destination or is broadcast + if frame.ethernet.dst_mac_addr == self.mac_address or frame.ethernet.dst_mac_addr == "ff:ff:ff:ff:ff:ff": + self._connected_node.receive_frame(frame=frame, from_nic=self) + return True + return False + + def __str__(self) -> str: + return f"{self.mac_address}/{self.ip_address}" + + class Router(Node): """ A class to represent a network router node. @@ -700,7 +741,7 @@ class Router(Node): """ num_ports: int - ethernet_ports: Dict[int, NIC] = {} + ethernet_ports: Dict[int, RouterNIC] = {} acl: AccessControlList route_table: RouteTable arp: RouterARPCache @@ -719,7 +760,7 @@ class Router(Node): kwargs["icmp"] = RouterICMP(sys_log=kwargs.get("sys_log"), arp_cache=kwargs.get("arp"), router=self) super().__init__(hostname=hostname, num_ports=num_ports, **kwargs) for i in range(1, self.num_ports + 1): - nic = NIC(ip_address="127.0.0.1", subnet_mask="255.0.0.0", gateway="0.0.0.0") + nic = RouterNIC(ip_address="127.0.0.1", subnet_mask="255.0.0.0", gateway="0.0.0.0") self.connect_nic(nic) self.ethernet_ports[i] = nic diff --git a/src/primaite/simulator/system/core/session_manager.py b/src/primaite/simulator/system/core/session_manager.py index 8658f155..a95846a3 100644 --- a/src/primaite/simulator/system/core/session_manager.py +++ b/src/primaite/simulator/system/core/session_manager.py @@ -1,6 +1,6 @@ from __future__ import annotations -from ipaddress import IPv4Address +from ipaddress import IPv4Address, IPv4Network from typing import Any, Dict, Optional, Tuple, TYPE_CHECKING, Union from prettytable import MARKDOWN, PrettyTable @@ -141,41 +141,76 @@ class SessionManager: def receive_payload_from_software_manager( self, payload: Any, - dst_ip_address: Optional[IPv4Address] = None, + dst_ip_address: Optional[Union[IPv4Address, IPv4Network]] = None, dst_port: Optional[Port] = None, session_id: Optional[str] = None, is_reattempt: bool = False, ) -> Union[Any, None]: """ - Receive a payload from the SoftwareManager. + Receive a payload from the SoftwareManager and send it to the appropriate NIC for transmission. - If no session_id, a Session is established. Once established, the payload is sent to ``send_payload_to_nic``. + This method supports both unicast and Layer 3 broadcast transmissions. If `dst_ip_address` is an + IPv4Network, a broadcast is initiated. For unicast, the destination MAC address is resolved via ARP. + A new session is established if `session_id` is not provided, and an existing session is used otherwise. :param payload: The payload to be sent. - :param session_id: The Session ID the payload is to originate from. Optional. If None, one will be created. + :param dst_ip_address: The destination IP address or network for broadcast. Optional. + :param dst_port: The destination port for the TCP packet. Optional. + :param session_id: The Session ID from which the payload originates. Optional. + :param is_reattempt: Flag to indicate if this is a reattempt after an ARP request. Default is False. + :return: The outcome of sending the frame, or None if sending was unsuccessful. """ + is_broadcast = False + outbound_nic = None + dst_mac_address = None + + # Use session details if session_id is provided if session_id: session = self.sessions_by_uuid[session_id] - dst_ip_address = self.sessions_by_uuid[session_id].with_ip_address - dst_port = self.sessions_by_uuid[session_id].dst_port + dst_ip_address = session.with_ip_address + dst_port = session.dst_port - dst_mac_address = self.arp_cache.get_arp_cache_mac_address(dst_ip_address) + # Determine if the payload is for broadcast or unicast - if dst_mac_address: - outbound_nic = self.arp_cache.get_arp_cache_nic(dst_ip_address) + # Handle broadcast transmission + if isinstance(dst_ip_address, IPv4Network): + is_broadcast = True + dst_ip_address = dst_ip_address.broadcast_address + if dst_ip_address: + # Find a suitable NIC for the broadcast + for nic in self.arp_cache.nics.values(): + if dst_ip_address in nic.ip_network and nic.enabled: + dst_mac_address = "ff:ff:ff:ff:ff:ff" + outbound_nic = nic else: - if not is_reattempt: - self.arp_cache.send_arp_request(dst_ip_address) - return self.receive_payload_from_software_manager( - payload=payload, - dst_ip_address=dst_ip_address, - dst_port=dst_port, - session_id=session_id, - is_reattempt=True, - ) - else: - return + # Resolve MAC address for unicast transmission + dst_mac_address = self.arp_cache.get_arp_cache_mac_address(dst_ip_address) + # Resolve outbound NIC for unicast transmission + if dst_mac_address: + outbound_nic = self.arp_cache.get_arp_cache_nic(dst_ip_address) + + # If MAC address not found, initiate ARP request + else: + if not is_reattempt: + self.arp_cache.send_arp_request(dst_ip_address) + # Reattempt payload transmission after ARP request + return self.receive_payload_from_software_manager( + payload=payload, + dst_ip_address=dst_ip_address, + dst_port=dst_port, + session_id=session_id, + is_reattempt=True, + ) + else: + # Return None if reattempt fails + return + + # Check if outbound NIC and destination MAC address are resolved + if not outbound_nic or not dst_mac_address: + return False + + # Construct the frame for transmission frame = Frame( ethernet=EthernetHeader(src_mac_addr=outbound_nic.mac_address, dst_mac_addr=dst_mac_address), ip=IPPacket( @@ -189,15 +224,17 @@ class SessionManager: payload=payload, ) - if not session_id: + # Manage session for unicast transmission + if not (is_broadcast and session_id): session_key = self._get_session_key(frame, inbound_frame=False) session = self.sessions_by_key.get(session_key) if not session: - # Create new session + # Create a new session if it doesn't exist session = Session.from_session_key(session_key) self.sessions_by_key[session_key] = session self.sessions_by_uuid[session.uuid] = session + # Send the frame through the NIC return outbound_nic.send_frame(frame) def receive_frame(self, frame: Frame): diff --git a/src/primaite/simulator/system/core/software_manager.py b/src/primaite/simulator/system/core/software_manager.py index 21a121c1..95948a1e 100644 --- a/src/primaite/simulator/system/core/software_manager.py +++ b/src/primaite/simulator/system/core/software_manager.py @@ -1,4 +1,4 @@ -from ipaddress import IPv4Address +from ipaddress import IPv4Address, IPv4Network from typing import Any, Dict, List, Optional, Tuple, TYPE_CHECKING, Union from prettytable import MARKDOWN, PrettyTable @@ -130,20 +130,28 @@ class SoftwareManager: def send_payload_to_session_manager( self, payload: Any, - dest_ip_address: Optional[IPv4Address] = None, + dest_ip_address: Optional[Union[IPv4Address, IPv4Network]] = None, dest_port: Optional[Port] = None, session_id: Optional[str] = None, ) -> bool: """ - Send a payload to the SessionManager. + Sends a payload to the SessionManager for network transmission. + + This method is responsible for initiating the process of sending network payloads. It supports both + unicast and Layer 3 broadcast transmissions. For broadcasts, the destination IP should be specified + as an IPv4Network. :param payload: The payload to be sent. - :param dest_ip_address: The ip address of the payload destination. - :param dest_port: The port of the payload destination. - :param session_id: The Session ID the payload is to originate from. Optional. + :param dest_ip_address: The IP address or network (for broadcasts) of the payload destination. + :param dest_port: The destination port for the payload. Optional. + :param session_id: The Session ID from which the payload originates. Optional. + :return: True if the payload was successfully sent, False otherwise. """ return self.session_manager.receive_payload_from_software_manager( - payload=payload, dst_ip_address=dest_ip_address, dst_port=dest_port, session_id=session_id + payload=payload, + dst_ip_address=dest_ip_address, + dst_port=dest_port, + session_id=session_id, ) def receive_payload_from_session_manager(self, payload: Any, port: Port, protocol: IPProtocol, session_id: str): diff --git a/src/primaite/simulator/system/software.py b/src/primaite/simulator/system/software.py index b393ffd8..d8aed2fb 100644 --- a/src/primaite/simulator/system/software.py +++ b/src/primaite/simulator/system/software.py @@ -2,8 +2,8 @@ import copy from abc import abstractmethod from datetime import datetime from enum import Enum -from ipaddress import IPv4Address -from typing import Any, Dict, Optional +from ipaddress import IPv4Address, IPv4Network +from typing import Any, Dict, Optional, Union from primaite.simulator.core import _LOGGER, RequestManager, RequestType, SimComponent from primaite.simulator.file_system.file_system import FileSystem, Folder @@ -317,19 +317,22 @@ class IOSoftware(Software): self, payload: Any, session_id: Optional[str] = None, - dest_ip_address: Optional[IPv4Address] = None, + dest_ip_address: Optional[Union[IPv4Address, IPv4Network]] = None, dest_port: Optional[Port] = None, **kwargs, ) -> bool: """ - Sends a payload to the SessionManager. + Sends a payload to the SessionManager for network transmission. + + This method is responsible for initiating the process of sending network payloads. It supports both + unicast and Layer 3 broadcast transmissions. For broadcasts, the destination IP should be specified + as an IPv4Network. It delegates the actual sending process to the SoftwareManager. :param payload: The payload to be sent. - :param dest_ip_address: The ip address of the payload destination. - :param dest_port: The port of the payload destination. - :param session_id: The Session ID the payload is to originate from. Optional. - - :return: True if successful, False otherwise. + :param dest_ip_address: The IP address or network (for broadcasts) of the payload destination. + :param dest_port: The destination port for the payload. Optional. + :param session_id: The Session ID from which the payload originates. Optional. + :return: True if the payload was successfully sent, False otherwise. """ if not self._can_perform_action(): return False diff --git a/tests/integration_tests/network/test_broadcast.py b/tests/integration_tests/network/test_broadcast.py new file mode 100644 index 00000000..b9ecb28b --- /dev/null +++ b/tests/integration_tests/network/test_broadcast.py @@ -0,0 +1,180 @@ +from ipaddress import IPv4Address, IPv4Network +from typing import Any, Dict, List, Tuple + +import pytest + +from primaite.simulator.network.container import Network +from primaite.simulator.network.hardware.nodes.computer import Computer +from primaite.simulator.network.hardware.nodes.server import Server +from primaite.simulator.network.hardware.nodes.switch import Switch +from primaite.simulator.network.transmission.network_layer import IPProtocol +from primaite.simulator.network.transmission.transport_layer import Port +from primaite.simulator.system.applications.application import Application +from primaite.simulator.system.services.service import Service + + +class BroadcastService(Service): + """A service for sending broadcast and unicast messages over a network.""" + + def __init__(self, **kwargs): + # Set default service properties for broadcasting + kwargs["name"] = "BroadcastService" + kwargs["port"] = Port.HTTP + kwargs["protocol"] = IPProtocol.TCP + super().__init__(**kwargs) + + def describe_state(self) -> Dict: + # Implement state description for the service + pass + + def unicast(self, ip_address: IPv4Address): + # Send a unicast payload to a specific IP address + super().send( + payload="unicast", + dest_ip_address=ip_address, + dest_port=Port.HTTP, + ) + + def broadcast(self, ip_network: IPv4Network): + # Send a broadcast payload to an entire IP network + super().send( + payload="broadcast", + dest_ip_address=ip_network, + dest_port=Port.HTTP, + ) + + +class BroadcastClient(Application): + """A client application to receive broadcast and unicast messages.""" + + payloads_received: List = [] + + def __init__(self, **kwargs): + # Set default client properties + kwargs["name"] = "BroadcastClient" + kwargs["port"] = Port.HTTP + kwargs["protocol"] = IPProtocol.TCP + super().__init__(**kwargs) + + def describe_state(self) -> Dict: + # Implement state description for the application + pass + + def receive(self, payload: Any, session_id: str, **kwargs) -> bool: + # Append received payloads to the list and print a message + self.payloads_received.append(payload) + print(f"Payload: {payload} received on node {self.sys_log.hostname}") + + +@pytest.fixture(scope="function") +def broadcast_network() -> Network: + network = Network() + + client_1 = Computer( + hostname="client_1", + ip_address="192.168.1.2", + subnet_mask="255.255.255.0", + default_gateway="192.168.1.1", + start_up_duration=0, + ) + client_1.power_on() + client_1.software_manager.install(BroadcastClient) + application_1 = client_1.software_manager.software["BroadcastClient"] + application_1.run() + + client_2 = Computer( + hostname="client_2", + ip_address="192.168.1.3", + subnet_mask="255.255.255.0", + default_gateway="192.168.1.1", + start_up_duration=0, + ) + client_2.power_on() + client_2.software_manager.install(BroadcastClient) + application_2 = client_2.software_manager.software["BroadcastClient"] + application_2.run() + + server_1 = Server( + hostname="server_1", + ip_address="192.168.1.1", + subnet_mask="255.255.255.0", + default_gateway="192.168.1.1", + start_up_duration=0, + ) + server_1.power_on() + + server_1.software_manager.install(BroadcastService) + service: BroadcastService = server_1.software_manager.software["BroadcastService"] + service.start() + + switch_1 = Switch(hostname="switch_1", num_ports=6, start_up_duration=0) + switch_1.power_on() + + network.connect(endpoint_a=client_1.ethernet_port[1], endpoint_b=switch_1.switch_ports[1]) + network.connect(endpoint_a=client_2.ethernet_port[1], endpoint_b=switch_1.switch_ports[2]) + network.connect(endpoint_a=server_1.ethernet_port[1], endpoint_b=switch_1.switch_ports[3]) + + return network + + +@pytest.fixture(scope="function") +def broadcast_service_and_clients(broadcast_network) -> Tuple[BroadcastService, BroadcastClient, BroadcastClient]: + client_1: BroadcastClient = broadcast_network.get_node_by_hostname("client_1").software_manager.software[ + "BroadcastClient" + ] + client_2: BroadcastClient = broadcast_network.get_node_by_hostname("client_2").software_manager.software[ + "BroadcastClient" + ] + service: BroadcastService = broadcast_network.get_node_by_hostname("server_1").software_manager.software[ + "BroadcastService" + ] + + return service, client_1, client_2 + + +def test_broadcast_correct_subnet(broadcast_service_and_clients): + service, client_1, client_2 = broadcast_service_and_clients + + assert not client_1.payloads_received + assert not client_2.payloads_received + + service.broadcast(IPv4Network("192.168.1.0/24")) + + assert client_1.payloads_received == ["broadcast"] + assert client_2.payloads_received == ["broadcast"] + + +def test_broadcast_incorrect_subnet(broadcast_service_and_clients): + service, client_1, client_2 = broadcast_service_and_clients + + assert not client_1.payloads_received + assert not client_2.payloads_received + + service.broadcast(IPv4Network("192.168.2.0/24")) + + assert not client_1.payloads_received + assert not client_2.payloads_received + + +def test_unicast_correct_address(broadcast_service_and_clients): + service, client_1, client_2 = broadcast_service_and_clients + + assert not client_1.payloads_received + assert not client_2.payloads_received + + service.unicast(IPv4Address("192.168.1.2")) + + assert client_1.payloads_received == ["unicast"] + assert not client_2.payloads_received + + +def test_unicast_incorrect_address(broadcast_service_and_clients): + service, client_1, client_2 = broadcast_service_and_clients + + assert not client_1.payloads_received + assert not client_2.payloads_received + + service.unicast(IPv4Address("192.168.2.2")) + + assert not client_1.payloads_received + assert not client_2.payloads_received