p2pool/util/switchprotocol.py

   1 from twisted.internet import protocol
   2
   3 class FirstByteSwitchProtocol(protocol.Protocol):
   4     p = None
   5     def dataReceived(self, data):
   6         if self.p is None:
   7             if not data: return
   8             serverfactory = self.factory.first_byte_to_serverfactory.get(data[0], self.factory.default_serverfactory)
   9             self.p = serverfactory.buildProtocol(self.transport.getPeer())
  10             self.p.makeConnection(self.transport)
  11         self.p.dataReceived(data)
  12     def connectionLost(self, reason):
  13         if self.p is not None:
  14             self.p.connectionLost(reason)
  15
  16 class FirstByteSwitchFactory(protocol.ServerFactory):
  17     protocol = FirstByteSwitchProtocol
  18
  19     def __init__(self, first_byte_to_serverfactory, default_serverfactory):
  20         self.first_byte_to_serverfactory = first_byte_to_serverfactory
  21         self.default_serverfactory = default_serverfactory
  22
  23     def startFactory(self):
  24         for f in list(self.first_byte_to_serverfactory.values()) + [self.default_serverfactory]:
  25             f.doStart()
  26
  27     def stopFactory(self):
  28         for f in list(self.first_byte_to_serverfactory.values()) + [self.default_serverfactory]:
  29             f.doStop()