manipulation: Return a callback that queues packets

Can be used by the manipulator when it wants to init connections

Author: roykuper13

linuxswitch/connections.py

@@ -6,8 +6,8 @@
 from contextlib import suppress
 import concurrent.futures
 
-from linuxswitch.util import PortType, add_vlan_tag, fix_checksums, apply_bpf_filter
-from linuxswitch.manipulation import ManipulateArgs, ManipulateActions
+from linuxswitch.util import PortType, DeviceType, add_vlan_tag, fix_checksums, apply_bpf_filter
+from linuxswitch.manipulation import ManipulateArgs
 
 
 DeviceEntry = namedtuple('DeviceEntry', [
@@ -43,52 +43,41 @@ async def _read_message_queue(self):
             await self._process_packet(packet, dev_entry)
 
     async def _process_packet(self, packet, src_device_entry):
-        # Should we deal with tagging. We left an option for the manipulator
-        # to tag the packet.
-        should_inject_raw = False
-
-        packets = [packet]
-
-        if self._manipulate_cb is not None:
+        """
+        If `src_device_entry` is None, we're not performing manipulation.
+        """
+        if src_device_entry is not None and self._manipulate_cb is not None:
             # We'll manipulate if there's no filter or if there's
             # filter and the packet matches the filter
             if self._manipulate_filter == NO_FILTER or \
                     apply_bpf_filter(packet, self._manipulate_filter):
 
-                # In case the packet matches the filter, and `duplicate` is not
-                # set, we remove the packet from the list, since the manipulator
-                # might drop the packet.
-                if not self._manipulate_dup:
-                    packets.remove(packet)
-
-                packet, action = await self._event_loop.run_in_executor(
+                await self._event_loop.run_in_executor(
                     self._executers,
                     self._manipulate_cb,
                     ManipulateArgs(packet, src_device_entry.port_type, src_device_entry.vlan))
 
-                if action == ManipulateActions.INJECT_RAW:
-                    should_inject_raw = True
-                    packets.append(packet)
-                elif action == ManipulateActions.HANDLE_ENCAP:
-                    should_inject_raw = False
-                    packets.append(packet)
-                # else - action is DROP, so we don't add the packet to the list.
-
-        for packet in packets:
+                # In case the packet matches the filter, and `duplicate` is not
+                # set, then only the manipulator should process the packet.
+                # Therefore, we abort here.
+                if not self._manipulate_dup:
+                    return None
 
-            src_vlan = self._get_vlan_by_mac(Ether(packet).src)
-            if src_vlan is None:
-                return None
+        src_vlan = self._get_vlan_by_mac(Ether(packet).src)
+        if src_vlan is None:
+            return None
 
-            # Looking for the destination device
-            for dst_device in self._devices:
-                if dst_device.dev.get_mac == Ether(packet).dst and dst_device.vlan == src_vlan:
+        # Looking for the destination device
+        for dst_device in self._devices:
+            if dst_device.dev.get_mac == Ether(packet).dst and dst_device.vlan == src_vlan:
 
-                    if dst_device.port_type == PortType.TRUNK and not should_inject_raw:
-                        packet = add_vlan_tag(packet, dst_device.vlan)
+                # TODO: currently we're taggint packets. Maybe the Device
+                # should state whether we should tag or inject raw.
+                if dst_device.port_type == PortType.TRUNK:
+                    packet = add_vlan_tag(packet, dst_device.vlan)
 
-                    packet = fix_checksums(packet)
-                    await self._event_loop.sock_sendall(dst_device.sock, packet)
+                packet = fix_checksums(packet)
+                await self._event_loop.sock_sendall(dst_device.sock, packet)
 
     def _read_raw_packet(self, device_entry):
         try:
@@ -137,12 +126,12 @@ def _get_vlan_by_mac(self, mac):
 
         return None
 
-    def append_device(self, dev, vlan, port_type):
+    def append_device(self, dev, vlan, port_type, iface_name):
         def _append_device_entry(self, new_dev_entry):
             self._devices.append(new_dev_entry)
 
         sock = socket.socket(socket.AF_PACKET, socket.SOCK_RAW, socket.htons(0x0003))
-        sock.bind(('br-{}'.format(dev.get_name), 0))
+        sock.bind((iface_name, 0))
 
         new_dev_entry = DeviceEntry(dev, sock, vlan, port_type)
 
@@ -168,15 +157,18 @@ def _remove_device_entry(self, dev_to_remove):
 
     def set_manipulation(self, cb, bpf_filter, duplicate):
         ''' Assumes cb is in valid form '''
-        def __set_manipulation(cb, bpf_filter):
+        def __set_manipulation(cb, bpf_filter, duplicate):
             self._manipulate_cb = cb
             self._manipulate_filter = bpf_filter
             self._manipulate_dup = duplicate
 
         # set_manipulation is called from the main thread.
         # Since we're affecting the event loop thread, we should call_soon_threadsafe,
         # so the call will be synchronized.
-        self._event_loop.call_soon_threadsafe(__set_manipulation, cb, bpf_filter)
+        self._event_loop.call_soon_threadsafe(__set_manipulation, cb, bpf_filter, duplicate)
+
+    def manipulator_queue_packet(self, packet):
+        self._message_queue.put_nowait((packet, None))
 
     def start_connections_thread(self):
         self._connections_thread.start()

linuxswitch/exception.py

@@ -12,3 +12,7 @@ class BridgeInterfaceCreationException(Exception):
 
 class NamespaceConnectionException(Exception):
     pass
+
+
+class ManipulationCallbackException(Exception):
+    pass

linuxswitch/manipulation.py

@@ -1,37 +1,14 @@
-import enum
 import inspect
 from collections import namedtuple
 
 
-class ManipulateActions(enum.Enum):
-    DROP = 0
-    # The manipulator will deal with tagging, if needed.
-    INJECT_RAW = 1
-    # The switch should deal with tagging, if needed.
-    HANDLE_ENCAP = 2
-
-
 ManipulateArgs = namedtuple('ManipulateInfo', [
     'packet',
     # can be used if the manipulator wants to deal with tagging.
     'is_trunk_port',
     'src_vlan',
 ])
 
-ManipulateRet = namedtuple('ManipulateRet', [
-    'packet',
-    # From ManipulateActions
-    'action',
-])
-
-
-def default_manipulation_cb(manipulate_args: ManipulateArgs) -> ManipulateRet:
-    '''
-    Can be used as an example of a valid signature of manipulation callback.
-    This default callback just returns the packet.
-    '''
-    return ManipulateRet(manipulate_args.packet, ManipulateActions.HANDLE_ENCAP)
-
 
 def validate_manipulation_cb(cb):
     # TODO: Currently annotations are required in order to check
@@ -43,4 +20,4 @@ def validate_manipulation_cb(cb):
 
     param = list(sig.parameters.values())[0]
 
-    return param.annotation == ManipulateArgs and sig.return_annotation == ManipulateRet
+    return param.annotation == ManipulateArgs

linuxswitch/switch.py

@@ -2,7 +2,9 @@
 from linuxswitch.util import shell_run_and_check, PortType
 from linuxswitch.manipulation import validate_manipulation_cb
 from linuxswitch.exception import (NamespaceCreationException,
-                                   BridgeInterfaceCreationException, NamespaceConnectionException)
+                                   BridgeInterfaceCreationException,
+                                   NamespaceConnectionException,
+                                   ManipulationCallbackException)
 
 
 class Switch(object):
@@ -56,43 +58,37 @@ def set_manipulation(self, cb, punt_policies_bpf=NO_FILTER, duplicate=False):
         Sets the manipulation routine.
 
         Before the switch checks to which device a packet should be forwarded,
-        it sends the packet the a manipulation routine.
+        it sends the packet to the manipulation routine.
         You can pass a callback that manipulate the packet, for example, change the
         destination and source addresses and the vlan tag, in order to perform VLAN-Hopping,
         or NAT.
 
+        The function returns a callback that can be used by the manipulator in order
+        to queue packets that should be processed by the `Switch`.
+
         In addition, In case you want to emulate "Punt-Policies" (decide what packets
         should be forwarded to the manipulation callback), you can pass a bpf filter.
         As mentioned, only packets that are filtered by this bpf filter will be forwarded
         to the manipultion routine before they get processed by the switch.
 
         :param cb: A callback the the manipulation routine.
-                   The callback must have arguments and return annotations,
+                   The callback must have arguments annotations,
                    since this function validates this callback using them.
-
-                   A valid callback has to match the following form:
-                   1. Recieve one argument. Its type is ManipulateArgs.
-                   2. Return ManipulateRet.
-
-                   An example of a valid callback (taken from manipulation.py:
-                   `default_manipulation_cb`):
-
-                def default_manipulation_cb(manipulate_args: ManipulateArgs) -> ManipulateRet:
-                    return ManipulateRet(manipulate_args.packet, ManipulateActions.HANDLE_ENCAP)
-
-                    This callback for example, just returns the packet.
+                   A valid callback recieves one argument.
+                   Its type is ManipulateArgs (see manipulation.py).
 
         :param punt_policies_bpf: The "Punt-Policies". Default is no-filter.
         :param duplicate: True if the packet should be processed by both switch
                           and manipulator, False if the packet should be processed
                           by the manipulator only.
-        :return: True if the callback is in a valid form, False otherwise.
+        :return: A callback that can be used by the manipulator in order to queue
+                 packets that should be processed by the `Switch`.
         """
-        if validate_manipulation_cb(cb):
-            self._connections.set_manipulation(cb, punt_policies_bpf, duplicate)
-            return True
+        if not validate_manipulation_cb(cb):
+            raise ManipulationCallbackException("Callback has invalid signature")
 
-        return False
+        self._connections.set_manipulation(cb, punt_policies_bpf, duplicate)
+        return self._connections.manipulator_queue_packet
 
     def connect_device_access(self, dev, vlan):
         """
@@ -115,7 +111,7 @@ def connect_device_access(self, dev, vlan):
             raise NamespaceConnectionException("failed to create connection to namespace {}".format(
                 dev.get_name))
 
-        self._connections.append_device(dev, vlan, PortType.ACCESS)
+        self._connections.append_device(dev, vlan, PortType.ACCESS, 'br-{}'.format(dev.get_name))
 
     def connect_device_trunk(self, dev, vlan):
         """
@@ -138,7 +134,7 @@ def connect_device_trunk(self, dev, vlan):
             raise NamespaceConnectionException("failed to create connection to namespace {}".format(
                 dev.get_name))
 
-        self._connections.append_device(dev, vlan, PortType.TRUNK)
+        self._connections.append_device(dev, vlan, PortType.TRUNK, 'br-{}'.format(dev.get_name))
 
     def disconnect_device(self, dev):
         """

linuxswitch/util.py

@@ -13,6 +13,11 @@ class PortType(enum.Enum):
     TRUNK = 1
 
 
+class DeviceType(enum.Enum):
+    DEVICE = 0
+    MANIPULATOR = 1
+
+
 def shell_run_and_check(cmd):
     # Pretty ugly but whatever
     res = (b'', b'') == run_shell_cmd(cmd)

tests/test_manipulation.py

@@ -1,8 +1,9 @@
+import time
 import threading
 from scapy.all import Ether, IP, ICMP, Raw
 
 from linuxswitch.device import Device
-from linuxswitch.manipulation import ManipulateArgs, ManipulateRet, ManipulateActions
+from linuxswitch.manipulation import ManipulateArgs
 from linuxswitch.util import add_vlan_tag
 
 
@@ -14,14 +15,20 @@
 GLOBAL_D3_MAC = None
 GLOBAL_D4_MAC = None
 
+GLOBAL_QUEUE_PKTS_CB = None
 
-def forward_d1_to_d3(arg: ManipulateArgs) -> ManipulateRet:
+GLOBAL_RECIEVED_ECHO_REPLY = False
+
+
+def forward_d1_to_d3(arg: ManipulateArgs):
     '''
     Replace addresses of d2 with d4, so the packet will be hopped
     to the second vlan.
     '''
     global GLOBAL_D3_MAC
     global GLOBAL_D4_MAC
+    global GLOBAL_QUEUE_PKTS_CB
+
     packet = Ether(arg.packet)
 
     packet[IP].src = D4_ADDR
@@ -30,19 +37,21 @@ def forward_d1_to_d3(arg: ManipulateArgs) -> ManipulateRet:
     packet[Ether].src = GLOBAL_D4_MAC
     packet[Ether].dst = GLOBAL_D3_MAC
 
-    return ManipulateRet(Raw(packet).load, ManipulateActions.HANDLE_ENCAP)
+    return GLOBAL_QUEUE_PKTS_CB(Raw(packet).load)
 
 
-def drop_packet(arg: ManipulateArgs) -> ManipulateRet:
-    return ManipulateRet(None, ManipulateActions.DROP)
+def drop_packet(arg: ManipulateArgs):
+    # Does nothing
+    pass
 
 
-def tag_packet_inject_raw(arg: ManipulateArgs) -> ManipulateRet:
-    return ManipulateRet(add_vlan_tag(arg.packet, arg.src_vlan), ManipulateActions.INJECT_RAW)
+def check_echo_reply(arg: ManipulateArgs):
+    global GLOBAL_RECIEVED_ECHO_REPLY
 
+    ECHO_REPLY = 0
+    packet = Ether(arg.packet)
 
-def tag_packet_handle_encap(arg: ManipulateArgs) -> ManipulateRet:
-    return ManipulateRet(add_vlan_tag(arg.packet, arg.src_vlan), ManipulateActions.HANDLE_ENCAP)
+    GLOBAL_RECIEVED_ECHO_REPLY = ICMP in packet and packet['ICMP'].type == ECHO_REPLY
 
 
 def test_vlan_hopping_and_punt_policies(switch):
@@ -60,6 +69,7 @@ def test_vlan_hopping_and_punt_policies(switch):
     """
     global GLOBAL_D3_MAC
     global GLOBAL_D4_MAC
+    global GLOBAL_QUEUE_PKTS_CB
 
     d1 = Device('a', D1_ADDR, '255.255.255.0')
     d2 = Device('b', D2_ADDR, '255.255.255.0')
@@ -77,8 +87,10 @@ def test_vlan_hopping_and_punt_policies(switch):
     GLOBAL_D4_MAC = d4.get_mac
 
     # Setting the manipulation routine, and punt-policies
-    switch.set_manipulation(forward_d1_to_d3, 'src host {} && dst host {}'.format(
-        D1_ADDR, D2_ADDR), False)
+    GLOBAL_QUEUE_PKTS_CB = switch.set_manipulation(
+        forward_d1_to_d3,
+        'src host {} && dst host {}'.format(D1_ADDR, D2_ADDR),
+        False)
 
     def _run_listening_nc(dev):
         out = dev.run_from_namespace('timeout 7s nc -lu 0.0.0.0 4444')
@@ -131,33 +143,27 @@ def test_manipulator_drop_packets(switch):
     switch.disconnect_device(d2)
 
 
-def test_manipulator_inject_raw(switch):
+def test_manipulator_init_connection(switch):
+    global GLOBAL_RECIEVED_ECHO_REPLY
     d1 = Device('a', '192.168.250.1', '255.255.255.0')
     d2 = Device('b', '192.168.250.2', '255.255.255.0')
 
     switch.connect_device_trunk(d1, 20)
     switch.connect_device_trunk(d2, 20)
 
-    # The manipulator will tag all packets, and will set INJECT_RAW
-    # as the action. The original packets won't get processed by switch,
-    # So if we get a valid connection, it means that the manipulator
-    # successfully tagged the packet, and the switch didn't deal with
-    # with tagging; it just send the packet as is (raw).
-    switch.set_manipulation(tag_packet_inject_raw, duplicate=False)
+    cb = switch.set_manipulation(check_echo_reply,
+                                 'icmp && src host 192.168.250.1',
+                                 duplicate=False)
 
-    out = d1.run_from_namespace('ping -c 1 192.168.250.2')
-    assert '1 packets transmitted, 1 received' in out
+    pkt = Ether(src=d2.get_mac, dst=d1.get_mac) / \
+        IP(src='192.168.250.2', dst='192.168.250.1', ttl=20) / ICMP()
 
-    out = d2.run_from_namespace('ping -c 1 192.168.250.1')
-    assert '1 packets transmitted, 1 received' in out
+    cb(Raw(pkt).load)
 
-    switch.set_manipulation(tag_packet_handle_encap, duplicate=False)
+    # Wait for the packet to get processed by the switch
+    time.sleep(1.5)
 
-    out = d1.run_from_namespace('ping -c 1 192.168.250.2 -W 2')
-    assert '1 packets transmitted, 0 received' in out
-
-    out = d2.run_from_namespace('ping -c 1 192.168.250.1 -W 2')
-    assert '1 packets transmitted, 0 received' in out
+    assert GLOBAL_RECIEVED_ECHO_REPLY
 
     switch.disconnect_device(d1)
     switch.disconnect_device(d2)