[WIP] BEP10 - Extension protocol

BitTornado/Client/Connecter.py

@@ -1,24 +1,32 @@
+from .Extension import PeerExtensions
+from ..Meta.bencode import bdecode
 from ..Types import Bitfield
 from BitTornado.clock import clock
 
 DEBUG1 = False
 DEBUG2 = False
 
 
-CHOKE = b'\x00'
-UNCHOKE = b'\x01'
-INTERESTED = b'\x02'
-NOT_INTERESTED = b'\x03'
-# index
-HAVE = b'\x04'
-# index, bitfield
-BITFIELD = b'\x05'
-# index, begin, length
-REQUEST = b'\x06'
-# index, begin, piece
-PIECE = b'\x07'
-# index, begin, piece
-CANCEL = b'\x08'
+# Message IDs
+CHOKE = b'\x00'             # no payload
+UNCHOKE = b'\x01'           # no payload
+INTERESTED = b'\x02'        # no payload
+NOT_INTERESTED = b'\x03'    # no payload
+HAVE = b'\x04'              # index
+BITFIELD = b'\x05'          # index, bitfield
+REQUEST = b'\x06'           # index, begin, length
+PIECE = b'\x07'             # index, begin, piece
+CANCEL = b'\x08'            # index, begin, piece
+DHT_PORT = b'\x09'          # port (2 bytes)
+FAST_SUGGEST = b'\x0d'      # index
+FAST_HAVE_ALL = b'\x0e'     # no payload
+FAST_HAVE_NONE = b'\x0f'    # no payload
+FAST_REJECT = b'\x10'       # index, begin, length
+ALLOW_FAST = b'\x11'        # index
+EXTENDED = b'\x14'          # msg_id, payload
+
+# Extended Message IDs
+EXT_HANDSHAKE = b'\x00'
 
 
 class Connection(object):
@@ -34,6 +42,7 @@ def __init__(self, connection, connecter, ccount):
         self.upload = None              # Uploader.Upload
         self.send_choke_queued = False  # Bool (togglable)
         self.just_unchoked = None       # None -> 0 <-> clock()
+        self.supported_exts = {}
 
         # Pass-through functions
         self.get_id = connection.get_id
@@ -97,6 +106,9 @@ def send_bitfield(self, bitfield):
     def send_have(self, index):
         self._send_message(HAVE + index.to_bytes(4, 'big'))
 
+    def send_extended(self, ext_id, payload):
+        self._send_message(EXTENDED + ext_id.to_bytes(1, 'big') + payload)
+
     def send_keepalive(self):
         self._send_message(b'')
 
@@ -167,6 +179,17 @@ def got_request(self, piece_num, pos, length):
             self.connecter.ratelimiter.ping(clock() - self.just_unchoked)
             self.just_unchoked = 0
 
