Count bytes bidirectionally between relays and dirauths

client.py

@@ -82,7 +82,7 @@ class Client:
         # Get a network address for client-side use only (do not bind it
         # to the network)
         self.netaddr = network.NetAddr()
-        self.perfstats = dirauth.PerfStats(dirauth.EntType.CLIENT)
+        self.perfstats = network.PerfStats(network.EntType.CLIENT)
         self.perfstats.name = "Client at %s" % self.netaddr
         self.perfstats.is_bootstrapping = True
         self.cellhandler = CellClient(self.netaddr, dirauthaddrs, self.perfstats)
@@ -123,7 +123,15 @@ class Client:
 
 
 if __name__ == '__main__':
-    perfstats = dirauth.PerfStats(dirauth.EntType.NONE)
+    perfstats = network.PerfStats(network.EntType.NONE)
+    totsent = 0
+    totrecv = 0
+    dirasent = 0
+    dirarecv = 0
+    relaysent = 0
+    relayrecv = 0
+    clisent = 0
+    clirecv = 0
 
     # Start some dirauths
     numdirauths = 9
@@ -201,6 +209,8 @@ if __name__ == '__main__':
         print("relay",r.descdict["addr"])
 
     relays[3].terminate()
+    relaysent += relays[3].perfstats.bytes_sent
+    relayrecv += relays[3].perfstats.bytes_received
     del relays[3]
 
     # Tick the epoch
@@ -219,19 +229,26 @@ if __name__ == '__main__':
 
     clients[0].cellhandler.new_circuit()
 
-totsent = 0
-totrecv = 0
-for d in dirauths:
-    print(d.perfstats)
-    totsent += d.perfstats.bytes_sent
-    totrecv += d.perfstats.bytes_received
-for r in relays:
-    print(r.perfstats)
-    totsent += r.perfstats.bytes_sent
-    totrecv += r.perfstats.bytes_received
-for c in clients:
-    print(c.perfstats)
-    totsent += c.perfstats.bytes_sent
-    totrecv += c.perfstats.bytes_received
-
-print("Total sent=%d recv=%d" % (totsent, totrecv))
+    for d in dirauths:
+        print(d.perfstats)
+        dirasent += d.perfstats.bytes_sent
+        dirarecv += d.perfstats.bytes_received
+    print("DirAuths sent=%d recv=%d" % (dirasent, dirarecv))
+    totsent += dirasent
+    totrecv += dirarecv
+    for r in relays:
+        print(r.perfstats)
+        relaysent += r.perfstats.bytes_sent
+        relayrecv += r.perfstats.bytes_received
+    print("Relays sent=%d recv=%d" % (relaysent, relayrecv))
+    totsent += relaysent
+    totrecv += relayrecv
+    for c in clients:
+        print(c.perfstats)
+        clisent += c.perfstats.bytes_sent
+        clirecv += c.perfstats.bytes_received
+    print("Client sent=%d recv=%d" % (clisent, clirecv))
+    totsent += clisent
+    totrecv += clirecv
+
+    print("Total sent=%d recv=%d" % (totsent, totrecv))

dirauth.py

@@ -7,46 +7,6 @@ import nacl.encoding
 import nacl.signing
 import network
 
-from enum import Enum
-
-class EntType(Enum):
-    """The different types of entities in the system."""
-    NONE = 0
-    DIRAUTH = 1
-    RELAY = 2
-    CLIENT = 3
-
-
-class PerfStats:
-    """A class to store performance statistics for a relay or client.
-    We keep track of bytes sent, bytes received, and counts of
-    public-key operations of various types.  We will reset these every
-    epoch."""
-
-    def __init__(self, ent_type):
-        # Which type of entity is this for (DIRAUTH, RELAY, CLIENT)
-        self.ent_type = ent_type
-        # A printable name for the entity
-        self.name = None
-        # True if bootstrapping this epoch
-        self.is_bootstrapping = False
-        # Bytes sent and received
-        self.bytes_sent = 0
-        self.bytes_received = 0
-        # Public-key operations: key generation, signing, verification,
-        # Diffie-Hellman
-        self.keygens = 0
-        self.sigs = 0
-        self.verifs = 0
-        self.dhs = 0
-
-    def __str__(self):
-        return "%s: type=%s boot=%s sent=%d recv=%d keygen=%d sig=%d verif=%d dh=%d" % \
-            (self.name, self.ent_type.name, self.is_bootstrapping, \
-            self.bytes_sent, self.bytes_received, self.keygens, \
-            self.sigs, self.verifs, self.dhs)
-
-
 # A relay descriptor is a dict containing:
 #  epoch: epoch id
 #  idkey: a public identity key
