From 7dbf42d899af7d243a25a12a741c45cce21a1a59 Mon Sep 17 00:00:00 2001
From: Sean Reifschneider <sean@realgo.com>
Date: Tue, 16 Mar 2021 15:52:55 -0600
Subject: [PATCH] Attempt to fix zts_recv for Python.

There were a few I'm attempting to fix in zts_py_recv():

Was allocating a static buffer of 4096, but taking whatever
length the user passed in.  This change allocates a PyBytes
object of the size the user requests.

Was converting to a string without giving a size.  Which probably
led to the UnicodeDecodeError I was seeing below, trying to
decode a character beyond the received bytes.  Would also lead
to problems reading binary data that included embedded NULs.
Used the number of bytes received as the size of the returned data.

Variable "err" was being used, changed that to "bytes_read" as
that's what lwip_recv() returns, with <0 bytes read indicating
error.

Read data was being returned as a Unicode string, leading to this
response when I tried talking to my SSH server:

    UnicodeDecodeError: 'utf-8' codec can't decode
    byte 0xa1 in position 42: invalid start byte

My ssh banner is 40 bytes long, so I think position 42 was outside
the received buffer.  Changed it to a PyBytes response as is
normal for network data.  This code was cribbed from the Python
socketmodule

This was producing this error, as it was still returning the tuple:

    SystemError: <built-in function zts_py_recv> returned a
    result with an error set

This indicates that the UnicodeDecodeError had set an exception,
but a tuple was being returned instead of NULL.

In the case of a lwip_recv() error, the error response was not
being sent back to the wrapper code.  Instead, a "return NULL;"
was done.  I changed this case to return the tuple (err, None)
to the wrapper.

NOTE: this code built using "./build.sh host-python release", but
there was some sort of build problem I didn't understand which
produced a _libzt.so that I couldn't import, due to:

    ImportError: dynamic module does not define module export
    function (PyInit__libzt)

So I don't have a way to test these changes.
---
 src/bindings/python/PythonSockets.cpp | 29 ++++++++++++++++++++-------
 1 file changed, 22 insertions(+), 7 deletions(-)

diff --git a/src/bindings/python/PythonSockets.cpp b/src/bindings/python/PythonSockets.cpp
index dc7aaf8..804bf1f 100644
--- a/src/bindings/python/PythonSockets.cpp
+++ b/src/bindings/python/PythonSockets.cpp
@@ -133,15 +133,30 @@ int zts_py_connect(int fd, int family, int type, PyObject *addr_obj)
 
 PyObject * zts_py_recv(int fd, int len, int flags)
 {
-	PyObject *t;
-	char buf[4096];
-	int err = zts_recv(fd, buf, len, flags);
-	if (err < 0) {
+	PyObject *t, *buf;
+	int bytes_read;
+
+	buf = PyBytes_FromStringAndSize((char *) 0, len);
+	if (buf == NULL)
 		return NULL;
-	}
+
+	bytes_read = zts_recv(fd, PyBytes_AS_STRING(buf), len, flags);
 	t = PyTuple_New(2);
-	PyTuple_SetItem(t, 0, PyLong_FromLong(err));
-	PyTuple_SetItem(t, 1, PyUnicode_FromString(buf));
+	PyTuple_SetItem(t, 0, PyLong_FromLong(bytes_read));
+
+	if (bytes_read < 0) {
+		Py_DECREF(buf);
+		Py_INCREF(Py_None);
+		PyTuple_SetItem(t, 1, Py_None);
+		Py_INCREF(t);
+		return t;
+	}
+
+	if (bytes_read != len) {
+		 _PyBytes_Resize(&buf, bytes_read);
+	}
+
+	PyTuple_SetItem(t, 1, buf);
 	Py_INCREF(t);
 	return t;
 }