diff --git a/src/wormhole_mailbox_server/server.py b/src/wormhole_mailbox_server/server.py index 6f8bf99..3c96a82 100644 --- a/src/wormhole_mailbox_server/server.py +++ b/src/wormhole_mailbox_server/server.py @@ -11,6 +11,8 @@ class CrowdedError(Exception): pass class ReclaimedError(Exception): pass +class UnknownNameplateError(Exception): + pass Usage = namedtuple("Usage", ["started", "waiting_time", "total_time", "result"]) TransitUsage = namedtuple("TransitUsage", @@ -236,7 +238,7 @@ def allocate_nameplate(self, side, when): mailbox_id = self.claim_nameplate(nameplate_id, side, when) return nameplate_id, mailbox_id - def claim_nameplate(self, name, side, when): + def claim_nameplate(self, name, side, when, allow_allocate=True): # when we're done: # * there will be one row for the nameplate # * there will be one 'side' attached to it, with claimed=True @@ -249,6 +251,8 @@ def claim_nameplate(self, name, side, when): " WHERE `app_id`=? AND `name`=?", (self._app_id, name)).fetchone() if not row: + if not allow_allocate: + raise UnknownNameplateError("Claimed nameplate does not exist, and `allow_allocate` is set to False") if self._log_requests: log.msg("creating nameplate#%s for app_id %s" % (name, self._app_id)) diff --git a/src/wormhole_mailbox_server/server_websocket.py b/src/wormhole_mailbox_server/server_websocket.py index 45ed440..b2bfa67 100644 --- a/src/wormhole_mailbox_server/server_websocket.py +++ b/src/wormhole_mailbox_server/server_websocket.py @@ -3,7 +3,7 @@ from twisted.internet import reactor from twisted.python import log from autobahn.twisted import websocket -from .server import CrowdedError, ReclaimedError, SidedMessage +from .server import CrowdedError, ReclaimedError, UnknownNameplateError, SidedMessage from .util import dict_to_bytes, bytes_to_dict # The WebSocket allows the client to send "commands" to the server, and the @@ -193,6 +193,10 @@ def handle_allocate(self, server_rx): def handle_claim(self, msg, server_rx): if "nameplate" not in msg: raise Error("claim requires 'nameplate'") + allow_allocate = True + if "allocate" in msg: + allow_allocate = msg["allocate"] + if self._did_claim: raise Error("only one claim per connection") self._did_claim = True @@ -201,11 +205,13 @@ def handle_claim(self, msg, server_rx): self._nameplate_id = nameplate_id try: mailbox_id = self._app.claim_nameplate(nameplate_id, self._side, - server_rx) + server_rx, allow_allocate) except CrowdedError: raise Error("crowded") except ReclaimedError: raise Error("reclaimed") + except UnknownNameplateError: + raise Error("unallocated") self.send("claimed", mailbox=mailbox_id) def handle_release(self, msg, server_rx): diff --git a/src/wormhole_mailbox_server/test/test_web.py b/src/wormhole_mailbox_server/test/test_web.py index ecf462b..f2c45c9 100644 --- a/src/wormhole_mailbox_server/test/test_web.py +++ b/src/wormhole_mailbox_server/test/test_web.py @@ -280,6 +280,11 @@ def test_claim(self): self.assertEqual(err["type"], "error") self.assertEqual(err["error"], "claim requires 'nameplate'") + c1.send("claim", nameplate="np1", allocate=False) # not allocated + err = yield c1.next_non_ack() + self.assertEqual(err["type"], "error") + self.assertEqual(err["error"], "unallocated") + c1.send("claim", nameplate="np1") m = yield c1.next_non_ack() self.assertEqual(m["type"], "claimed")