Prevent message truncation by populate function

[ric-plt/lib/rmr.git] / src / bindings / rmr-python / rmr / rmr.py

diff --git a/src/bindings/rmr-python/rmr/rmr.py b/src/bindings/rmr-python/rmr/rmr.py
index 2ae40c8..99c2088 100644 (file)
--- a/src/bindings/rmr-python/rmr/rmr.py
+++ b/src/bindings/rmr-python/rmr/rmr.py
@@ -1,3 +1,4 @@
+# vim: expandtab ts=4 sw=4:
 # ==================================================================================
 #       Copyright (c) 2019 Nokia
 #       Copyright (c) 2018-2019 AT&T Intellectual Property.
@@ -19,6 +20,7 @@ import json
 from ctypes import RTLD_GLOBAL, Structure, c_int, POINTER, c_char, c_char_p, c_void_p, memmove, cast
 from ctypes import CDLL
 from ctypes import create_string_buffer
+from rmr.exceptions import BadBufferAllocation
 
 # https://docs.python.org/3.7/library/ctypes.html
 # https://stackoverflow.com/questions/2327344/ctypes-loading-a-c-shared-library-that-has-dependencies/30845750#30845750
@@ -214,12 +216,50 @@ _rmr_alloc_msg.argtypes = [c_void_p, c_int]
 _rmr_alloc_msg.restype = POINTER(rmr_mbuf_t)
 
 
-def rmr_alloc_msg(vctx, size):
+def rmr_alloc_msg(vctx, size, payload=None, gen_transaction_id=False, mtype=None, meid=None):
     """
     Refer to the rmr C documentation for rmr_alloc_msg
     extern rmr_mbuf_t* rmr_alloc_msg(void* vctx, int size)
+
+    if payload is not None, attempts to set the payload
+    if gen_transaction_id is True, it generates and sets a transaction id
+    if mtype is not None, sets the sbuf's message type
+    if meid is not None, sets the sbuf's meid
+
     """
-    return _rmr_alloc_msg(vctx, size)
+    sbuf = _rmr_alloc_msg(vctx, size)
+    # make sure it's good
+    try:
+        sbuf.contents
+        if payload:
+            set_payload_and_length(payload, sbuf)
+
+        if gen_transaction_id:
+            generate_and_set_transaction_id(sbuf)
+
+        if mtype:
+            sbuf.contents.mtype = mtype
+
+        if meid:
+            rmr_set_meid(sbuf, meid)
+
+        return sbuf
+
+    except ValueError:
+        raise BadBufferAllocation
+
+
+_rmr_realloc_payload = rmr_c_lib.rmr_realloc_payload
+_rmr_realloc_payload.argtypes = [POINTER(rmr_mbuf_t), c_int, c_int, c_int]  # new_len, copy, clone
+_rmr_realloc_payload.restype = POINTER(rmr_mbuf_t)
+
+
+def rmr_realloc_payload(ptr_mbuf, new_len, copy=False, clone=False):
+    """
+    Refer to the rmr C documentation for rmr_realloc_payload().
+    extern rmr_mbuf_t* rmr_realloc_payload(rmr_mbuf_t*, int, int, int)
+    """
+    return _rmr_realloc_payload(ptr_mbuf, new_len, copy, clone)
 
 
 _rmr_free_msg = rmr_c_lib.rmr_free_msg
@@ -230,7 +270,7 @@ _rmr_free_msg.restype = None
 def rmr_free_msg(mbuf):
     """
     Refer to the rmr C documentation for rmr_free_msg
-    extern void rmr_free_msg( rmr_mbuf_t* mbuf )
+    extern void rmr_free_msg(rmr_mbuf_t* mbuf )
     """
     if mbuf is not None:
         _rmr_free_msg(mbuf)
@@ -324,16 +364,13 @@ _rmr_bytes2meid.argtypes = [POINTER(rmr_mbuf_t), c_char_p, c_int]
 _rmr_bytes2meid.restype = c_int
 
 
-def rmr_bytes2meid(ptr_mbuf, src, length):
+def rmr_set_meid(ptr_mbuf, byte_str):
     """
     Refer to the rmr C documentation for rmr_bytes2meid
     extern int rmr_bytes2meid(rmr_mbuf_t* mbuf, unsigned char const* src, int len);
     """
-    return _rmr_bytes2meid(ptr_mbuf, src, length)
-
+    return _rmr_bytes2meid(ptr_mbuf, byte_str, len(byte_str))
 
-# this is an alias to rmr_bytes2meid using familiar set/get terminoloigy
-rmr_set_meid = rmr_bytes2meid
 
 # CAUTION:  Some of the C functions expect a mutable buffer to copy the bytes into;
 #           if there is a get_* function below, use it to set up and return the
@@ -364,7 +401,7 @@ def rmr_get_meid(ptr_mbuf):
     sz = _get_constants().get("RMR_MAX_MEID", 64)  # size for buffer to fill
     buf = create_string_buffer(sz)
     _rmr_get_meid(ptr_mbuf, buf)
-    return buf.value.decode()  # decode turns into a string
+    return buf.value
 
 
 _rmr_get_src = rmr_c_lib.rmr_get_src
@@ -438,7 +475,7 @@ def message_summary(ptr_mbuf):
         dict message summary
     """
     return {
-        "payload": get_payload(ptr_mbuf),
+        "payload": get_payload(ptr_mbuf) if ptr_mbuf.contents.state == RMR_OK else None,
         "payload length": ptr_mbuf.contents.len,
         "message type": ptr_mbuf.contents.mtype,
         "subscription id": ptr_mbuf.contents.sub_id,
@@ -464,6 +501,9 @@ def set_payload_and_length(byte_str, ptr_mbuf):
     ptr_mbuf: ctypes c_void_p
         Pointer to an rmr message buffer
     """
+    if rmr_payload_size(ptr_mbuf) < len(byte_str):  # existing message payload too small
+        ptr_mbuf = rmr_realloc_payload(ptr_mbuf, len(byte_str), True)
+
     memmove(ptr_mbuf.contents.payload, byte_str, len(byte_str))
     ptr_mbuf.contents.len = len(byte_str)