Handle EXTEND/EXTENDED cells. Circuits can now be built successfully!

client.py

@@ -27,10 +27,12 @@ class VanillaCreatedExtendedHandler:
         circhandler.add_crypt_layer(enckey, deckey)
         if len(circhandler.circuit_descs) == 0:
             # This was a VanillaCreatedCircuitCell
-            del circhandler.cell_dispatch_table[relay.VanillaCreatedCircuitCell]
+            circhandler.replace_celltype_handler( \
+                    relay.VanillaCreatedCircuitCell, None)
         else:
             # This was a VanillaExtendedCircuitCell
-            del circhandler.cell_dispatch_table[relay.VanillaExtendedCircuitCell]
+            circhandler.replace_celltype_handler( \
+                    relay.VanillaExtendedCircuitCell, None)
         circhandler.circuit_descs.append(self.expecteddesc)
 
         # Are we done building the circuit?
@@ -42,7 +44,9 @@ class VanillaCreatedExtendedHandler:
         while nexthop is None:
             nexthop = self.cellhandler.consensus.select_weighted_relay( \
                     self.cellhandler.consensus_cdf)
-            if nexthop.descdict['addr'] in circhandler.circuit_descs:
+            if nexthop.descdict['addr'] in \
+                    [ desc.descdict['addr'] \
+                        for desc in circhandler.circuit_descs ]:
                 nexthop = None
 
         # Construct the VanillaExtendCircuitCell
@@ -52,8 +56,9 @@ class VanillaCreatedExtendedHandler:
                 nexthop.descdict['addr'], ntor_request)
 
         # Set up the reply handler
-        circhandler.cell_dispatch_table[relay.VanillaExtendedCircuitCell] = \
-                VanillaCreatedExtendedHandler(self, ntor, nexthop)
+        circhandler.replace_celltype_handler( \
+                relay.VanillaExtendedCircuitCell, \
+                VanillaCreatedExtendedHandler(self.cellhandler, ntor, nexthop))
 
         # Send the cell
         circhandler.send_cell(circextendmsg)
@@ -117,8 +122,9 @@ class CellClient(relay.CellHandler):
         circcreatemsg = relay.VanillaCreateCircuitMsg(circid, ntor_request)
 
         # Set up the reply handler
-        circhandler.cell_dispatch_table[relay.VanillaCreatedCircuitCell] = \
-                VanillaCreatedExtendedHandler(self, ntor, self.guard)
+        circhandler.replace_celltype_handler( \
+                relay.VanillaCreatedCircuitCell, \
+                VanillaCreatedExtendedHandler(self, ntor, self.guard))
 
         # Send the message
         guardchannel.send_msg(circcreatemsg)
@@ -290,7 +296,7 @@ if __name__ == '__main__':
 
     clients[0].cellhandler.new_circuit()
     circhandler = clients[0].cellhandler.channels[clients[0].cellhandler.guardaddr].circuithandlers[2]
-    circhandler.send_cell(network.StringNetMsg("hello world"))
+    circhandler.send_cell(relay.StringCell("hello world"))
 
     # See what channels exist and do a consistency check
     for r in relays:

network.py

@@ -230,8 +230,8 @@ class NetMsg:
 
 class StringNetMsg(NetMsg):
     """Send an arbitratry string as a NetMsg."""
-    def __init__(self, str):
-        self.data = str
+    def __init__(self, data):
+        self.data = data
 
     def __str__(self):
         return self.data.__str__()

relay.py

@@ -58,6 +58,15 @@ class RelayCell(RelayNetMsg):
     need their own circuitid) should be a subclass of this class."""
 
 
+class StringCell(RelayCell):
+    """Send an arbitrary string as a cell."""
+    def __init__(self, str):
+        self.data = str
+
+    def __str__(self):
+        return self.data.__str__()
+
+
 class VanillaCreateCircuitMsg(RelayNetMsg):
     """The message for requesting circuit creation in Vanilla Onion
     Routing."""
@@ -145,8 +154,7 @@ class NTor:
                 server_ephem_key.public_key.encode(encoder=nacl.encoding.RawEncoder)
         A = nacl.hash.sha256(M + b'verify', encoder=nacl.encoding.RawEncoder)
         S = nacl.hash.sha256(M + b'secret', encoder=nacl.encoding.RawEncoder)
-        return ((server_ephem_key.public_key, onion_privkey.public_key, A), \
-                S)
+        return ((server_ephem_key.public_key, onion_privkey.public_key, A), S)
 
     def verify(self, reply, onion_pubkey, idpubkey):
         """The client calls this method to verify the ntor reply
