Plugged some memory leakage

7 files changed, 85 insertions, 3 deletions

diff --git a/libvncserver/corre.c b/libvncserver/corre.c
index 75f1211..bb07c77 100755
--- a/libvncserver/corre.c
+++ b/libvncserver/corre.c
@@ -50,6 +50,17 @@ static uint32_t getBgColour(char *data, int size, int bpp);
 static rfbBool rfbSendSmallRectEncodingCoRRE(rfbClientPtr cl, int x, int y,
                                           int w, int h);
 
+void rfbCoRRECleanup(rfbScreenInfoPtr screen)
+{
+  if (rreBeforeBufSize) {
+    free(rreBeforeBuf);
+    rreBeforeBufSize=0;
+  }
+  if (rreAfterBufSize) {
+    free(rreAfterBuf);
+    rreAfterBufSize=0;
+  }
+}
 
 /*
  * rfbSendRectEncodingCoRRE - send an arbitrary size rectangle using CoRRE
diff --git a/libvncserver/main.c b/libvncserver/main.c
index ee5af4a..3af90ca 100644
--- a/libvncserver/main.c
+++ b/libvncserver/main.c
@@ -982,7 +982,12 @@ void rfbScreenCleanup(rfbScreenInfoPtr screen)
   TINI_MUTEX(screen->cursorMutex);
   if(screen->cursor && screen->cursor->cleanup)
     rfbFreeCursor(screen->cursor);
+
+  rfbRRECleanup(screen);
+  rfbCoRRECleanup(screen);
+  rfbUltraCleanup(screen);
 #ifdef LIBVNCSERVER_HAVE_LIBZ
+  rfbZlibCleanup(screen);
 #ifdef LIBVNCSERVER_HAVE_LIBJPEG
   rfbTightCleanup(screen);
 #endif
diff --git a/libvncserver/private.h b/libvncserver/private.h
index 8865e81..71370b5 100644
--- a/libvncserver/private.h
+++ b/libvncserver/private.h
@@ -17,11 +17,28 @@ rfbClientPtr rfbClientIteratorHead(rfbClientIteratorPtr i);
 #ifdef LIBVNCSERVER_HAVE_LIBJPEG
 extern void rfbTightCleanup(rfbScreenInfoPtr screen);
 #endif
-#endif
+
+/* from zlib.c */
+extern void rfbZlibCleanup(rfbScreenInfoPtr screen);
 
 /* from zrle.c */
+void rfbFreeZrleData(rfbClientPtr cl);
+
+#endif
+
+
+/* from ultra.c */
+
+extern void rfbUltraCleanup(rfbScreenInfoPtr screen);
+extern void rfbFreeUltraData(rfbClientPtr cl);
+
+/* from rre.c */
+
+extern void rfbRRECleanup(rfbScreenInfoPtr screen);
+
+/* from corre.c */
 
-extern void rfbFreeZrleData(rfbClientPtr cl);
+extern void rfbCoRRECleanup(rfbScreenInfoPtr screen);
 
 #endif
 
diff --git a/libvncserver/rfbserver.c b/libvncserver/rfbserver.c
index c86d7ca..57ba5d9 100644
--- a/libvncserver/rfbserver.c
+++ b/libvncserver/rfbserver.c
@@ -484,6 +484,8 @@ rfbClientConnectionGone(rfbClientPtr cl)
     rfbFreeZrleData(cl);
 #endif
 
+    rfbFreeUltraData(cl);
+
 #ifdef LIBVNCSERVER_HAVE_LIBPTHREAD
     if(cl->screen->backgroundLoop != FALSE) {
       int i;
diff --git a/libvncserver/rre.c b/libvncserver/rre.c
index c0759aa..b43561a 100755
--- a/libvncserver/rre.c
+++ b/libvncserver/rre.c
@@ -48,6 +48,19 @@ static int subrectEncode32(uint32_t *data, int w, int h);
 static uint32_t getBgColour(char *data, int size, int bpp);
 
 
+void rfbRRECleanup(rfbScreenInfoPtr screen)
+{
+  if (rreBeforeBufSize) {
+    free(rreBeforeBuf);
+    rreBeforeBufSize=0;
+  }
+  if (rreAfterBufSize) {
+    free(rreAfterBuf);
+    rreAfterBufSize=0;
+  }
+}
+
+
 /*
  * rfbSendRectEncodingRRE - send a given rectangle using RRE encoding.
  */
diff --git a/libvncserver/ultra.c b/libvncserver/ultra.c
index 1e51446..a802026 100644
--- a/libvncserver/ultra.c
+++ b/libvncserver/ultra.c
@@ -23,7 +23,7 @@ static char *lzoBeforeBuf = NULL;
 
 static int lzoAfterBufSize = 0;
 static char *lzoAfterBuf = NULL;
-static int lzoAfterBufLen;
+static int lzoAfterBufLen = 0;
 
 /*
  * rfbSendOneRectEncodingZlib - send a given rectangle using one Zlib
@@ -32,6 +32,26 @@ static int lzoAfterBufLen;
 
 #define MAX_WRKMEM ((LZO1X_1_MEM_COMPRESS) + (sizeof(lzo_align_t) - 1)) / sizeof(lzo_align_t)
 
+void rfbUltraCleanup(rfbScreenInfoPtr screen)
+{
+  if (lzoBeforeBufSize) {
+    free(lzoBeforeBuf);
+    lzoBeforeBufSize=0;
+  }
+  if (lzoAfterBufSize) {
+    free(lzoAfterBuf);
+    lzoAfterBufSize=0;
+  }
+}
+
+void rfbFreeUltraData(rfbClientPtr cl) {
+  if (cl->compStreamInitedLZO) {
+    free(cl->lzoWrkMem);
+    cl->compStreamInitedLZO=FALSE;
+  }
+}
+
+
 static rfbBool
 rfbSendOneRectEncodingUltra(rfbClientPtr cl,
                            int x,
@@ -125,6 +145,7 @@ rfbSendOneRectEncodingUltra(rfbClientPtr cl,
     memcpy(&cl->updateBuf[cl->ublen], (char *)&hdr, sz_rfbZlibHeader);
     cl->ublen += sz_rfbZlibHeader;
 
+    /* We might want to try sending the data directly... */
     for (i = 0; i < lzoAfterBufLen;) {
 
 	int bytesToCopy = UPDATE_BUF_SIZE - cl->ublen;
diff --git a/libvncserver/zlib.c b/libvncserver/zlib.c
index 56bcc6e..7b20f74 100644
--- a/libvncserver/zlib.c
+++ b/libvncserver/zlib.c
@@ -47,6 +47,19 @@ static int zlibAfterBufSize = 0;
 static char *zlibAfterBuf = NULL;
 static int zlibAfterBufLen;
 
+void rfbZlibCleanup(rfbScreenInfoPtr screen)
+{
+  if (zlibBeforeBufSize) {
+    free(zlibBeforeBuf);
+    zlibBeforeBufSize=0;
+  }
+  if (zlibAfterBufSize) {
+    zlibAfterBufSize=0;
+    free(zlibAfterBuf);
+  }
+}
+
+
 /*
  * rfbSendOneRectEncodingZlib - send a given rectangle using one Zlib
  *                              rectangle encoding.