Allow a richer set of transformations by passing a mutable

structure into in2out and out2in and allowing handler to
move or replace the pointer and/or length.

Author: sobomax

python/AsyncProxy.py

@@ -23,7 +23,7 @@
 
 
 from ctypes import cdll, c_int, c_char_p, c_ushort, c_void_p, CFUNCTYPE, \
-  POINTER, pointer, Structure, Union, byref
+  POINTER, pointer, Structure, Union, byref, c_size_t
 
 from sysconfig import get_config_var
 from site import getsitepackages
@@ -58,7 +58,13 @@ class asyncproxy_ctor_args(Structure):
         ("_anon_union", _AnonUnion),
     ]
 
-_asp_data_cb = CFUNCTYPE(None, c_void_p, c_int)
+class transform_res(Structure):
+    _fields_ = [
+        ("buf", c_void_p),
+        ("len", c_size_t),
+    ]
+
+_asp_data_cb = CFUNCTYPE(None, POINTER(transform_res))
 
 _esuf = get_config_var('EXT_SUFFIX')
 if not _esuf:

src/asyncproxy.c

@@ -74,7 +74,7 @@ struct asyncproxy {
         socklen_t alen;
     } destaddr;
     int last_seen_alive;
-    void (*transform[2])(void *, int);
+    void (*transform[2])(struct transform_res *);
     int needsjoin;
     char addrbuf[FILENAME_MAX];
 };
@@ -163,7 +163,7 @@ asp_sock_setnonblock(int fd)
 
 struct io_buf {
     unsigned char data[16 * 1024];
-    int len;
+    size_t len;
 };
 
 #define BUF_FREE(ibp) (sizeof((ibp)->data) - (ibp)->len)
@@ -278,17 +278,32 @@ asyncproxy_run(void *args)
                     goto out;
                 }
                 pthread_mutex_lock(&ap->mutex);
-                void (*transform)(void *, int) = ap->transform[i];
+                __typeof(ap->transform[i]) transform = ap->transform[i];
                 pthread_mutex_unlock(&ap->mutex);
                 if (transform != NULL) {
 #if defined(PYTHON_AWARE)
                     PyGILState_STATE gstate;
                     gstate = PyGILState_Ensure();
 #endif
-                    transform(BUF_P(&bufs[i]), r.len);
+                    struct transform_res tr = {BUF_P(&bufs[i]), r.len};
+                    transform(&tr);
 #if defined(PYTHON_AWARE)
                     PyGILState_Release(gstate);
 #endif
+                    if ((ssize_t)tr.len != r.len) {
+                        assert(BUF_FREE(&bufs[i]) >= tr.len);
+                        r.len = tr.len;
+                    }
+                    if (tr.buf != BUF_P(&bufs[i])) {
+                        if (tr.len > 0) {
+                            assert(BUF_FREE(&bufs[i]) >= tr.len);
+                            memmove(BUF_P(&bufs[i]), tr.buf, tr.len);
+                        }
+                        r.len = tr.len;
+                    } else if ((ssize_t)tr.len != r.len) {
+                        assert((ssize_t)tr.len < r.len);
+                        r.len = tr.len;
+                    }
                 }
                 bufs[i].len += r.len;
                 if (BUF_FREE(&bufs[i]) == 0) {
@@ -308,12 +323,12 @@ asyncproxy_run(void *args)
                     fprintf(stderr, "asyncproxy_run(%p): sent %ld bytes to %d\n", ap, rlen, pfds[j].fd);
                     fflush(stderr);
                 }
