diff --git a/tests/test_uwsgi_structs.py b/tests/test_uwsgi_structs.py new file mode 100644 index 0000000..8fa6618 --- /dev/null +++ b/tests/test_uwsgi_structs.py @@ -0,0 +1,55 @@ +import unittest + +from uwsgi_tools import uwsgi_structs as uw + +""" +From definitions available at +http://uwsgi-docs.readthedocs.io/en/latest/Protocol.html +""" + + +class TestUwsgiPacketHeader(unittest.TestCase): + header = uw.UwsgiPacketHeader(0, 166, 0) + header_bytes = b'\x00\xa6\x00\x00' + + def test_packet_header_serializes(self): + expected = self.header_bytes + + self.assertEqual(expected, bytearray(self.header)) + + def test_packet_header_deserializes(self): + buffer = bytearray(self.header_bytes) + expected_header = uw.UwsgiPacketHeader(0, 166, 0) + + header = uw.UwsgiPacketHeader.from_buffer(buffer) + self.assertEqual( + bytearray(expected_header), + bytearray(header) + ) + + +class TestUwsgiVar(unittest.TestCase): + + key = b'foo' + key_size = len(key) + val = b'bar' + val_size = len(val) + + var = uw.UwsgiVar(key_size, key, val_size, val) + var_bytes = b'\x03\x00foo\x03\x00bar' + + def test_var_serializes(self): + self.assertEqual( + self.var_bytes, + bytearray(self.var) + ) + + def test_var_deserializes(self): + buffer = bytearray(self.var_bytes) + expected = self.var + + var = uw.UwsgiVar.from_buffer(buffer) + self.assertEqual( + bytearray(expected), + bytearray(var) + ) diff --git a/uwsgi_tools/compat.py b/uwsgi_tools/compat.py index 521b979..aa48279 100644 --- a/uwsgi_tools/compat.py +++ b/uwsgi_tools/compat.py @@ -2,7 +2,7 @@ __all__ = [ 'BaseHTTPRequestHandler', 'TCPServer', 'get_content_type', 'urlsplit', - 'hex2bytes', + 'struct2bytes', ] PY3 = sys.version_info[0] == 3 @@ -25,8 +25,8 @@ def get_content_type(headers): return headers.typeheader -def hex2bytes(s): +def struct2bytes(s): if PY3: - return bytes.fromhex(s) + return bytes(s) else: - return s.decode('hex') + return bytes(bytearray(s)) diff --git a/uwsgi_tools/utils.py b/uwsgi_tools/utils.py index 1b153c4..9b324a9 100644 --- a/uwsgi_tools/utils.py +++ b/uwsgi_tools/utils.py @@ -1,17 +1,18 @@ -from .compat import hex2bytes, urlsplit - - -def sz(x): - s = hex(x if isinstance(x, int) else len(x))[2:].rjust(4, '0') - s = hex2bytes(s) - return s[::-1] +from .compat import struct2bytes, urlsplit +from .uwsgi_structs import UwsgiPacketHeader, UwsgiVar def pack_uwsgi_vars(var): - pk = b'' - for k, v in var.items() if hasattr(var, 'items') else var: - pk += sz(k) + k.encode('utf8') + sz(v) + v.encode('utf8') - return b'\x00' + sz(pk) + b'\x00' + pk + encoded_vars = [ + (k.encode('utf-8'), v.encode('utf-8')) + for k, v in var.items() + ] + packed_vars = b''.join( + struct2bytes(UwsgiVar(len(k), k, len(v), v)) + for k, v in encoded_vars + ) + packet_header = struct2bytes(UwsgiPacketHeader(0, len(packed_vars), 0)) + return packet_header + packed_vars def parse_addr(addr, default_port=3030): diff --git a/uwsgi_tools/uwsgi_structs.py b/uwsgi_tools/uwsgi_structs.py new file mode 100644 index 0000000..73d0982 --- /dev/null +++ b/uwsgi_tools/uwsgi_structs.py @@ -0,0 +1,57 @@ +""" +From definitions available at +http://uwsgi-docs.readthedocs.io/en/latest/Protocol.html +""" + +import ctypes + + +class UwsgiPacketHeader(ctypes.Structure): + """ + struct uwsgi_packet_header { + uint8_t modifier1; + uint16_t datasize; + uint8_t modifier2; + }; + """ + _pack_ = 1 + _fields_ = [ + ("modifier1", ctypes.c_int8), + ("datasize", ctypes.c_int16), + ("modifier2", ctypes.c_int8), + ] + + +class UwsgiVar(object): + """ + struct uwsgi_var { + uint16_t key_size; + uint8_t key[key_size]; + uint16_t val_size; + uint8_t val[val_size]; + } + """ + + def __new__(self, key_size, key, val_size, val): + class UwsgiVar(ctypes.Structure): + _pack_ = 1 + _fields_ = [ + ("key_size", ctypes.c_int16), + ("key", ctypes.c_char * key_size), + ("val_size", ctypes.c_int16), + ("val", ctypes.c_char * val_size), + ] + + return UwsgiVar(key_size, key, val_size, val) + + @classmethod + def from_buffer(cls, buffer, offset=0): + key_size = ctypes.c_int16.from_buffer(buffer, offset).value + offset += ctypes.sizeof(ctypes.c_int16) + key = (ctypes.c_char * key_size).from_buffer(buffer, offset).value + offset += ctypes.sizeof(ctypes.c_char * key_size) + val_size = ctypes.c_int16.from_buffer(buffer, offset).value + offset += ctypes.sizeof(ctypes.c_int16) + val = (ctypes.c_char * val_size).from_buffer(buffer, offset).value + + return cls(key_size, key, val_size, val)