@@ -169,6 +177,54 @@ class NTor:
             raise ValueError("NTor auth mismatch")
         return S
 
+
+class VanillaExtendCircuitHandler:
+    """A handler for VanillaExtendCircuitCell cells.  It allocates a new
+    circuit id on the Channel to the requested next hop, connects the
+    existing and new circuits together, and forwards a
+    VanillaCreateCircuitMsg to the next hop."""
+
+    def received_cell(self, circhandler, cell):
+        # Remove ourselves from handling a second
+        # VanillaExtendCircuitCell on this circuit
+        circhandler.replace_celltype_handler(VanillaExtendCircuitCell, None)
+
+        # Allocate a new circuit id to the requested next hop
+        cellhandler = circhandler.channel.cellhandler
+        nexthopchannel = cellhandler.get_channel_to(cell.hopaddr)
+        newcircid, newcirchandler = nexthopchannel.new_circuit()
+
+        # Connect the existing and new circuits together
+        circhandler.adjacent_circuit_handler = newcirchandler
+        newcirchandler.adjacent_circuit_handler = circhandler
+
+        # Set up a handler for when the VanillaCreatedCircuitCell comes
+        # back
+        newcirchandler.replace_celltype_handler(
+                VanillaCreatedCircuitCell,
+                VanillaCreatedRelayHandler())
+
+        # Forward a VanillaCreateCircuitMsg to the next hop
+        nexthopchannel.send_msg(
+                VanillaCreateCircuitMsg(newcircid, cell.ntor_request))
+
+
+class VanillaCreatedRelayHandler:
+    """Handle a VanillaCreatedCircuitCell received by a _relay_ that
+    recently received a VanillaExtendCircuitCell from a client, and so
+    forwarded a VanillaCreateCircuitCell to the next hop."""
+
+    def received_cell(self, circhandler, cell):
+        # Remove ourselves from handling a second
+        # VanillaCreatedCircuitCell on this circuit
+        circhandler.replace_celltype_handler(VanillaCreatedCircuitCell, None)
+
+        # Just forward a VanillaExtendedCircuitCell back towards the
+        # client
+        circhandler.adjacent_circuit_handler.send_cell(
+            VanillaExtendedCircuitCell(cell.ntor_reply))
+
+
 class CircuitHandler:
     """A class for managing sending and receiving encrypted cells on a
     particular circuit."""
@@ -203,6 +259,10 @@ class CircuitHandler:
         # furthest away relay in the circuit) and are decrypted starting
         # with the first layer that was added (the keys for the guard).
         self.crypt_layer = self.NoCryptLayer()
+        # The adjacent CircuitHandler that's connected to this one.  If
+        # we get a cell on one, we forward it to the other (if it's not
+        # meant for us to handle directly)
+        self.adjacent_circuit_handler = None
 
     def send_cell(self, cell):
         """Send a cell on this circuit, encrypting as needed."""
@@ -221,9 +281,36 @@ class CircuitHandler:
 
         print("CircuitHandler: %s received cell %s on circuit %d from %s" % (self.channel.cellhandler.myaddr, cell, self.circid, self.channel.peer.cellhandler.myaddr))
 
