diff --git a/RNS/Destination.py b/RNS/Destination.py index 87296855..7ce6239f 100755 --- a/RNS/Destination.py +++ b/RNS/Destination.py @@ -155,6 +155,7 @@ class Destination: self.accept_link_requests = True self.callbacks = Callbacks() self.request_handlers = {} + self.max_request_size = None self.type = type self.direction = direction self.proof_strategy = Destination.PROVE_NONE @@ -362,10 +363,20 @@ class Destination: :param proof_strategy: One of ``RNS.Destination.PROVE_NONE``, ``RNS.Destination.PROVE_ALL`` or ``RNS.Destination.PROVE_APP``. If ``RNS.Destination.PROVE_APP`` is set, the `proof_requested_callback` will be called to determine whether a proof should be sent or not. """ - if not proof_strategy in Destination.proof_strategies: - raise TypeError("Unsupported proof strategy") - else: - self.proof_strategy = proof_strategy + if not proof_strategy in Destination.proof_strategies: raise TypeError("Unsupported proof strategy") + else: self.proof_strategy = proof_strategy + + def set_max_request_size(self, max_request_size): + """ + Registers a request handler. + + :param max_request_size: The path for the request handler to be registered. + :raises: ``TypeError`` or ``ValueError`` if any of the argument is invalid. + """ + try: max_request_size = int(max_request_size) + except: raise TypeError("Invalid maximum request size specified") + if max_request_size < 0: raise ValueError("Maximum request size cannot be negative") + self.max_request_size = max_request_size def register_request_handler(self, path, response_generator = None, allow = ALLOW_NONE, allowed_list = None, auto_compress = True): """ diff --git a/RNS/Link.py b/RNS/Link.py index 09d3c44e..be4d8202 100644 --- a/RNS/Link.py +++ b/RNS/Link.py @@ -981,9 +981,13 @@ class Link: request_id = packet.getTruncatedHash() packed_request = self.decrypt(packet.data) if packed_request != None: - unpacked_request = umsgpack.unpackb(packed_request) - def job(): self.handle_request(request_id, unpacked_request) - threading.Thread(target=job, daemon=True).start() + if self.destination.max_request_size == None: size_ok = True + else: size_ok = len(packed_request) <= self.destination.max_request_size + if not size_ok: RNS.log(f"Ignored request with excessive size {RNS.prettysize(len(packed_request))} on {self.destination}", RNS.LOG_DEBUG) if RNS.sl(RNS.LOG_DEBUG) else None + else: + unpacked_request = umsgpack.unpackb(packed_request) + def job(): self.handle_request(request_id, unpacked_request) + threading.Thread(target=job, daemon=True).start() self.__update_phy_stats(packet, query_shared=True) except Exception as e: RNS.log("Error occurred while handling request. The contained exception was: "+str(e), RNS.LOG_ERROR) @@ -1016,7 +1020,12 @@ class Link: try: if RNS.ResourceAdvertisement.is_request(packet): if self.destination.request_handlers: - RNS.Resource.accept(packet, callback=self.request_resource_concluded) + if self.destination.max_request_size == None: size_ok = True + else: size_ok = RNS.ResourceAdvertisement.read_size(packet) <= self.destination.max_request_size + if size_ok: RNS.Resource.accept(packet, callback=self.request_resource_concluded) + else: + RNS.Resource.reject(packet) + RNS.log(f"Ignored request with excessive size {RNS.prettysize(RNS.ResourceAdvertisement.read_size(packet))} on {self.destination}", RNS.LOG_DEBUG) if RNS.sl(RNS.LOG_DEBUG) else None elif RNS.ResourceAdvertisement.is_response(packet): request_id = RNS.ResourceAdvertisement.read_request_id(packet) for pending_request in self.pending_requests: