Merge pull request #8 from tomdottom/easier-to-read-pack-uwsgi-vars

Easier to read pack uwsgi vars

tests/test_uwsgi_structs.py

@@ -0,0 +1,48 @@
+import unittest
+
+from uwsgi_tools import uwsgi_structs as uw
+
+
+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)
+        )

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))

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):

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)