@@ -204,7 +164,7 @@ class DirAuthConnection(network.ClientConnection):
     """The subclass of Connection for connections to directory
     authorities."""
 
-    def __init__(self, peer = None):
+    def __init__(self, peer):
         super().__init__(peer)
 
     def uploaddesc(self, desc):
@@ -250,7 +210,7 @@ class DirAuth(network.Server):
         self.me = me
         self.tot = tot
         self.name = "Dirauth %d of %d" % (me+1, tot)
-        self.perfstats = PerfStats(EntType.DIRAUTH)
+        self.perfstats = network.PerfStats(network.EntType.DIRAUTH)
         self.perfstats.is_bootstrapping = True
 
         # Create the dirauth signature keypair
@@ -297,6 +257,7 @@ class DirAuth(network.Server):
         DirAuth.consensus.sign(self.sigkey, self.me, self.perfstats)
 
     def received(self, client, msg):
+        self.perfstats.bytes_received += msg.size()
         if isinstance(msg, DirAuthUploadDescMsg):
             # Check the uploaded descriptor for sanity
             epoch = msg.desc.descdict['epoch']
@@ -330,9 +291,15 @@ class DirAuth(network.Server):
                     (DirAuth.uploadeddescs[epoch][descstr][0]-1,
                      DirAuth.uploadeddescs[epoch][descstr][1])
         elif isinstance(msg, DirAuthGetConsensusMsg):
-            client.reply(DirAuthConsensusMsg(DirAuth.consensus))
+            replymsg = DirAuthConsensusMsg(DirAuth.consensus)
+            msgsize = replymsg.size()
+            self.perfstats.bytes_sent += msgsize
+            client.reply(replymsg)
         elif isinstance(msg, DirAuthGetENDIVEMsg):
-            client.reply(DirAuthENDIVEMsg(DirAuth.endive))
+            replymsg = DirAuthENDIVEMsg(DirAuth.endive)
+            msgsize = replymsg.size()
+            self.perfstats.bytes_sent += msgsize
+            client.reply(replymsg)
         else:
             raise TypeError('Not a client-originating DirAuthNetMsg', msg)
 

network.py

@@ -1,6 +1,7 @@
 #!/usr/bin/env python3
 
 import random
+import pickle
 from enum import Enum
 
 class WOMode(Enum):
@@ -17,6 +18,44 @@ class SNIPAuthMode(Enum):
     THRESHSIG   = 2  # Threshold signatures
 
 
+class EntType(Enum):
+    """The different types of entities in the system."""
+    NONE = 0
+    DIRAUTH = 1
+    RELAY = 2
+    CLIENT = 3
+
+
+class PerfStats:
+    """A class to store performance statistics for a relay or client.
+    We keep track of bytes sent, bytes received, and counts of
+    public-key operations of various types.  We will reset these every
+    epoch."""
+
+    def __init__(self, ent_type):
+        # Which type of entity is this for (DIRAUTH, RELAY, CLIENT)
+        self.ent_type = ent_type
+        # A printable name for the entity
+        self.name = None
+        # True if bootstrapping this epoch
+        self.is_bootstrapping = False
+        # Bytes sent and received
+        self.bytes_sent = 0
+        self.bytes_received = 0
+        # Public-key operations: key generation, signing, verification,
+        # Diffie-Hellman
+        self.keygens = 0
+        self.sigs = 0
+        self.verifs = 0
+        self.dhs = 0
+
+    def __str__(self):
+        return "%s: type=%s boot=%s sent=%d recv=%d keygen=%d sig=%d verif=%d dh=%d" % \
+            (self.name, self.ent_type.name, self.is_bootstrapping, \
+            self.bytes_sent, self.bytes_received, self.keygens, \
+            self.sigs, self.verifs, self.dhs)
+
+
 class NetAddr:
     """A class representing a network address"""
     nextaddr = 1
