Add pb_release() function

5 files changed, 185 insertions, 38 deletions

diff --git a/pb_decode.c b/pb_decode.c
index 65e22af5..68497406 100644
--- a/pb_decode.c
+++ b/pb_decode.c
@@ -487,6 +487,23 @@ static bool checkreturn allocate_field(pb_istream_t *stream, void *pData, size_t
     *(void**)pData = ptr;
     return true;
 }
+
+/* Clear a newly allocated item in case it contains a pointer, or is a submessage. */
+static void initialize_pointer_field(void *pItem, pb_field_iterator_t *iter)
+{
+    if (PB_LTYPE(iter->pos->type) == PB_LTYPE_STRING)
+    {
+        *(char**)pItem = NULL;
+    }
+    else if (PB_LTYPE(iter->pos->type) == PB_LTYPE_BYTES)
+    {
+        memset(pItem, 0, iter->pos->data_size); 
+    }
+    else if (PB_LTYPE(iter->pos->type) == PB_LTYPE_SUBMESSAGE)
+    {
+        pb_message_set_to_defaults((const pb_field_t *) iter->pos->ptr, pItem);
+    }
+}
 #endif
 
 static bool checkreturn decode_pointer_field(pb_istream_t *stream, pb_wire_type_t wire_type, pb_field_iterator_t *iter)
@@ -515,6 +532,7 @@ static bool checkreturn decode_pointer_field(pb_istream_t *stream, pb_wire_type_
                 if (!allocate_field(stream, iter->pData, iter->pos->data_size, 1))
                     return false;
                 
+                initialize_pointer_field(*(void**)iter->pData, iter);
                 return func(stream, iter->pos, *(void**)iter->pData);
             }
     
