diff --git a/quanttp/__main__.py b/quanttp/__main__.py index c0fb48e..89fe540 100644 --- a/quanttp/__main__.py +++ b/quanttp/__main__.py @@ -58,6 +58,10 @@ def api_randnormal(): def api_randbytes(): try: length = int(request.args.get('length')) + if length < 1: + return Response('length must be greater than 0', status=400, content_type='text/plain') + elif length > 8192: + return Response('length must be less than or equal to 8192', status=400, content_type='text/plain') qng_wrapper.clear() return Response(qng_wrapper.randbytes(length), content_type='application/octet-stream') except ValueError as e: @@ -85,6 +89,10 @@ def handle_ws_message(message, websocket, subscribed): websocket.send(str(qng_wrapper.randnormal())) elif split_message[0] == 'RANDBYTES': length = int(split_message[1]) + if length < 1: + raise Exception('length must be greater than 0') + elif length > 8192: + raise Exception('length must be less than or equal to 8192') qng_wrapper.clear() websocket.send(qng_wrapper.randbytes(length)) elif split_message[0] == 'SUBSCRIBEINT32': @@ -107,6 +115,10 @@ def handle_ws_message(message, websocket, subscribed): websocket.send(str(qng_wrapper.randnormal())) elif split_message[0] == 'SUBSCRIBEBYTES': chunk = int(split_message[1]) + if chunk < 1: + raise Exception('chunk must be greater than 0') + elif chunk > 8192: + raise Exception('chunk must be less than or equal to 8192') qng_wrapper.clear() if not subscribed[0]: subscribed[0] = True