-                if (rlen < bufs[i].len) {
+                if (rlen < (ssize_t)bufs[i].len) {
                     pfds[j].events |= POLLOUT;
                 }
                 if (rlen <= 0)
                     continue;
-                if (rlen < bufs[i].len) {
+                if (rlen < (ssize_t)bufs[i].len) {
                     memmove(bufs[i].data, bufs[i].data + rlen, bufs[i].len - rlen);
                     bufs[i].len -= rlen;
                 } else {
@@ -552,7 +567,7 @@ asyncproxy_isalive(void *_ap)
 }
 
 void
-asyncproxy_set_i2o(void *_ap, void (*i2ofp)(void *, int))
+asyncproxy_set_i2o(void *_ap, void (*i2ofp)(struct transform_res *))
 {
     struct asyncproxy *ap;
 
@@ -564,7 +579,7 @@ asyncproxy_set_i2o(void *_ap, void (*i2ofp)(void *, int))
 }
 
 void
-asyncproxy_set_o2i(void *_ap, void (*o2ifp)(void *, int))
+asyncproxy_set_o2i(void *_ap, void (*o2ifp)(struct transform_res *))
 {
     struct asyncproxy *ap;
 

src/asyncproxy.h

@@ -24,11 +24,16 @@ struct asyncproxy_ctor_args {
     };
 };
 
+struct transform_res {
+    void *buf;
+    size_t len;
+};
+
 void * asyncproxy_ctor(const struct asyncproxy_ctor_args *);
 int asyncproxy_start(void *);
 int asyncproxy_isalive(void *);
-void asyncproxy_set_i2o(void *, void (*)(void *, int));
-void asyncproxy_set_o2i(void *, void (*)(void *, int));
+void asyncproxy_set_i2o(void *, void (*)(struct transform_res *));
+void asyncproxy_set_o2i(void *, void (*)(struct transform_res *));
 void asyncproxy_join(void *, int);
 void asyncproxy_dtor(void *);
 const char * asyncproxy_describe(void *);

tests/AsyncProxy2FD_test.py

@@ -1,24 +1,33 @@
 import socket
 import unittest
 from ctypes import string_at, memmove
-from libasyncproxy.AsyncProxy import AsyncProxy2FD
+from libasyncproxy.AsyncProxy import AsyncProxy2FD, transform_res
 
 class NosyProxy(AsyncProxy2FD):
-    def in2out(self, ptr, length):
-        # Read original data from the pointer
+    def in2out(self, res_p):
+        # unpack the struct
+        tr = res_p.contents
+        ptr, length = tr.buf, tr.len
+        # read, transform, write back
         original = string_at(ptr, length)
-        # Transform data to uppercase (note: transformation must retain length)
-        transformed = original.upper()
+        # Transform data to uppercase
+        length -= 1
+        transformed = original.upper()[:length]
         memmove(ptr, transformed, length)
         print("in2out hook: transformed", original, "to", transformed)
+        tr.len = length
 
-    def out2in(self, ptr, length):
-        # Read original data from the pointer
-        original = string_at(ptr, length)
-        # Reverse the data bytes (again, ensuring the length remains unchanged)
-        transformed = original[::-1]
+    def out2in(self, res_p):
+        # unpack the struct
+        tr = res_p.contents
+        ptr, length = tr.buf, tr.len
+        # read, transform, write back
+        original    = string_at(ptr, length)
+        length -= 1
+        transformed = original[::-1][1:]
         memmove(ptr, transformed, length)
         print("out2in hook: transformed", original, "to", transformed)
+        tr.len = length
 
 class AsyncProxy2FDTest(unittest.TestCase):
     def test_AsyncProxy(self):
@@ -45,7 +54,7 @@ def test_AsyncProxy(self):
 
         server_received = server_socket.recv(1024)
         print("Server received:", server_received.decode())
-        self.assertEqual(client_message.upper(), server_received)
+        self.assertEqual(client_message.upper()[:-1], server_received)
 
         # Now send a message from server back to client.
         server_message = b"Hello from Server!"
@@ -54,7 +63,7 @@ def test_AsyncProxy(self):
 
         client_received = client_socket.recv(1024)
         print("Client received:", client_received.decode())
-        self.assertEqual(server_message[::-1], client_received)
+        self.assertEqual(server_message[::-1][1:], client_received)
 
         # Shutdown the proxy worker and cleanup.
         proxy_fd.join(shutdown=True)