-        celltype = type(cell)
-        if celltype in self.cell_dispatch_table:
-            self.cell_dispatch_table[celltype].received_cell(self, cell)
+        # If it's still encrypted, it's for sure meant for forwarding to
+        # our adjacent hop, which had better exist.
+        if isinstance(cell, EncryptedCell):
+            self.adjacent_circuit_handler.send_cell(cell)
+        else:
+            # This is a plaintext cell meant for us.  Handle it
+            # according to the table.
+            celltype = type(cell)
+            if celltype in self.cell_dispatch_table:
+                self.cell_dispatch_table[celltype].received_cell(self, cell)
+            elif isinstance(cell, StringCell):
+                # Default handler; just print the message in the cell
+                print("CircuitHandler: %s received '%s' on circuit %d from %s" \
+                        % (self.channel.cellhandler.myaddr, cell,
+                            self.circid, self.channel.peer.cellhandler.myaddr))
+            else:
+                # I don't know how to handle this cell?
+                raise ValueError("CircuitHandler: %s received unknown cell type %s on circuit %d from %s" \
+                        % (self.channel.cellhandler.myaddr, cell,
+                            self.circid, self.channel.peer.cellhandler.myaddr))
+
+    def replace_celltype_handler(self, celltype, handler):
+        """Add an object with a received_cell(circhandler, cell) method
+        to the cell dispatch table.  It replaces anything that's already
+        there.  Passing None as the handler removes the dispatcher for
+        that cell type."""
+        if handler is None:
+            del self.cell_dispatch_table[celltype]
+        else:
+            self.cell_dispatch_table[celltype] = handler
 
     def add_crypt_layer(self, enckey, deckey):
         """Add a processing layer to this CircuitHandler so that cells
@@ -343,7 +430,7 @@ class CellHandler:
             return self.channels[addr]
 
         # Create the new channel
-        newchannel = network.thenetwork.connect(self.myaddr, addr, \
+        newchannel = network.thenetwork.connect(self.myaddr, addr,
                 self.perfstats)
         self.channels[addr] = newchannel
         newchannel.closer = lambda: self.channels.pop(addr)
@@ -385,7 +472,7 @@ class CellRelay(CellHandler):
         a = random.choice(self.dirauthaddrs)
         c = network.thenetwork.connect(self, a, self.perfstats)
         self.consensus = c.getconsensus()
-        dirauth.Consensus.verify(self.consensus, \
+        dirauth.Consensus.verify(self.consensus,
                 network.thenetwork.dirauthkeys(), self.perfstats)
         c.close()
 
@@ -405,14 +492,18 @@ class CellRelay(CellHandler):
             # A new circuit has arrived
             circhandler = channel.new_circuit_with_circid(msg.circid)
             # Create the ntor reply
-            reply, secret = NTor.reply(self.onionkey, self.idpubkey, \
+            reply, secret = NTor.reply(self.onionkey, self.idpubkey,
                     msg.ntor_request, self.perfstats)
             # Set up the circuit to use the shared secret
             enckey = nacl.hash.sha256(secret + b'downstream')
             deckey = nacl.hash.sha256(secret + b'upstream')
             circhandler.add_crypt_layer(enckey, deckey)
+            # Add a handler for if an Extend Cell arrives (there should
+            # be at most one on this circuit).
+            circhandler.replace_celltype_handler(
+                    VanillaExtendCircuitCell, VanillaExtendCircuitHandler())
             # Send the ntor reply
-            self.send_msg(CircuitCellMsg(msg.circid, \
+            self.send_msg(CircuitCellMsg(msg.circid,
                     VanillaCreatedCircuitCell(reply)), peeraddr)
         else:
             return super().received_msg(msg, peeraddr, channel)
@@ -450,7 +541,7 @@ class Relay(network.Server):
         network.thenetwork.wantepochticks(self, True)
 
         # Create the CellRelay connection manager
-        self.cellhandler = CellRelay(self.netaddr, dirauthaddrs, \
+        self.cellhandler = CellRelay(self.netaddr, dirauthaddrs,
                 self.onionkey, self.idkey.verify_key, self.perfstats)
 
         # Initially, we're not a fallback relay
@@ -579,7 +670,7 @@ if __name__ == '__main__':
     # Tick the epoch
     network.thenetwork.nextepoch()
 
-    dirauth.Consensus.verify(dirauth.DirAuth.consensus, \
+    dirauth.Consensus.verify(dirauth.DirAuth.consensus,
             network.thenetwork.dirauthkeys(), perfstats)
 
     print('ticked; epoch=', network.thenetwork.getepoch())
@@ -625,7 +716,7 @@ if __name__ == '__main__':
     circid, circhandler = channel.new_circuit()
     peerchannel = relays[5].cellhandler.get_channel_to(relays[3].netaddr)
     peerchannel.new_circuit_with_circid(circid)
-    relays[3].cellhandler.send_cell(circid, network.StringNetMsg("test"), relays[5].netaddr)
+    relays[3].cellhandler.send_cell(circid, StringCell("test"), relays[5].netaddr)
 
     idpubkey = dirauth.DirAuth.consensus.consdict["relays"][1].descdict["idkey"]
     onionpubkey = dirauth.DirAuth.consensus.consdict["relays"][1].descdict["onionkey"]