express: add Datagram add_blob and add_blob32, et al.

This is for writing Python 2/3 agnostic code for writing binary data to a datagram, and reading from it using DatagramIterator.

direct/src/distributed/PyDatagram.py

@@ -25,8 +25,8 @@ class PyDatagram(Datagram):
         STUint64: (Datagram.addUint64, int),
         STFloat64: (Datagram.addFloat64, None),
         STString: (Datagram.addString, None),
-        STBlob: (Datagram.addString, None),
-        STBlob32: (Datagram.addString32, None),
+        STBlob: (Datagram.addBlob, None),
+        STBlob32: (Datagram.addBlob32, None),
         }
 
     #def addChannel(self, channelId):

direct/src/distributed/PyDatagramIterator.py

@@ -23,8 +23,8 @@ class PyDatagramIterator(DatagramIterator):
         STUint64: DatagramIterator.getUint64,
         STFloat64: DatagramIterator.getFloat64,
         STString: DatagramIterator.getString,
-        STBlob: DatagramIterator.getString,
-        STBlob32: DatagramIterator.getString32,
+        STBlob: DatagramIterator.getBlob,
+        STBlob32: DatagramIterator.getBlob32,
         }
 
     getChannel = DatagramIterator.getUint64

panda/src/express/datagram.I

@@ -270,6 +270,35 @@ add_fixed_string(const std::string &str, size_t size) {
   }
 }
 
+/**
+ * Adds a variable-length binary blob to the datagram.  This actually adds a
+ * count followed by n bytes.
+ */
+INLINE void Datagram::
+add_blob(const vector_uchar &data) {
+  // The max sendable size for a blob is 2^16.
+  nassertv(data.size() <= (uint16_t)0xffff);
+
+  // Blobs always are preceded by their size
+  add_uint16((uint16_t)data.size());
+
+  // Add the blob
+  append_data(data.data(), data.size());
+}
+
+/**
+ * Adds a variable-length binary blob to the datagram, using a 32-bit length
+ * field to allow very long blobs.
+ */
+INLINE void Datagram::
+add_blob32(const vector_uchar &data) {
+  // Blobs always are preceded by their size
+  add_uint32((uint32_t)data.size());
+
+  // Add the blob
+  append_data(data.data(), data.size());
+}
+
 /**
  * Appends some more raw data to the end of the datagram.
  */
@@ -445,3 +474,8 @@ INLINE void
 generic_write_datagram(Datagram &dest, const std::wstring &value) {
   dest.add_wstring(value);
 }
+
+INLINE void
+generic_write_datagram(Datagram &dest, const vector_uchar &value) {
+  dest.add_blob(value);
+}

panda/src/express/datagram.h

@@ -81,6 +81,9 @@ PUBLISHED:
   INLINE void add_fixed_string(const std::string &str, size_t size);
   void add_wstring(const std::wstring &str);
 
+  INLINE void add_blob(const vector_uchar &);
+  INLINE void add_blob32(const vector_uchar &);
+
   void pad_bytes(size_t size);
   void append_data(const void *data, size_t size);
   INLINE void append_data(const vector_uchar &data);
@@ -158,6 +161,8 @@ INLINE void
 generic_write_datagram(Datagram &dest, const std::string &value);
 INLINE void
 generic_write_datagram(Datagram &dest, const std::wstring &value);
+INLINE void
+generic_write_datagram(Datagram &dest, const vector_uchar &value);
 
 
 #include "datagram.I"

panda/src/express/datagramIterator.I

@@ -400,6 +400,22 @@ get_be_float64() {
   return tempvar;
 }
 
+/**
+ * Extracts a variable-length binary blob.
+ */
+INLINE vector_uchar DatagramIterator::
+get_blob() {
+  return extract_bytes(get_uint16());
+}
+
+/**
+ * Extracts a variable-length binary blob with a 32-bit size field.
+ */
+INLINE vector_uchar DatagramIterator::
+get_blob32() {
+  return extract_bytes(get_uint32());
+}
+
 /**
  * Skips over the indicated number of bytes in the datagram.
  */
@@ -485,3 +501,8 @@ INLINE void
 generic_read_datagram(std::wstring &result, DatagramIterator &source) {
   result = source.get_wstring();
 }
+
+INLINE void
+generic_read_datagram(vector_uchar &result, DatagramIterator &source) {
+  result = source.get_blob();
+}

panda/src/express/datagramIterator.h

@@ -61,6 +61,9 @@ PUBLISHED:
   std::string get_fixed_string(size_t size);
   std::wstring get_wstring();
 
+  INLINE vector_uchar get_blob();
+  INLINE vector_uchar get_blob32();
+
   INLINE void skip_bytes(size_t size);
   vector_uchar extract_bytes(size_t size);
   size_t extract_bytes(unsigned char *into, size_t size);

panda/src/recorder/socketStreamRecorder.cxx

@@ -82,10 +82,7 @@ play_frame(DatagramIterator &scan, BamReader *manager) {
 
   int num_packets = scan.get_uint16();
   for (int i = 0; i < num_packets; i++) {
-    size_t size = scan.get_uint16();
-    vector_uchar packet(size);
-    scan.extract_bytes(&packet[0], size);
-    _data.push_back(Datagram(std::move(packet)));
+    _data.push_back(Datagram(scan.get_blob()));
   }
 }
 

tests/putil/test_datagram.py

@@ -28,6 +28,9 @@ def datagram_small(request):
     dg.add_string32('this is another string')
     dg.add_string('this is yet a third string')
 
+    dg.add_blob(b'blob data \x00\xf2\xa0\x00\x00')
+    dg.add_blob32(b'\xc9\x8f\x00 test blob32')
+
     dg.add_stdfloat(800.2)
     dg.add_stdfloat(3.1415926)
     dg.add_stdfloat(2.7182818)
@@ -49,6 +52,9 @@ def datagram_small(request):
         assert dgi.get_string32() == 'this is another string'
         assert dgi.get_string() == 'this is yet a third string'
 
+        assert dgi.get_blob() == b'blob data \x00\xf2\xa0\x00\x00'
+        assert dgi.get_blob32() == b'\xc9\x8f\x00 test blob32'
+
         assert dgi.get_stdfloat() == pytest.approx(800.2)
         assert dgi.get_stdfloat() == pytest.approx(3.1415926)
         assert dgi.get_stdfloat() == pytest.approx(2.7182818)
@@ -88,6 +94,12 @@ def test_datagram_bytes():
     dgi.get_remaining_bytes() == b'abc\x00\xff123'
 
 
+def test_datagram_get_message():
+    dg = core.Datagram(b'abc\x00')
+    dg.append_data(b'\xff123')
+    assert dg.get_message() == b'abc\x00\xff123'
+
+
 def test_iterator(datagram_small):
     """This tests Datagram/DatagramIterator, and sort of serves as a self-check
     of the test fixtures too."""