Changeset 166 for etherws

Timestamp:

06/27/12 01:38:46 (12 years ago)

Author:

atzm

Message:

• restructure

File:

• etherws/trunk/etherws.py (modified) (12 diffs)

etherws/trunk/etherws.py

@@ -51 +51 @@
 import traceback
 
-import pytun
 import websocket
 import tornado.web
@@ -58 +57 @@
 import tornado.httpserver
 
+from pytun import TunTapDevice, IFF_TAP, IFF_NO_PI
+
 
 class DebugMixIn(object):
-    def dprintf(self, msg, func):
+    def dprintf(self, msg, func=lambda: ()):
         if self._debug:
             prefix = '[%s] %s - ' % (time.asctime(), self.__class__.__name__)
@@ -93 +94 @@
 
 
-class SwitchingTable(DebugMixIn):
+class FDB(DebugMixIn):
     def __init__(self, ageout=300, debug=False):
         self._ageout = ageout
         self._debug = debug
-        self._table = {}
-
-    def learn(self, frame, port):
-        mac = frame.src_mac
-        vid = frame.vid
-
-        if vid not in self._table:
-            self._table[vid] = {}
-
-        self._table[vid][mac] = {'time': time.time(), 'port': port}
-        self.dprintf('learned: [%d] %s\n',
-                     lambda: (vid, mac.encode('hex')))
+        self._dict = {}
 
     def lookup(self, frame):
@@ -114 +104 @@
         vid = frame.vid
 
-        group = self._table.get(vid, None)
+        group = self._dict.get(vid, None)
         if not group:
             return None
@@ -123 +113 @@
 
         if time.time() - entry['time'] > self._ageout:
-            del self._table[vid][mac]
-            if not self._table[vid]:
-                del self._table[vid]
+            del self._dict[vid][mac]
+            if not self._dict[vid]:
+                del self._dict[vid]
             self.dprintf('aged out: [%d] %s\n',
                          lambda: (vid, mac.encode('hex')))
@@ -132 +122 @@
         return entry['port']
 
+    def learn(self, port, frame):
+        mac = frame.src_mac
+        vid = frame.vid
+
+        if vid not in self._dict:
+            self._dict[vid] = {}
+
+        self._dict[vid][mac] = {'time': time.time(), 'port': port}
+        self.dprintf('learned: [%d] %s\n',
+                     lambda: (vid, mac.encode('hex')))
+
     def delete(self, port):
-        for vid in self._table.keys():
-            for mac in self._table[vid].keys():
-                if self._table[vid][mac]['port'] is port:
-                    del self._table[vid][mac]
+        for vid in self._dict.keys():
+            for mac in self._dict[vid].keys():
+                if self._dict[vid][mac]['port'] is port:
+                    del self._dict[vid][mac]
                     self.dprintf('deleted: [%d] %s\n',
                                  lambda: (vid, mac.encode('hex')))
-            if not self._table[vid]:
-                del self._table[vid]
-
-
-class TapHandler(DebugMixIn):
-    READ_SIZE = 65535
-
-    def __init__(self, dev, debug=False):
+            if not self._dict[vid]:
+                del self._dict[vid]
+
+
+class SwitchingHub(DebugMixIn):
+    def __init__(self, fdb, debug=False):
+        self._fdb = fdb
         self._debug = debug
-        self._clients = []
-        self._table = SwitchingTable(debug=debug)
-        self._tap = pytun.TunTapDevice(dev, pytun.IFF_TAP | pytun.IFF_NO_PI)
-        self._tap.up()
-        self.register_client(self)
-
-    def fileno(self):
-        return self._tap.fileno()
-
-    def register_client(self, client):
-        self._clients.append(client)
-
-    def unregister_client(self, client):
-        self._table.delete(client)
-        self._clients.remove(client)
-
-    def write_message(self, message, binary=False):
-        self._tap.write(message)
-
-    def write(self, caller, message):
-        frame = EthernetFrame(message)
-
-        self._table.learn(frame, caller)
-
-        if not frame.multicast:
-            dst = self._table.lookup(frame)
-
-            if dst:
-                dst.write_message(frame.data, True)
-                self.dprintf('sent unicast: [%d] %s -> %s\n',
-                             lambda: (frame.vid,
-                                      frame.src_mac.encode('hex'),
-                                      frame.dst_mac.encode('hex')))
-                return
-
-        clients = self._clients[:]
-        clients.remove(caller)
-
-        for c in clients:
-            c.write_message(frame.data, True)
-
+        self._ports = []
+
+    def register_port(self, port):
+        self._ports.append(port)
+
+    def unregister_port(self, port):
+        self._fdb.delete(port)
+        self._ports.remove(port)
+
+    def forward(self, src_port, frame):
+        try:
+            self._fdb.learn(src_port, frame)
+
+            if not frame.multicast:
+                dst_port = self._fdb.lookup(frame)
+
+                if dst_port:
+                    self._unicast(frame, dst_port)
+                    return
+
+            self._broadcast(frame, src_port)
+
+        except:  # ex. received invalid frame
+            traceback.print_exc()
+
+    def _unicast(self, frame, port):
+        port.write_message(frame.data, True)
+        self.dprintf('sent unicast: [%d] %s -> %s\n',
+                     lambda: (frame.vid,
+                              frame.src_mac.encode('hex'),
+                              frame.dst_mac.encode('hex')))
+
+    def _broadcast(self, frame, *except_ports):
+        ports = self._ports[:]
+        for port in except_ports:
+            ports.remove(port)
+        for port in ports:
+            port.write_message(frame.data, True)
         self.dprintf('sent broadcast: [%d] %s -> %s\n',
                      lambda: (frame.vid,
@@ -194 +191 @@
                               frame.dst_mac.encode('hex')))
 
