Added ability to return file resources for request responses. Added option to specify request response auto-compression limits.

This commit is contained in:
Mark Qvist 2025-05-11 16:37:57 +02:00
parent 2912fb2184
commit 594f5fba1e
2 changed files with 44 additions and 23 deletions

View file

@ -377,7 +377,7 @@ class Destination:
else:
self.proof_strategy = proof_strategy
def register_request_handler(self, path, response_generator = None, allow = ALLOW_NONE, allowed_list = None):
def register_request_handler(self, path, response_generator = None, allow = ALLOW_NONE, allowed_list = None, auto_compress = True):
"""
Registers a request handler.
@ -385,17 +385,15 @@ class Destination:
:param response_generator: A function or method with the signature *response_generator(path, data, request_id, link_id, remote_identity, requested_at)* to be called. Whatever this funcion returns will be sent as a response to the requester. If the function returns ``None``, no response will be sent.
:param allow: One of ``RNS.Destination.ALLOW_NONE``, ``RNS.Destination.ALLOW_ALL`` or ``RNS.Destination.ALLOW_LIST``. If ``RNS.Destination.ALLOW_LIST`` is set, the request handler will only respond to requests for identified peers in the supplied list.
:param allowed_list: A list of *bytes-like* :ref:`RNS.Identity<api-identity>` hashes.
:param auto_compress: If ``True`` or ``False``, determines whether automatic compression of responses should be carried out. If set to an integer value, responses will only be auto-compressed if under this size in bytes. If omitted, the default compression settings will be followed.
:raises: ``ValueError`` if any of the supplied arguments are invalid.
"""
if path == None or path == "":
raise ValueError("Invalid path specified")
elif not callable(response_generator):
raise ValueError("Invalid response generator specified")
elif not allow in Destination.request_policies:
raise ValueError("Invalid request policy")
if path == None or path == "": raise ValueError("Invalid path specified")
elif not callable(response_generator): raise ValueError("Invalid response generator specified")
elif not allow in Destination.request_policies: raise ValueError("Invalid request policy")
else:
path_hash = RNS.Identity.truncated_hash(path.encode("utf-8"))
request_handler = [path, response_generator, allow, allowed_list]
request_handler = [path, response_generator, allow, allowed_list, auto_compress]
self.request_handlers[path_hash] = request_handler
def deregister_request_handler(self, path):
@ -491,7 +489,6 @@ class Destination:
self.latest_ratchet_time = 0
self._reload_ratchets(ratchets_path)
# TODO: Remove at some point
RNS.log("Ratchets enabled on "+str(self), RNS.LOG_DEBUG)
return True

View file

@ -40,6 +40,7 @@ import struct
import math
import time
import RNS
import io
class LinkCallbacks:
def __init__(self):
@ -194,7 +195,7 @@ class Link:
link.mode = Link.mode_from_lr_packet(packet)
# TODO: Remove
# TODO: Remove debug
RNS.log(f"Incoming link request with mode {Link.MODE_DESCRIPTIONS[link.mode]}", RNS.LOG_DEBUG)
link.update_mdu()
@ -851,6 +852,7 @@ class Link:
response_generator = request_handler[1]
allow = request_handler[2]
allowed_list = request_handler[3]
auto_compress = request_handler[4]
allowed = False
if not allow == RNS.Destination.ALLOW_NONE:
@ -869,18 +871,29 @@ class Link:
else:
raise TypeError("Invalid signature for response generator callback")
if response != None:
packed_response = umsgpack.packb([request_id, response])
file_response = False
file_handle = None
if type(response) == list or type(response) == tuple:
metadata = None
if len(response) > 0 and type(response[0]) == io.BufferedReader:
if len(response) > 1: metadata = response[1]
file_handle = response[0]
file_response = True
if response != None:
if file_response:
response_resource = RNS.Resource(file_handle, self, metadata=metadata, request_id = request_id, is_response = True, auto_compress=auto_compress)
else:
packed_response = umsgpack.packb([request_id, response])
if len(packed_response) <= self.mdu:
RNS.Packet(self, packed_response, RNS.Packet.DATA, context = RNS.Packet.RESPONSE).send()
else:
response_resource = RNS.Resource(packed_response, self, request_id = request_id, is_response = True)
response_resource = RNS.Resource(packed_response, self, request_id = request_id, is_response = True, auto_compress=auto_compress)
else:
identity_string = str(self.get_remote_identity()) if self.get_remote_identity() != None else "<Unknown>"
RNS.log("Request "+RNS.prettyhexrep(request_id)+" from "+identity_string+" not allowed for: "+str(path), RNS.LOG_DEBUG)
def handle_response(self, request_id, response_data, response_size, response_transfer_size):
def handle_response(self, request_id, response_data, response_size, response_transfer_size, metadata=None):
if self.status == Link.ACTIVE:
remove = None
for pending_request in self.pending_requests:
@ -891,7 +904,7 @@ class Link:
if pending_request.response_transfer_size == None:
pending_request.response_transfer_size = 0
pending_request.response_transfer_size += response_transfer_size
pending_request.response_received(response_data)
pending_request.response_received(response_data, metadata)
except Exception as e:
RNS.log("Error occurred while handling response. The contained exception was: "+str(e), RNS.LOG_ERROR)
@ -914,12 +927,21 @@ class Link:
def response_resource_concluded(self, resource):
if resource.status == RNS.Resource.COMPLETE:
# If the response resource has metadata, this
# is a file response, and we'll pass the open
# file handle directly.
if resource.has_metadata:
self.handle_response(resource.request_id, resource.data, resource.total_size, resource.size, metadata=resource.metadata)
# If not, we'll unpack the response data and
# pass the unpacked structure to the handler
else:
packed_response = resource.data.read()
unpacked_response = umsgpack.unpackb(packed_response)
request_id = unpacked_response[0]
response_data = unpacked_response[1]
self.handle_response(request_id, response_data, resource.total_size, resource.size)
else:
RNS.log("Incoming response resource failed with status: "+RNS.hexrep([resource.status]), RNS.LOG_DEBUG)
for pending_request in self.pending_requests:
@ -1343,6 +1365,7 @@ class RequestReceipt():
self.response = None
self.response_transfer_size = None
self.response_size = None
self.metadata = None
self.status = RequestReceipt.SENT
self.sent_at = time.time()
self.progress = 0
@ -1429,10 +1452,11 @@ class RequestReceipt():
resource.cancel()
def response_received(self, response):
def response_received(self, response, metadata=None):
if not self.status == RequestReceipt.FAILED:
self.progress = 1.0
self.response = response
self.metadata = metadata
self.status = RequestReceipt.READY
self.response_concluded_at = time.time()