diff --git a/direct/src/distributed/PyDatagram.py b/direct/src/distributed/PyDatagram.py index eebc06e440..f45c3c52c1 100755 --- a/direct/src/distributed/PyDatagram.py +++ b/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): diff --git a/direct/src/distributed/PyDatagramIterator.py b/direct/src/distributed/PyDatagramIterator.py index 60267a1ab9..6ce96e77e4 100755 --- a/direct/src/distributed/PyDatagramIterator.py +++ b/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 diff --git a/panda/src/express/datagram.I b/panda/src/express/datagram.I index 350fa3d30c..42cc6b9cf1 100644 --- a/panda/src/express/datagram.I +++ b/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); +} diff --git a/panda/src/express/datagram.h b/panda/src/express/datagram.h index 420175c114..9611645eed 100644 --- a/panda/src/express/datagram.h +++ b/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" diff --git a/panda/src/express/datagramIterator.I b/panda/src/express/datagramIterator.I index 763850b60c..140151089a 100644 --- a/panda/src/express/datagramIterator.I +++ b/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(); +} diff --git a/panda/src/express/datagramIterator.h b/panda/src/express/datagramIterator.h index 570e34c3c4..867c7de58c 100644 --- a/panda/src/express/datagramIterator.h +++ b/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); diff --git a/panda/src/recorder/socketStreamRecorder.cxx b/panda/src/recorder/socketStreamRecorder.cxx index 5ee6168d44..8418485b20 100644 --- a/panda/src/recorder/socketStreamRecorder.cxx +++ b/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())); } } diff --git a/tests/putil/test_datagram.py b/tests/putil/test_datagram.py index 171363109a..7919a1360f 100644 --- a/tests/putil/test_datagram.py +++ b/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."""