+
+class TapHandler(DebugMixIn):
+    READ_SIZE = 65535
+
+    def __init__(self, switch, dev, debug=False):
+        self._switch = switch
+        self._dev = dev
+        self._debug = debug
+        self._tap = None
+
+    @property
+    def closed(self):
+        return not self._tap
+
+    def open(self):
+        if not self.closed:
+            raise ValueError('already opened')
+        self._tap = TunTapDevice(self._dev, IFF_TAP | IFF_NO_PI)
+        self._tap.up()
+        self._switch.register_port(self)
+
+    def close(self):
+        if self.closed:
+            raise ValueError('I/O operation on closed tap')
+        self._switch.unregister_port(self)
+        self._tap.close()
+        self._tap = None
+
+    def fileno(self):
+        if self.closed:
+            raise ValueError('I/O operation on closed tap')
+        return self._tap.fileno()
+
+    def write_message(self, message, binary=False):
+        if self.closed:
+            raise ValueError('I/O operation on closed tap')
+        self._tap.write(message)
+
     def __call__(self, fd, events):
+        try:
+            self._switch.forward(self, EthernetFrame(self._read()))
+            return
+        except:
+            traceback.print_exc()
+        tornado.ioloop.IOLoop.instance().stop()
+
+    def _read(self):
+        if self.closed:
+            raise ValueError('I/O operation on closed tap')
         buf = []
-
         while True:
-            data = self._tap.read(self.READ_SIZE)
-
-            if data:
-                buf.append(data)
-
-            if len(data) < self.READ_SIZE:
+            buf.append(self._tap.read(self.READ_SIZE))
+            if len(buf[-1]) < self.READ_SIZE:
                 break
-
-        self.write(self, ''.join(buf))
+        return ''.join(buf)
 
 
 class EtherWebSocketHandler(tornado.websocket.WebSocketHandler, DebugMixIn):
-    def __init__(self, app, req, tap, debug=False):
+    def __init__(self, app, req, switch, debug=False):
         super(EtherWebSocketHandler, self).__init__(app, req)
-        self._tap = tap
+        self._switch = switch
         self._debug = debug
 
     def open(self):
-        self._tap.register_client(self)
+        self._switch.register_port(self)
         self.dprintf('connected: %s\n', lambda: self.request.remote_ip)
 
     def on_message(self, message):
-        self._tap.write(self, message)
+        self._switch.forward(self, EthernetFrame(message))
 
     def on_close(self):
-        self._tap.unregister_client(self)
+        self._switch.unregister_port(self)
         self.dprintf('disconnected: %s\n', lambda: self.request.remote_ip)
 
 
 class EtherWebSocketClient(DebugMixIn):
-    def __init__(self, tap, url, user=None, passwd=None, debug=False):
-        self._sock = None
-        self._tap = tap
+    def __init__(self, switch, url, user=None, passwd=None, debug=False):
+        self._switch = switch
         self._url = url
         self._debug = debug
+        self._sock = None
         self._options = {}
 
@@ -249 +288 @@
         self._sock = websocket.WebSocket()
         self._sock.connect(self._url, **self._options)
+        self._switch.register_port(self)
         self.dprintf('connected: %s\n', lambda: self._url)
 
@@ -254 +294 @@
         if self.closed:
             raise websocket.WebSocketException('already closed')
+        self._switch.unregister_port(self)
         self._sock.close()
         self._sock = None
@@ -276 +317 @@
             data = self._sock.recv()
             if data is not None:
-                self._tap.write(self, data)
+                self._switch.forward(self, EthernetFrame(data))
                 return
         except:
@@ -389 +430 @@
                               load_htpasswd(args.htpasswd))
 
-    tap = TapHandler(args.device, debug=args.debug)
+    fdb = FDB(debug=args.debug)
+    switch = SwitchingHub(fdb, debug=args.debug)
+    tap = TapHandler(switch, args.device, debug=args.debug)
     app = tornado.web.Application([
-        (args.path, handler, {'tap': tap, 'debug': args.debug}),
+        (args.path, handler, {'switch': switch, 'debug': args.debug}),
     ])
     server = tornado.httpserver.HTTPServer(app, ssl_options=ssl_options)
+
+    tap.open()
     server.listen(args.port, address=args.address)
 
@@ -419 +464 @@
         args.passwd = getpass.getpass()
 
-    tap = TapHandler(args.device, debug=args.debug)
-    client = EtherWebSocketClient(tap, args.uri,
+    fdb = FDB(debug=args.debug)
+    switch = SwitchingHub(fdb, debug=args.debug)
+    tap = TapHandler(switch, args.device, debug=args.debug)
+    client = EtherWebSocketClient(switch, args.uri,
                                   args.user, args.passwd, args.debug)
 
-    tap.register_client(client)
+    tap.open()
     client.open()