Better wrapping for libz and liblzma

5 files changed, 152 insertions, 61 deletions

diff --git a/src/wrapped/generated/functions_list.txt b/src/wrapped/generated/functions_list.txt
index f7097031..bcb6f0af 100644
--- a/src/wrapped/generated/functions_list.txt
+++ b/src/wrapped/generated/functions_list.txt
@@ -4419,8 +4419,12 @@ wrappedlibz:
 - iFpiiiiipi:
   - deflateInit2_
 wrappedlzma:
+- vFp:
+  - lzma_end
 - vFpp:
   - lzma_index_end
+- iFpi:
+  - lzma_code
 - iFpU:
   - lzma_alone_decoder
 - iFpp:
diff --git a/src/wrapped/generated/wrappedlzmatypes.h b/src/wrapped/generated/wrappedlzmatypes.h
index a2d6d8cc..43862c5f 100644
--- a/src/wrapped/generated/wrappedlzmatypes.h
+++ b/src/wrapped/generated/wrappedlzmatypes.h
@@ -11,7 +11,9 @@
 #define ADDED_FUNCTIONS() 
 #endif
 
+typedef void (*vFp_t)(void*);
 typedef void (*vFpp_t)(void*, void*);
+typedef int32_t (*iFpi_t)(void*, int32_t);
 typedef int32_t (*iFpU_t)(void*, uint64_t);
 typedef int32_t (*iFpp_t)(void*, void*);
 typedef int32_t (*iFpui_t)(void*, uint32_t, int32_t);
@@ -22,7 +24,9 @@ typedef int32_t (*iFpppppL_t)(void*, void*, void*, void*, void*, uintptr_t);
 typedef int32_t (*iFpupppLppL_t)(void*, uint32_t, void*, void*, void*, uintptr_t, void*, void*, uintptr_t);
 
 #define SUPER() ADDED_FUNCTIONS() \
+	GO(lzma_end, vFp_t) \
 	GO(lzma_index_end, vFpp_t) \
+	GO(lzma_code, iFpi_t) \
 	GO(lzma_alone_decoder, iFpU_t) \
 	GO(lzma_alone_encoder, iFpp_t) \
 	GO(lzma_raw_decoder, iFpp_t) \
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;
 }
 
 
diff --git a/src/wrapped/wrappedlzma.c b/src/wrapped/wrappedlzma.c
index b7839e3c..ed578b5a 100755
--- a/src/wrapped/wrappedlzma.c
+++ b/src/wrapped/wrappedlzma.c
@@ -81,6 +81,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;                               \
@@ -103,109 +113,143 @@ 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
 