@@ -114,14 +153,16 @@ class Network:
         server.bound(addr, lambda: self.servers.pop(addr))
         return addr
 
-    def connect(self, client, srvaddr):
+    def connect(self, client, srvaddr, perfstats):
         """Connect the given client to the server bound to addr.  Throw
         an exception if there is no server bound to that address."""
         try:
             server = self.servers[srvaddr]
         except KeyError:
             raise NetNoServer()
-        return server.connected(client)
+        conn = server.connected(client)
+        conn.perfstats = perfstats
+        return conn
 
     def setfallbackrelays(self, fallbackrelays):
         """Set the list of globally known fallback relays.  Clients use
@@ -157,6 +198,16 @@ class NetMsg:
     """The parent class of network messages.  Subclass this class to
     implement specific kinds of network messages."""
 
+    def size(self):
+        """Return the size of this network message.  For now, just
+        pickle it and return the length of that.  There's some
+        unnecessary overhead in this method; if you want specific
+        messages to have more accurate sizes, override this method in
+        the subclass."""
+        sz = len(pickle.dumps(self))
+        print('size',sz,type(self))
+        return sz
+
 
 class StringNetMsg(NetMsg):
     """Send an arbitratry string as a NetMsg."""
@@ -193,17 +244,22 @@ class ClientConnection(Connection):
         """Create a ClientConnection object with the given peer.  The
         peer must have a received(client, msg) method."""
         self.peer = peer
+        self.perfstats = None
 
     def sendmsg(self, netmsg):
         assert(isinstance(netmsg, NetMsg))
+        msgsize = netmsg.size()
+        self.perfstats.bytes_sent += msgsize
         self.peer.received(self, netmsg)
 
     def reply(self, netmsg):
         assert(isinstance(netmsg, NetMsg))
+        msgsize = netmsg.size()
+        self.perfstats.bytes_received += msgsize
         self.receivedfromserver(netmsg)
 
-    def received(self, netmsg):
-        print("received", netmsg, "from server")
+    #def received(self, netmsg):
+        #print("received", netmsg, "from server")
 
 
 class ServerConnection(Connection):

relay.py

@@ -194,7 +194,8 @@ 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)
         newchannel.cellhandler = self
@@ -231,7 +232,7 @@ class CellRelay(CellHandler):
     def get_consensus(self):
         """Download a fresh consensus from a random dirauth."""
         a = random.choice(self.dirauthaddrs)
-        c = network.thenetwork.connect(self, a)
+        c = network.thenetwork.connect(self, a, self.perfstats)
         self.consensus = c.getconsensus()
         dirauth.Consensus.verify(self.consensus, \
                 network.thenetwork.dirauthkeys(), self.perfstats)
@@ -263,7 +264,7 @@ class Relay(network.Server):
 
     def __init__(self, dirauthaddrs, bw, flags):
         # Gather performance statistics
-        self.perfstats = dirauth.PerfStats(dirauth.EntType.RELAY)
+        self.perfstats = network.PerfStats(network.EntType.RELAY)
         self.perfstats.is_bootstrapping = True
 
         # Create the identity and onion keys
@@ -351,7 +352,7 @@ class Relay(network.Server):
 
         # Upload them
         for a in self.cellhandler.dirauthaddrs:
-            c = network.thenetwork.connect(self, a)
+            c = network.thenetwork.connect(self, a, self.perfstats)
             c.sendmsg(descmsg)
             c.close()
 
@@ -381,7 +382,7 @@ class Relay(network.Server):
         return peerchannel
 
 if __name__ == '__main__':
-    perfstats = dirauth.PerfStats(dirauth.EntType.NONE)
+    perfstats = network.PerfStats(network.EntType.NONE)
 
     # Start some dirauths
     numdirauths = 9