fix python crash issue (#2)

gidra/proxy/kcp_socket.py

@@ -95,6 +95,7 @@ def bind(self, addr: _Address) -> bool:
         )
         self.kcp.setmtu(1200)
         self.kcp.wndsize(1024, 1024)
+        self.kcp.nodelay(1, 10, 2, 1)
 
         threading.Thread(target=self._kcp_update).start()
         threading.Thread(target=self._kcp_recv).start()

python-kcp/lkcp/core.pyx

@@ -59,13 +59,12 @@ cdef extern from "compat.h":
 
 cdef struct UsrInfo:
     int handle
+    char* recv_buffer
 
 g_KcpAgentCbs = {}
 
 RECV_BUFFER_LEN = 4*1024*1024
 
-cdef char* recv_buffer = <char *>PyMem_Malloc(sizeof(char)*RECV_BUFFER_LEN)
-
 cdef int kcp_output_callback(const char *buf, int len, ikcpcb *kcp, void *arg):
     global g_KcpAgentCbs
     cdef UsrInfo *c = <UsrInfo *>arg;
@@ -80,6 +79,7 @@ cdef void del_kcp(PyObject *obj):
     if ckcp.user != NULL:
         global g_KcpAgentCbs
         c = <UsrInfo *>ckcp.user
+        PyMem_Free(c.recv_buffer)
         uid = <object>c.handle
         del g_KcpAgentCbs[uid]
         PyMem_Free(c)
@@ -92,18 +92,19 @@ def lkcp_create(conv, token, uid, cb):
     g_KcpAgentCbs[uid] = cb
     cdef UsrInfo *c = <UsrInfo *>PyMem_Malloc(sizeof(UsrInfo))
     c.handle = <int>uid
+    c.recv_buffer = <char *>PyMem_Malloc(sizeof(char)*RECV_BUFFER_LEN)
     cdef ikcpcb* ckcp = ikcp_create(conv, token, c)
     ckcp.output = kcp_output_callback
     return make_capsule(ckcp, NULL, del_kcp)
 
 def lkcp_recv(kcp):
-    global recv_buffer
     cdef ikcpcb* ckcp = <ikcpcb*>get_pointer(<object>kcp, NULL)
-    hr = ikcp_recv(ckcp, recv_buffer, RECV_BUFFER_LEN)
+    c = <UsrInfo *>ckcp.user
+    hr = ikcp_recv(ckcp, c.recv_buffer, RECV_BUFFER_LEN)
     if hr <= 0:
         return hr,None
     else:
-        return hr,recv_buffer[:hr]
+        return hr,c.recv_buffer[:hr]
 
 def lkcp_send(kcp, data):
     cdef ikcpcb* ckcp = <ikcpcb*>get_pointer(<object>kcp, NULL)

python-kcp/lkcp/kcp.py

@@ -1,5 +1,5 @@
 from . import core
-import sys
+import sys, threading
 
 __all__ = ["KcpObj"]
 
@@ -9,6 +9,7 @@ class KcpObj:
     def __init__(self, conv, token, callback):
         self.conv = conv
         self.token = token
+        self.mutex = threading.Lock()
 
         global i
         self.cobj = core.lkcp_create(conv, token, i, callback)
@@ -21,24 +22,41 @@ def nodelay(self, nodelay, interval, resend, nc):
         return core.lkcp_nodelay(self.cobj, nodelay, interval, resend, nc)
 
     def check(self, current):
-        return core.lkcp_check(self.cobj, current)
+        self.mutex.acquire()
+        ret = core.lkcp_check(self.cobj, current)
+        self.mutex.release()
+        return ret
 
     def update(self, current):
+        self.mutex.acquire()
         core.lkcp_update(self.cobj, current)
+        self.mutex.release()
 
     def send(self, data):
+        self.mutex.acquire()
         if sys.version_info.major == 3 and isinstance(data, str):
             data = data.encode("UTF-8")
-        return core.lkcp_send(self.cobj, data)
+        ret = core.lkcp_send(self.cobj, data)
+        self.mutex.release()
+        return ret
 
     def input(self, data):
-        return core.lkcp_input(self.cobj, data)
+        self.mutex.acquire()
+        ret = core.lkcp_input(self.cobj, data)
+        self.mutex.release()
+        return ret
 
     def recv(self):
-        return core.lkcp_recv(self.cobj)
+        self.mutex.acquire()
+        ret = core.lkcp_recv(self.cobj)
+        self.mutex.release()
+        return ret
 
     def flush(self):
+        self.mutex.acquire()
         core.lkcp_flush(self.cobj)
+        self.mutex.release()
+
 
     def setmtu(self, mtu):
         core.lkcp_setmtu(self.cobj, mtu)