-static void wrap_alloc_struct(lzma_allocator_t* dst, lzma_allocator_t* src)
+static void wrap_alloc_struct(lzma_allocator_t* a)
 {
-    if(!src)
+    if(!a)
         return;
-    dst->opaque = src->opaque;
-    dst->alloc = find_alloc_Fct(src->alloc);
-    dst->free = find_free_Fct(src->free);
+    a->alloc = find_alloc_Fct(a->alloc);
+    a->free = find_free_Fct(a->free);
+
+}
+static void unwrap_alloc_struct(lzma_allocator_t* a)
+{
+    if(!a)
+        return;
+    a->alloc = reverse_alloc_Fct(a->alloc);
+    a->free = reverse_free_Fct(a->free);
 
 }
 
 EXPORT int my_lzma_index_buffer_decode(x64emu_t* emu, void* i, void* memlimit, lzma_allocator_t* alloc, void* in_, void* in_pos, size_t in_size)
 {
-    lzma_allocator_t allocator = {0};
-    wrap_alloc_struct(&allocator, alloc);
-    return my->lzma_index_buffer_decode(i, memlimit, alloc?&allocator:NULL, in_, in_pos, in_size);
+    wrap_alloc_struct(alloc);
+    int ret = my->lzma_index_buffer_decode(i, memlimit, alloc, in_, in_pos, in_size);
+    unwrap_alloc_struct(alloc);
+    return ret;
 }
 
 EXPORT void my_lzma_index_end(x64emu_t* emu, void* i, lzma_allocator_t* alloc)
 {
-    lzma_allocator_t allocator = {0};
-    wrap_alloc_struct(&allocator, alloc);
-    return my->lzma_index_end(i,alloc?&allocator:NULL);
+    wrap_alloc_struct(alloc);
+    my->lzma_index_end(i,alloc);
+    unwrap_alloc_struct(alloc);
 }
 
 EXPORT int my_lzma_stream_buffer_decode(x64emu_t* emu, void* memlimit, uint32_t flags, lzma_allocator_t* alloc, void* in_, void* in_pos, size_t in_size, void* out_, void* out_pos, size_t out_size)
 {
-    lzma_allocator_t allocator = {0};
-    wrap_alloc_struct(&allocator, alloc);
-    return my->lzma_stream_buffer_decode(memlimit, flags, alloc?&allocator:NULL, in_, in_pos, in_size, out_, out_pos, out_size);
+    wrap_alloc_struct(alloc);
+    int ret = my->lzma_stream_buffer_decode(memlimit, flags, alloc, in_, in_pos, in_size, out_, out_pos, out_size);
+    unwrap_alloc_struct(alloc);
+    return ret;
 }
 
 EXPORT int my_lzma_stream_decoder(x64emu_t* emu, lzma_stream_t* stream, uint64_t memlimit, uint32_t flags)
 {
-    // not restoring the allocator after, so lzma_code and lzma_end can be used without "GOM" wrapping
-    if(stream->allocator)
-        wrap_alloc_struct(stream->allocator, stream->allocator);
-    return my->lzma_stream_decoder(stream, memlimit, flags);
+    wrap_alloc_struct(stream->allocator);
+    int ret = my->lzma_stream_decoder(stream, memlimit, flags);
+    unwrap_alloc_struct(stream->allocator);
+    return ret;
 }
 
 EXPORT int my_lzma_stream_encoder(x64emu_t* emu, lzma_stream_t* stream, void* filters, int check)
 {
-    // not restoring the allocator after, so lzma_code and lzma_end can be used without "GOM" wrapping
-    if(stream->allocator)
-        wrap_alloc_struct(stream->allocator, stream->allocator);
-    return my->lzma_stream_encoder(stream, filters, check);
+    wrap_alloc_struct(stream->allocator);
+    int ret = my->lzma_stream_encoder(stream, filters, check);
+    unwrap_alloc_struct(stream->allocator);
+    return ret;
 }
 
 
 EXPORT int my_lzma_easy_encoder(x64emu_t* emu, lzma_stream_t* stream, uint32_t precheck, uint32_t check)
 {
-    // not restoring the allocator after, so lzma_code and lzma_end can be used without "GOM" wrapping
-    if(stream->allocator)
-        wrap_alloc_struct(stream->allocator, stream->allocator);
-    return my->lzma_easy_encoder(stream, precheck, check);
+    wrap_alloc_struct(stream->allocator);
+    int ret = my->lzma_easy_encoder(stream, precheck, check);
+    unwrap_alloc_struct(stream->allocator);
+    return ret;
 }
 
 EXPORT int my_lzma_raw_encoder(x64emu_t* emu, lzma_stream_t* stream, void* filters)
 {
-    // not restoring the allocator after, so lzma_code and lzma_end can be used without "GOM" wrapping
-    if(stream->allocator)
-        wrap_alloc_struct(stream->allocator, stream->allocator);
-    return my->lzma_raw_encoder(stream, filters);
+    wrap_alloc_struct(stream->allocator);
+    int ret = my->lzma_raw_encoder(stream, filters);
+    unwrap_alloc_struct(stream->allocator);
+    return ret;
 }
 
 EXPORT int my_lzma_raw_decoder(x64emu_t* emu, lzma_stream_t* stream, void* filters)
 {
-    // not restoring the allocator after, so lzma_code and lzma_end can be used without "GOM" wrapping
-    if(stream->allocator)
-        wrap_alloc_struct(stream->allocator, stream->allocator);
-    return my->lzma_raw_decoder(stream, filters);
+    wrap_alloc_struct(stream->allocator);
+    int ret = my->lzma_raw_decoder(stream, filters);
+    unwrap_alloc_struct(stream->allocator);
+    return ret;
 }
 
 EXPORT int my_lzma_properties_decode(x64emu_t* emu, void* filters, lzma_allocator_t* allocator, void* props, size_t size)
 {
-    lzma_allocator_t alloc = {0};
-    wrap_alloc_struct(&alloc, allocator);
-    return my->lzma_properties_decode(filters, &alloc, props, size);
+    wrap_alloc_struct(allocator);
+    int ret = my->lzma_properties_decode(filters, allocator, props, size);
+    unwrap_alloc_struct(allocator);
+    return ret;
 }
 
 EXPORT int my_lzma_alone_decoder(x64emu_t* emu, lzma_stream_t* stream, uint64_t memlimit)
 {
-    // not restoring the allocator after, so lzma_code and lzma_end can be used without "GOM" wrapping
-    if(stream->allocator)
-        wrap_alloc_struct(stream->allocator, stream->allocator);
-    return my->lzma_alone_decoder(stream, memlimit);
+    wrap_alloc_struct(stream->allocator);
+    int ret = my->lzma_alone_decoder(stream, memlimit);
+    unwrap_alloc_struct(stream->allocator);
+    return ret;
 }
 
 EXPORT int my_lzma_alone_encoder(x64emu_t* emu, lzma_stream_t* stream, void* options)
 {
-    // not restoring the allocator after, so lzma_code and lzma_end can be used without "GOM" wrapping
-    if(stream->allocator)
-        wrap_alloc_struct(stream->allocator, stream->allocator);
-    return my->lzma_alone_encoder(stream, options);
+    wrap_alloc_struct(stream->allocator);
+    int ret = my->lzma_alone_encoder(stream, options);
+    unwrap_alloc_struct(stream->allocator);
+    return ret;
 }
 
 EXPORT int my_lzma_stream_encoder_mt(x64emu_t* emu, lzma_stream_t* stream, void* options)
 {
-    // not restoring the allocator after, so lzma_code and lzma_end can be used without "GOM" wrapping
-    if(stream->allocator)
-        wrap_alloc_struct(stream->allocator, stream->allocator);
-    return my->lzma_stream_encoder_mt(stream, options);
+    wrap_alloc_struct(stream->allocator);
+    int ret = my->lzma_stream_encoder_mt(stream, options);
+    unwrap_alloc_struct(stream->allocator);
+    return ret;
+}
+
+EXPORT int my_lzma_code(x64emu_t* emu, lzma_stream_t* stream, int a)
+{
+    wrap_alloc_struct(stream->allocator);
+    int ret = my->lzma_code(stream, a);
+    unwrap_alloc_struct(stream->allocator);
+    return ret;
+}
+
+EXPORT void my_lzma_end(x64emu_t* emu, lzma_stream_t* stream)
+{
+    wrap_alloc_struct(stream->allocator);
+    my->lzma_end(stream);
 }
 
 #define CUSTOM_INIT \
diff --git a/src/wrapped/wrappedlzma_private.h b/src/wrapped/wrappedlzma_private.h
index 274bf88a..6f8b4765 100755
--- a/src/wrapped/wrappedlzma_private.h
+++ b/src/wrapped/wrappedlzma_private.h
@@ -19,7 +19,7 @@ GO(lzma_block_uncomp_encode, uFppLppL)
 //GO(lzma_block_unpadded_size, 
 GO(lzma_check_is_supported, CFu)
 GO(lzma_check_size, uFu)
-GO(lzma_code, iFpi)
+GOM(lzma_code, iFEpi)
 GO(lzma_cputhreads, uFv)
 GO(lzma_crc32, uFpLu)
 GO(lzma_crc64, LFpLL)
@@ -27,7 +27,7 @@ GO(lzma_crc64, LFpLL)
 GO(lzma_easy_decoder_memusage, LFu)
 GOM(lzma_easy_encoder, iFEpui)
 //GO(lzma_easy_encoder_memusage, 
-GO(lzma_end, vFp)
+GOM(lzma_end, vFEp)
 GO(lzma_filter_decoder_is_supported, CFL)
 //GO(lzma_filter_encoder_is_supported, 
 //GO(lzma_filter_flags_decode,