diff --git a/src/wormhole_transit_relay/test/test_transit_server.py b/src/wormhole_transit_relay/test/test_transit_server.py index f07b275..325ee04 100644 --- a/src/wormhole_transit_relay/test/test_transit_server.py +++ b/src/wormhole_transit_relay/test/test_transit_server.py @@ -396,6 +396,23 @@ class _Transit: a1.transport.loseConnection() + @defer.inlineCallbacks + def test_short_handshake(self): + ep = clientFromString(reactor, self.transit) + a1 = yield connectProtocol(ep, Accumulator()) + + # hang up before sending a complete handshake + a1.transport.write(b"short") + a1.transport.loseConnection() + + @defer.inlineCallbacks + def test_empty_handshake(self): + ep = clientFromString(reactor, self.transit) + a1 = yield connectProtocol(ep, Accumulator()) + + # hang up before sending anything + a1.transport.loseConnection() + class TransitWithLogs(_Transit, ServerBase, unittest.TestCase): log_requests = True @@ -412,6 +429,45 @@ class Usage(ServerBase, unittest.TestCase): total_time, waiting_time)) self._transit_server.recordUsage = record + @defer.inlineCallbacks + def test_empty(self): + ep = clientFromString(reactor, self.transit) + a1 = yield connectProtocol(ep, Accumulator()) + + # hang up before sending anything + a1.transport.loseConnection() + yield a1._disconnect + # give the server a chance to react. in most of the other tests, the + # server hangs up on us, so this test needs extra synchronization + while len(self._usage) == 0: + d = defer.Deferred() + reactor.callLater(0.01, d.callback, None) + yield d + # that will log the "empty" usage event + self.assertEqual(len(self._usage), 1, self._usage) + (started, result, total_bytes, total_time, waiting_time) = self._usage[0] + self.assertEqual(result, "empty", self._usage) + + @defer.inlineCallbacks + def test_short(self): + ep = clientFromString(reactor, self.transit) + a1 = yield connectProtocol(ep, Accumulator()) + + # hang up before sending a complete handshake + a1.transport.write(b"short") + a1.transport.loseConnection() + yield a1._disconnect + # give the server a chance to react. in most of the other tests, the + # server hangs up on us, so this test needs extra synchronization + while len(self._usage) == 0: + d = defer.Deferred() + reactor.callLater(0.01, d.callback, None) + yield d + # that will log the "empty" usage event + self.assertEqual(len(self._usage), 1, self._usage) + (started, result, total_bytes, total_time, waiting_time) = self._usage[0] + self.assertEqual(result, "empty", self._usage) + @defer.inlineCallbacks def test_errory(self): ep = clientFromString(reactor, self.transit) diff --git a/src/wormhole_transit_relay/transit_server.py b/src/wormhole_transit_relay/transit_server.py index 1c7f0e2..4075979 100644 --- a/src/wormhole_transit_relay/transit_server.py +++ b/src/wormhole_transit_relay/transit_server.py @@ -29,7 +29,7 @@ class TransitConnection(protocol.Protocol): self._got_side = False self._token_buffer = b"" self._sent_ok = False - self._mood = None + self._mood = "empty" self._buddy = None self._total_sent = 0 @@ -182,7 +182,8 @@ class TransitConnection(protocol.Protocol): finished = time.time() total_time = finished - self._started - # Record usage. There are seven cases: + # Record usage. There are eight cases: + # * n0: we haven't gotten a full handshake yet (empty) # * n1: the handshake failed, not a real client (errory) # * n2: real client disconnected before any buddy appeared (lonely) # * n3: real client closed as redundant after buddy appears (redundant) @@ -191,11 +192,15 @@ class TransitConnection(protocol.Protocol): # * n6: real client connected last, buddy closes first (jilted) # * n7: real client connected last, buddy closes last (happy) - # * non-connected clients (1,2,3) always write a usage record + # * non-connected clients (0,1,2,3) always write a usage record # * for connected clients, whoever disconnects first gets to write the # usage record (5, 7). The last disconnect doesn't write a record. - if self._mood == "errory": # 1 + if self._mood == "empty": # 0 + assert not self._buddy + self.factory.recordUsage(self._started, "empty", 0, + total_time, None) + elif self._mood == "errory": # 1 assert not self._buddy self.factory.recordUsage(self._started, "errory", 0, total_time, None) @@ -213,7 +218,7 @@ class TransitConnection(protocol.Protocol): total_time, None) else: # 5 or 7 # we were connected, we hung up first. We record the event. - assert self._mood == "happy", self._mood # TODO: mood==None + assert self._mood == "happy", self._mood assert self._buddy starts = [self._started, self._buddy._started] total_time = finished - min(starts)