@@ -550,6 +568,7 @@ static bool checkreturn decode_pointer_field(pb_istream_t *stream, pb_wire_type_
 
                     /* Decode the array entry */
                     pItem = *(uint8_t**)iter->pData + iter->pos->data_size * (*size);
+                    initialize_pointer_field(pItem, iter);
                     if (!func(&substream, iter->pos, pItem))
                     {
                         status = false;
@@ -567,26 +586,12 @@ static bool checkreturn decode_pointer_field(pb_istream_t *stream, pb_wire_type_
                 size_t *size = (size_t*)iter->pSize;
                 void *pItem;
                 
-                if (!allocate_field(stream, iter->pData, iter->pos->data_size, *size + 1))
+                (*size)++;
+                if (!allocate_field(stream, iter->pData, iter->pos->data_size, *size))
                     return false;
             
-                pItem = *(uint8_t**)iter->pData + iter->pos->data_size * (*size);
-                
-                /* Clear the new item in case it contains a pointer, or is a submessage. */
-                if (PB_LTYPE(type) == PB_LTYPE_STRING)
-                {
-                    *(char**)pItem = NULL;
-                }
-                else if (PB_LTYPE(type) == PB_LTYPE_BYTES)
-                {
-                    memset(pItem, 0, iter->pos->data_size); 
-                }
-                else if (PB_LTYPE(type) == PB_LTYPE_SUBMESSAGE)
-                {
-                    pb_message_set_to_defaults((const pb_field_t *) iter->pos->ptr, pItem);
-                }
-                
-                (*size)++;
+                pItem = *(uint8_t**)iter->pData + iter->pos->data_size * (*size - 1);
+                initialize_pointer_field(pItem, iter);
                 return func(stream, iter->pos, pItem);
             }
             
@@ -908,6 +913,79 @@ bool pb_decode_delimited(pb_istream_t *stream, const pb_field_t fields[], void *
     return status;
 }
 
+#ifdef PB_ENABLE_MALLOC
+void pb_release(const pb_field_t fields[], void *dest_struct)
+{
+    pb_field_iterator_t iter;
+    pb_field_init(&iter, fields, dest_struct);
+    
+    do
+    {
+        pb_type_t type;
+        type = iter.pos->type;
+    
+        /* Avoid crash on empty message types (zero fields) */
+        if (iter.pos->tag == 0)
+            continue;
+        
+        if (PB_ATYPE(type) == PB_ATYPE_POINTER)
+        {
+            if (PB_LTYPE(type) == PB_LTYPE_STRING &&
+                PB_HTYPE(type) == PB_HTYPE_REPEATED)
+            {
+                /* Release entries in repeated string array */
+                void **pItem = *(void***)iter.pData;
+                size_t count = *(size_t*)iter.pSize;
+                while (count--)
+                {
+                    free(*pItem);
+                    *pItem++ = NULL;
+                }
+            }
+            else if (PB_LTYPE(type) == PB_LTYPE_BYTES)
+            {
+                /* Release entries in repeated bytes array */
+                pb_bytes_ptr_t *pItem = *(pb_bytes_ptr_t**)iter.pData;
+                size_t count = (pItem ? 1 : 0);
+                
+                if (PB_HTYPE(type) == PB_HTYPE_REPEATED)
+                {
+                    count = *(size_t*)iter.pSize;   
+                }
+                
+                while (count--)
+                {
+                    free(pItem->bytes);
+                    pItem->bytes = NULL;
+                    pItem++;
+                }
+            }
+            else if (PB_LTYPE(type) == PB_LTYPE_SUBMESSAGE)
+            {
+                /* Release fields in submessages */
+                void *pItem = *(void**)iter.pData;
+                size_t count = (pItem ? 1 : 0);
+                
+                if (PB_HTYPE(type) == PB_HTYPE_REPEATED)
+                {
+                    count = *(size_t*)iter.pSize;   
+                }
+                
+                while (count--)
+                {
+                    pb_release((const pb_field_t*)iter.pos->ptr, pItem);
+                    pItem = (uint8_t*)pItem + iter.pos->data_size;
+                }
+            }
+            
+            /* Release main item */
+            free(*(void**)iter.pData);
+            *(void**)iter.pData = NULL;
+        }
+    } while (pb_field_next(&iter));
+}
+#endif
+
 /* Field decoders */
 
 bool pb_decode_svarint(pb_istream_t *stream, int64_t *dest)
diff --git a/pb_decode.h b/pb_decode.h
index f71b5f1b..2e1da5a6 100644
--- a/pb_decode.h
+++ b/pb_decode.h
@@ -82,6 +82,14 @@ bool pb_decode_noinit(pb_istream_t *stream, const pb_field_t fields[], void *des
  */
 bool pb_decode_delimited(pb_istream_t *stream, const pb_field_t fields[], void *dest_struct);
 
+#ifdef PB_ENABLE_MALLOC
+/* Release any allocated pointer fields. If you use dynamic allocation, you should
+ * call this for any decoded message when you are done with it. You also need to
+ * free messages even if pb_decode() returned with error.
+ */
+void pb_release(const pb_field_t fields[], void *dest_struct);
+#endif
+
 
 /**************************************
  * Functions for manipulating streams *
diff --git a/tests/SConstruct b/tests/SConstruct
index abc6e7cb..eedb694e 100644
--- a/tests/SConstruct
+++ b/tests/SConstruct
@@ -58,7 +58,12 @@ if not env.GetOption('clean'):
         if stdint: conf.env.Append(CPPDEFINES = {'HAVE_STDINT_H': 1})
         if stddef: conf.env.Append(CPPDEFINES = {'HAVE_STDDEF_H': 1})
         if string: conf.env.Append(CPPDEFINES = {'HAVE_STRING_H': 1})
-        
+    
+    # Check if we have mallinfo for memory leak tests
+    mallinfo = conf.CheckFunc('mallinfo', '#include <malloc.h>\n')
+    if mallinfo:
+        conf.env.Append(CPPDEFINES = {'HAVE_MALLINFO': 1})
+    
     # Check if we can use pkg-config to find protobuf include path
     status, output = conf.TryAction('pkg-config protobuf --variable=includedir > $TARGET')
     if status:
@@ -70,7 +75,7 @@ if not env.GetOption('clean'):
     if 'gcc' in env['CC']:
         if conf.CheckLib('mudflap'):
             conf.env.Append(CCFLAGS = '-fmudflap')
-            conf.env.Append(LINKFLAGS = '-lmudflap -fmudflap')
+            conf.env.Append(LINKFLAGS = '-fmudflap')
     
     # Check if we can use extra strict warning flags (only with GCC)
     extra = '-Wcast-qual -Wlogical-op -Wconversion'
diff --git a/tests/alltypes_pointer/SConscript b/tests/alltypes_pointer/SConscript
index 05b4e52d..45985ff5 100644
--- a/tests/alltypes_pointer/SConscript
+++ b/tests/alltypes_pointer/SConscript
@@ -4,16 +4,26 @@
 Import("env")
 
 # We need our own pb_decode.o for the malloc support
+env = env.Clone()
+env.Append(CPPDEFINES = {'PB_ENABLE_MALLOC': 1});
+
+# Disable libmudflap, because it will confuse e.g. valgrind
+# and other memory leak detection tools.
+if '-fmudflap' in env["CCFLAGS"]:
+    env["CCFLAGS"].remove("-fmudflap")
+    env["LINKFLAGS"].remove("-fmudflap")
+    env["LIBS"].remove("mudflap")
+
 strict = env.Clone()
 strict.Append(CFLAGS = strict['CORECFLAGS'])
-strict.Append(CPPDEFINES = {'PB_ENABLE_MALLOC': 1});
 strict.Object("pb_decode_with_malloc.o", "$NANOPB/pb_decode.c")
+strict.Object("pb_encode_with_malloc.o", "$NANOPB/pb_encode.c")
 
 c = Copy("$TARGET", "$SOURCE")
 env.Command("alltypes.proto", "#alltypes/alltypes.proto", c)
 
 env.NanopbProto(["alltypes", "alltypes.options"])
-enc = env.Program(["encode_alltypes_pointer.c", "alltypes.pb.c", "$COMMON/pb_encode.o"])
+enc = env.Program(["encode_alltypes_pointer.c", "alltypes.pb.c", "pb_encode_with_malloc.o"])
 dec = env.Program(["decode_alltypes_pointer.c", "alltypes.pb.c", "pb_decode_with_malloc.o"])
 
 refdec = "$BUILD/alltypes/decode_alltypes$PROGSUFFIX"
diff --git a/tests/alltypes_pointer/decode_alltypes_pointer.c b/tests/alltypes_pointer/decode_alltypes_pointer.c
index 3db48114..29495ac4 100644
--- a/tests/alltypes_pointer/decode_alltypes_pointer.c
+++ b/tests/alltypes_pointer/decode_alltypes_pointer.c
@@ -5,8 +5,12 @@
 #include "alltypes.pb.h"
 #include "test_helpers.h"
 
+#ifdef HAVE_MALLINFO
+#include <malloc.h>
+#endif
+
 #define TEST(x) if (!(x)) { \
-    printf("Test " #x " failed.\n"); \
+    fprintf(stderr, "Test " #x " failed.\n"); \
     status = false; \
     }
 
@@ -21,7 +25,10 @@ bool check_alltypes(pb_istream_t *stream, int mode)
     memset(&alltypes, 0xAA, sizeof(alltypes));
     
     if (!pb_decode(stream, AllTypes_fields, &alltypes))
+    {
+        pb_release(AllTypes_fields, &alltypes);
         return false;
+    }
     
     TEST(alltypes.req_int32     && *alltypes.req_int32         == -1001);
     TEST(alltypes.req_int64     && *alltypes.req_int64         == -1002);
@@ -184,31 +191,70 @@ bool check_alltypes(pb_istream_t *stream, int mode)
     TEST(alltypes.end == 1099);
 #endif
 
+    pb_release(AllTypes_fields, &alltypes);
+
     return status;
 }
 
 int main(int argc, char **argv)
 {
-    uint8_t buffer[1024];
-    size_t count;
-    pb_istream_t stream;
-
-    /* Whether to expect the optional values or the default values. */
-    int mode = (argc > 1) ? atoi(argv[1]) : 0;
+    bool status;
+    int orig_allocations;
     
-    /* Read the data into buffer */
-    SET_BINARY_MODE(stdin);
-    count = fread(buffer, 1, sizeof(buffer), stdin);
+#ifdef HAVE_MALLINFO
+    /* Dynamic library loader etc. may have some malloc()ed memory also. */
+    {
+        struct mallinfo m = mallinfo();
+        orig_allocations = m.uordblks;
+    }
+#endif
+
+    {    
+        uint8_t buffer[1024];
+        size_t count;
+        pb_istream_t stream;
+        
+        /* Whether to expect the optional values or the default values. */
+        int mode = (argc > 1) ? atoi(argv[1]) : 0;
+        
+        /* Read the data into buffer */
+        SET_BINARY_MODE(stdin);
+        count = fread(buffer, 1, sizeof(buffer), stdin);
+        
+        /* Construct a pb_istream_t for reading from the buffer */
+        stream = pb_istream_from_buffer(buffer, count);
+        
+        /* Decode and verify the message */
+        status = check_alltypes(&stream, mode);
+        
+        if (!status)
+            fprintf(stderr, "Parsing failed: %s\n", PB_GET_ERROR(&stream));
+    }
     
-    /* Construct a pb_istream_t for reading from the buffer */
-    stream = pb_istream_from_buffer(buffer, count);
+#ifdef HAVE_MALLINFO
+    /* Check for memory leaks */
+    {
+        struct mallinfo m = mallinfo();
+        int leak = m.uordblks - orig_allocations;
+        
+        if (leak > 0)
+        {
+            fprintf(stderr, "Memory leak: %d bytes\n", leak);
+            return 1;
+        }
+        else
+        {
+            fprintf(stderr, "Ok, no memory leaks\n");
+        }
+    }
+#endif
     
-    /* Decode and print out the stuff */
-    if (!check_alltypes(&stream, mode))
+    if (!status)
     {
-        printf("Parsing failed: %s\n", PB_GET_ERROR(&stream));
         return 1;
-    } else {
+    }
+    else
+    {
         return 0;
     }
 }