Better wrapping for libz and liblzma

1 files changed, 49 insertions, 10 deletions

diff --git a/src/wrapped/wrappedlibz.c b/src/wrapped/wrappedlibz.c
index a8184e11..5934b27a 100755
--- a/src/wrapped/wrappedlibz.c
+++ b/src/wrapped/wrappedlibz.c
@@ -51,6 +51,16 @@ static void* find_alloc_Fct(void* fct)
     printf_log(LOG_NONE, "Warning, no more slot for zlib alloc callback\n");
     return NULL;
 }
+static void* reverse_alloc_Fct(void* fct)
+{
+    if(!fct) return fct;
+    if(CheckBridged(my_lib->w.bridge, fct))
+        return (void*)CheckBridged(my_lib->w.bridge, fct);
+    #define GO(A) if(my_alloc_##A == fct) return (void*)my_alloc_fct_##A;
+    SUPER()
+    #undef GO
+    return (void*)AddBridge(my_lib->w.bridge, pFpuu, fct, 0, NULL);
+}
 // free ...
 #define GO(A)   \
 static uintptr_t my_free_fct_##A = 0;                               \
@@ -73,6 +83,16 @@ static void* find_free_Fct(void* fct)
     printf_log(LOG_NONE, "Warning, no more slot for zlib free callback\n");
     return NULL;
 }
+static void* reverse_free_Fct(void* fct)
+{
+    if(!fct) return fct;
+    if(CheckBridged(my_lib->w.bridge, fct))
+        return (void*)CheckBridged(my_lib->w.bridge, fct);
+    #define GO(A) if(my_free_##A == fct) return (void*)my_free_fct_##A;
+    SUPER()
+    #undef GO
+    return (void*)AddBridge(my_lib->w.bridge, vFpp, fct, 0, NULL);
+}
 #undef SUPER
 
 typedef struct z_stream_s {
@@ -99,28 +119,41 @@ static void wrapper_stream_z(x64emu_t* emu, void* str)
     stream->zalloc = find_alloc_Fct(stream->zalloc);
     stream->zfree = find_free_Fct(stream->zfree);
 }
+static void unwrapper_stream_z(x64emu_t* emu, void* str)
+{
+    (void)emu;
+    z_stream *stream = (z_stream*)str;
+    stream->zalloc = reverse_alloc_Fct(stream->zalloc);
+    stream->zfree = reverse_free_Fct(stream->zfree);
+}
 
 EXPORT int my_inflateInit_(x64emu_t* emu, void* str, void* version, int size)
 {
     wrapper_stream_z(emu, str);
-    return my->inflateInit_(str, version, size);
+    int ret = my->inflateInit_(str, version, size);
+    unwrapper_stream_z(emu, str);
+    return ret;
 }
 
 EXPORT int my_inflateInit2_(x64emu_t* emu, void* str, int windowBits, void* version, int stream_size)
 {
     wrapper_stream_z(emu, str);
-    return my->inflateInit2_(str, windowBits, version, stream_size);
+    int ret = my->inflateInit2_(str, windowBits, version, stream_size);
+    unwrapper_stream_z(emu, str);
+    return ret;
 }
 
-EXPORT int my_inflateBackInit_(x64emu_t* emu, void* strm, int windowBits, void *window, void* version, int size)
+EXPORT int my_inflateBackInit_(x64emu_t* emu, void* str, int windowBits, void *window, void* version, int size)
 {
-    wrapper_stream_z(emu, strm);
-    return my->inflateBackInit_(strm, windowBits, window, version, size);
+    wrapper_stream_z(emu, str);
+    int ret = my->inflateBackInit_(str, windowBits, window, version, size);
+    unwrapper_stream_z(emu, str);
+    return ret;
 }
 
-// TODO: remove this?
 EXPORT int my_inflateEnd(x64emu_t* emu, void* str)
 {
+    wrapper_stream_z(emu, str);
     int r = my->inflateEnd(str);
     return r;
 }
@@ -128,18 +161,23 @@ EXPORT int my_inflateEnd(x64emu_t* emu, void* str)
 EXPORT int my_deflateInit_(x64emu_t* emu, void* str, int level, void* version, int stream_size)
 {
     wrapper_stream_z(emu, str);
-    return my->deflateInit_(str, level, version, stream_size);
+    int ret = my->deflateInit_(str, level, version, stream_size);
+    unwrapper_stream_z(emu, str);
+    return ret;
 }
 
 EXPORT int my_deflateInit2_(x64emu_t* emu, void* str, int level, int method, int windowBits, int memLevel, int strategy, void* version, int stream_size)
 {
     wrapper_stream_z(emu, str);
-    return my->deflateInit2_(str, level, method, windowBits, memLevel, strategy, version, stream_size);
+    int ret = my->deflateInit2_(str, level, method, windowBits, memLevel, strategy, version, stream_size);
+    unwrapper_stream_z(emu, str);
+    return ret;
 }
 
 // TODO: remove this?
 EXPORT int my_deflateEnd(x64emu_t* emu, void* str)
 {
+    wrapper_stream_z(emu, str);
     int r = my->deflateEnd(str);
     return r;
 }
@@ -147,8 +185,9 @@ EXPORT int my_deflateEnd(x64emu_t* emu, void* str)
 EXPORT int my_inflate(x64emu_t* emu, void* str, int flush)
 {
     wrapper_stream_z(emu, str);
-    return my->inflate(str, flush);
-    //TODO: should unwrap the stream
+    int ret = my->inflate(str, flush);
+    unwrapper_stream_z(emu, str);
+    return ret;
 }