diff --git a/src/wormhole/relay.py b/src/wormhole/relay.py index 8e62e28..6a5fc6e 100644 --- a/src/wormhole/relay.py +++ b/src/wormhole/relay.py @@ -1,9 +1,13 @@ -import re, json +import os, re, json, binascii from collections import defaultdict from twisted.python import log +from twisted.internet import protocol from twisted.application import strports, service from twisted.web import server, static, resource, http +SECONDS = 1.0 +MB = 1000*1000 + class Channel(resource.Resource): isLeaf = True @@ -97,6 +101,128 @@ class Relay(resource.Resource): log.msg("freed %d, now have %d channels" % (channel_id, len(self.channels))) +class Transit(resource.Resource, protocol.ServerFactory, service.MultiService): + # Transit manages two simultaneous connections to a secondary TCP port, + # both forwarded to the other. Transit will allocate you a token when the + # ports are free, and will inform you of the MAXLENGTH and MAXTIME + # limits. Connect to the port, send "TOKEN\n", receive "ok\n", and all + # subsequent data you send will be delivered to the other side. All data + # you get after the "ok" will be from the other side. You will not + # receive "ok" until the other side has also connected and submitted a + # valid token. The token is different for each side. The connections will + # be dropped after MAXLENGTH bytes have been sent by either side, or + # MAXTIME seconds after the token is issued, whichever is reached first. + + # These relay connections are not half-closeable (unlike full TCP + # connections, applications will not receive any data after half-closing + # their outgoing side). Applications must negotiate shutdown with their + # peer and not close the connection until all data has finished + # transferring in both directions. Applications which only need to send + # data in one direction can use close as usual. + + MAXLENGTH = 10*MB + MAXTIME = 60*SECONDS + + def __init__(self): + resource.Resource.__init__(self) + service.MultiService.__init__(self) + self.pending_requests = [] + self.active_token = None + self.active_connection = None + self.active_timer = None + + def make_token(self): + return binascii.hexlify(os.urandom(8)) + + def render_POST(self, request): + if self.active_connection: + self.pending_requests.append(request) + return server.NOT_DONE_YET + self.active_token = self.make_token() + request.setHeader("content-type", "application/json; charset=utf-8") + t = service.TimerService(self.MAXTIME, self.timer_expired) + self.active_timer = t + t.setServiceParent(self) + r = { "token": self.active_token, + "maxlength": self.MAXLENGTH, + "maxtime": self.MAXTIME, + } + return json.dumps(r)+"\n" + + def timer_expired(self): + self.remove_timer() + for c in self.active_connections: + c.STOPSTOP() + self.active_connections[:] = [] + self.active_token = None + + def remove_timer(self): + self.active_timer.disownServiceParent() + self.active_timer = None + + # ServerFactory methods, which manage the two TransitConnection protocols + + def buildProtocol(self, addr): + p = TransitConnection(self.active_token) + p.factory = self + return p + + def connection_got_token(self, p): + pass + def transitFinished(self, p): + pass + def transitFailed(self): + pass + +# after getting a token, both transit clients connect to one of these + +class TransitConnection(protocol.Protocol): + def __init__(self, expected_token): + self.expected_token = expected_token + self.got_token = False + self.token_buffer = b"" + self.sent_ok = False + self.buddy = None + + def dataReceived(self, data): + if self.sent_ok: + # TODO: connect as producer/consumer + self.buddy.transport.write(data) + return + if self.got_token: # but not yet sent_ok + return self.disconnect() # impatience yields failure + # else this should be (part of) the token + self.token_buffer += data + if b"\n" not in self.token_buffer: + return + lines = self.token_buffer.split(b"\n") + if len(lines) > 1: + return self.disconnect() # impatience yields failure + token = lines[0] + if token != self.expected_token: + return self.disconnect() # incorrectness yields failure + self.got_token = True + self.factory.connection_got_token(self) + + def buddy_connected(self, them): + self.buddy = them + self.transport.write(b"ok\n") + self.sent_ok = True + + def buddy_disconnected(self): + self.buddy = None + self.transport.loseConnection() + self.factory.transitFinished(self) + + def connectionLost(self, reason): + if self.buddy: + self.buddy.buddy_disconnected() + + def disconnect(self): + self.transport.loseConnection() + self.factory.transitFailed() + + class Root(resource.Resource): # child_FOO is a nevow thing, not a twisted.web.resource thing def __init__(self): @@ -104,17 +230,19 @@ class Root(resource.Resource): self.putChild("", static.Data("Wormhole Relay\n", "text/plain")) class RelayServer(service.MultiService): - def __init__(self, listenport): + def __init__(self, relayport, transitport): service.MultiService.__init__(self) self.root = Root() site = server.Site(self.root) - self.port_service = strports.service(listenport, site) - self.port_service.setServiceParent(self) + self.relayport_service = strports.service(relayport, site) + self.relayport_service.setServiceParent(self) self.relay = Relay() # for tests self.root.putChild("relay", self.relay) - - def get_root(self): - return self.root + self.transit = Transit() + self.root.putChild("transit", self.transit) + self.transit.setServiceParent(self) # for the timer + self.transport_service = strports.service(transitport, self.transit) + self.transport_service.setServiceParent(self) application = service.Application("foo") -RelayServer("tcp:8009").setServiceParent(application) +RelayServer("tcp:8009", "tcp:8010").setServiceParent(application)