+    def got_extended(self, message):
+        message_id = message[:1]
+        if message_id == EXT_HANDSHAKE:
+            payload = bdecode(message[1:])
+            supported_exts = payload['m']
+            self.supported_exts.update((k, v)
+                                       for k, v in supported_exts.items()
+                                       if k in SUPPORTED_EXTS)
+        else:
+            pass
+
 
 class Connecter(object):
     def __init__(self, make_upload, downloader, choker, numpieces, totalup,
@@ -311,5 +334,7 @@ def got_message(self, connection, message):
             if c.download.got_piece(i, int.from_bytes(message[5:9], 'big'),
                                     message[9:]):
                 self.got_piece(i)
+        elif t == EXTENDED:
+            c.got_extended(message)
         else:
             connection.close()

BitTornado/Client/Extension.py

@@ -0,0 +1,127 @@
+from BitTornado.Meta.bencode import bencode, bdecode
+# Design
+#
+# Fundamental:
+# extension_name, msg_id, handler
+#
+# ext_table[msg_id] = (extension_name, handler)
+# extensions[extension_name] = msg_id
+#
+# Ours:
+#   msg_id = index
+#   ext_table = [(extension_name, handler)]
+#
+#   ext_lookup = {extension_name: msg_id}
+#
+# Theirs:
+#   msg_id = key
+#   handlers =   {msg_id: (extension_name, handler)}
+#   extensions = {extension_name: msg_id}
+#
+#   Default:
+#   handlers = {0: (handshake, handshake_handler)}
+#   extensions = {handshake: 0}
+
+
+class _SupportedExtensions(object):
+    def __init__(self):
+        self.table = [Handshake(None)]
+        self.lookup = {}
+
+    def register(self, handler):
+        """Register a handler (may be used as decorator)"""
+        assert isinstance(handler.name, (bytes, str))
+        if handler.name in self.lookup:
+            # Previously registered, do nothing
+            if self.table[self.lookup[handler.name]] == handler:
+                return
+
+            # Previously unregistered, restore
+            try:
+                msg_id = self.table.index(handler, 1)
+                self.lookup[handler.name] = msg_id
+            except ValueError:
+                pass
+
+        # New handler (possibly replacing existing handler)
+        self.lookup[handler.name] = len(self.table)
+        self.table.append(handler)
+        return handler
+
+    def unregister(self, item):
+        if isinstance(item, int):
+            msg_id = item
+            handler = self.table[item]
+        elif isinstance(item, ExtensionHandler):
+            handler = item
+            try:
+                msg_id = self.table.index(item, 1)
+            except ValueError:
+                raise ValueError("Handler not registered!")
+
+EXTENSIONS = _SupportedExtensions()
+
+
+class ExtensionHandler(object):
+    name = None
+
+    def prepare(self):
+        raise NotImplementedError
+
+    def receive(self, payload):
+        raise NotImplementedError
+
+
+class Handshake(ExtensionHandler):
+    """Handle the BEP 10 handshake, registering mutually supported handshakes
+    for a peer.
+    """
+    def __init__(self, extensions=None):
+        self.exts = extensions
+        self.payload = {'m': ext_lookup}
+
+    def prepare(self):
+        return bencode(self.payload)
+
+    def receive(self, payload):
+        if self.exts is None:
+            raise ValueError('This Handshake is not associated with a client.')
+        message = bdecode(payload)
+        to_add = []
+        to_remove = []
+
+        for ext_name, msg_id in message['m']:
+            if msg_id == 0:
+                to_remove.append(ext_name)
+            elif ext_name not in ext_table:
+                continue
+            elif ext_name not in self.exts.extensions:
+                to_add.append((ext_name, msg_id))
+            elif self.exts.extensions[ext_name] != msg_id:
+                to_remove.append(ext_name)
+                to_add.append((ext_name, msg_id))
+
+        if len(to_add) + len(to_remove) == 0:
+            return
+
+        handlers = self.exts.handlers.copy()
+        extensions = self.exts.extensions.copy()
+
+        for ext_name in to_remove:
+            old_id = extensions[ext_name]
+            assert handlers[old_id].name == ext_name
+            del handlers[old_id]
+            del extensions[ext_name]
+
+        for ext_name, msg_id in to_add:
+            extensions[ext_name] = msg_id
+            handlers[msg_id] = ext_table[ext_lookup[ext_name]]
+
+
+class PeerExtensions(object):
+    def __init__(self):
+        self.handlers = {0: Handshake(self)}
+        self.extensions = {}
+
+    def __getitem__(self, key):
+        return self.handlers[key]

BitTornado/Network/Encrypter.py

@@ -1,14 +1,12 @@
 import urllib
 from binascii import hexlify
 from .BTcrypto import Crypto, padding
+from .Protocol import protocol_name, reserved
 
 DEBUG = False
 
 MAX_INCOMPLETE = 8
 
-protocol_name = b'\x13BitTorrent protocol'
-option_pattern = bytes(8)
-
 
 def make_readable(s):
     if not s:
@@ -63,7 +61,7 @@ def __init__(self, Encoder, connection, peerid,
                 self.write(self.encrypter.padded_pubkey())
             else:
                 self.encrypted = False
-                self.write(protocol_name + option_pattern +
+                self.write(protocol_name + bytes(reserved) +
                            self.Encoder.download_id)
             self.next_len = len(protocol_name)
             self.next_func = self.read_header
@@ -285,7 +283,7 @@ def read_crypto_block4done(self):
             if not self.buffer:  # oops; check for exceptions to this
                 return None
             self._end_crypto()
-        self.write(protocol_name + option_pattern + self.Encoder.download_id)
+        self.write(protocol_name + bytes(reserved) + self.Encoder.download_id)
         return len(protocol_name), self.read_encrypted_header
 
     ### START PROTOCOL OVER ENCRYPTED CONNECTION ###
@@ -306,7 +304,7 @@ def read_download_id(self, s):
         if not self.locally_initiated:
             if not self.encrypted:
                 self.Encoder.connecter.external_connection_made += 1
-            self.write(protocol_name + option_pattern +
+            self.write(protocol_name + bytes(reserved) +
                        self.Encoder.download_id + self.Encoder.my_id)
         return 20, self.read_peer_id
 

BitTornado/Network/NatCheck.py

@@ -1,5 +1,5 @@
 from .BTcrypto import Crypto, CRYPTO_OK, padding
-from .Encrypter import protocol_name, option_pattern
+from .Protocol import protocol_name, reserved
 
 CHECK_PEER_ID_ENCRYPTED = True
 
@@ -126,7 +126,7 @@ def read_crypto_block4done(self):
             if not self.buffer:  # oops; check for exceptions to this
                 return None
             self._end_crypto()
-        self.write(protocol_name + option_pattern + self.Encoder.download_id)
+        self.write(protocol_name + bytes(reserved) + self.Encoder.download_id)
         return len(protocol_name), self.read_encrypted_header
 
     ### START PROTOCOL OVER ENCRYPTED CONNECTION ###

BitTornado/Network/Protocol.py

@@ -0,0 +1,21 @@
+from ..Types import Bitfield
+
+protocol_name = b'\x13BitTorrent protocol'
+reserved = Bitfield(64)
+
+# Flags (See: http://www.bittorrent.org/beps/bep_0004.html)
+# Byte Index: 0 1 2 3 4 5 6 7
+DIST_HASH = 0x0000000000000001  # BEP 5: Distributed Hash Table
+XBT_PEERX = 0x0000000000000002  # XBT Peer Exchange
+FAST_EXTN = 0x0000000000000004  # BEP 6: Fast Extension
+NAT_PUNCH = 0x0000000000000008  # NAT Traversal
+# Byte Index: 0 1 2 3 4 5 6 7
+EXT_NEG_1 = 0x0000000000010000  # Extension Negotiation Protocol (Deprecated)
+EXT_NEG_2 = 0x0000000000020000  # Extension Negotiation Protocol (Deprecated)
+EXTENSION = 0x0000000000100000  # BEP 10: Libtorrent Extension Protocol
+# Byte Index: 0 1 2 3 4 5 6 7
+LOC_AWARE = 0x0008000000000000  # BitTorrent Location-aware Protocol
+COMET_BYT = 0x00ff000000000000  # BitComet Extension Protocol
+# Byte Index: 0 1 2 3 4 5 6 7
+AZURE_MSG = 0x8000000000000000  # Azureus Messaging Protocol
+COMET_BYT = 0xff00000000000000  # BitComet Extension Protocol

BitTornado/Network/ServerPortHandler.py

@@ -1,6 +1,5 @@
 from .BTcrypto import Crypto
-
-from .Encrypter import protocol_name
+from .Protocol import protocol_name
 
 
 class SingleRawServer(object):
@@ -75,7 +74,7 @@ def __init__(self, multihandler, connection):
         self.complete = False
         self.read = self._read
         self.write = connection.write
-        self.next_len = 1 + len(protocol_name)
+        self.next_len = len(protocol_name)
         self.next_func = self.read_header
         self.multihandler.rawserver.add_task(self._auto_close, 30)