From fcc07dad38b37d38081128f1d38ec50fca67de6f Mon Sep 17 00:00:00 2001
From: Night12138 <www2277511@163.com>
Date: Wed, 26 Oct 2022 22:53:07 +0800
Subject: [PATCH] fix python crash issue

---
 gidra/proxy/kcp_socket.py |  1 +
 python-kcp/lkcp/core.pyx  | 11 ++++++-----
 python-kcp/lkcp/kcp.py    | 28 +++++++++++++++++++++++-----
 3 files changed, 30 insertions(+), 10 deletions(-)

diff --git a/gidra/proxy/kcp_socket.py b/gidra/proxy/kcp_socket.py
index 6709f0ea..f5d54ae6 100644
--- a/gidra/proxy/kcp_socket.py
+++ b/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()
diff --git a/python-kcp/lkcp/core.pyx b/python-kcp/lkcp/core.pyx
index b8dbc029..f1e92bb8 100644
--- a/python-kcp/lkcp/core.pyx
+++ b/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)
diff --git a/python-kcp/lkcp/kcp.py b/python-kcp/lkcp/kcp.py
index a913a761..9e5cc5e8 100644
--- a/python-kcp/lkcp/kcp.py
+++ b/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)