Rewrite FreeBSD/Dragonfly C code from scratch by using patch http://bugs.python.org/issue10882 patch. This should also fix issue 3.

2 files changed, 148 insertions, 127 deletions

diff --git a/sendfilemodule.c b/sendfilemodule.c
index ed87402..639fee5 100644
--- a/sendfilemodule.c
+++ b/sendfilemodule.c
@@ -24,139 +24,155 @@
 #include <Python.h>
 #include <stdlib.h>
 
+
+
+/* --- begin FreeBSD / Dragonfly --- */
 #if defined(__FreeBSD__) || defined(__DragonFly__)
 #include <sys/types.h>
 #include <sys/socket.h>
 #include <sys/uio.h>
 
-static PyObject *
-method_sendfile(PyObject *self, PyObject *args)
+int
+PyParse_off_t(PyObject* arg, void* addr)
 {
-    int fd, s, sts;
-    off_t offset, sbytes;
-    size_t nbytes;
-    PyObject *headers;
-    struct iovec *header_iovs;
-    struct iovec *footer_iovs;
-    struct sf_hdtr hdtr;
+#if !defined(HAVE_LARGEFILE_SUPPORT)
+    *((off_t*)addr) = PyLong_AsLong(arg);
+#else
+    *((off_t*)addr) = PyLong_Check(arg) ? PyLong_AsLongLong(arg)
+ : PyLong_AsLong(arg);
+#endif
+    if (PyErr_Occurred())
+        return 0;
+    return 1;
+}
 
-    headers = NULL;
 
-    if (!PyArg_ParseTuple(args, "iiLi|O:sendfile", &fd, &s, &offset, &nbytes, &headers))
-        return NULL;
+static int
+iov_setup(struct iovec **iov, Py_buffer **buf, PyObject *seq, int cnt, int type)
+{
+    int i, j;
+    *iov = PyMem_New(struct iovec, cnt);
+    if (*iov == NULL) {
+        PyErr_NoMemory();
+        return 0;
+    }
+    *buf = PyMem_New(Py_buffer, cnt);
+    if (*buf == NULL) {
+        PyMem_Del(*iov);
+        PyErr_NoMemory();
+        return 0;
+    }
 
-    if (headers != NULL) {
-        int i, listlen;
-        PyObject *string;
-        char *buf;
-        int headerlist_len;
-        int headerlist_string = 0;
-        int footerlist_len;
-        int footerlist_string = 0;
-        PyObject *headerlist;
-        PyObject *footerlist;
-
-        if (PyTuple_Check(headers) && PyTuple_Size(headers) > 1) {
-            //printf("arg is tuple length %d\n", PyTuple_Size(headers));
-            headerlist = PyTuple_GetItem(headers, 0);
-            if (PyList_Check(headerlist)) {
-                headerlist_len = PyList_Size(headerlist);
-            } else if (PyString_Check(headerlist)) {
-                headerlist_string = 1;
-                headerlist_len = 1;
-            } else {
-                headerlist_len = 0;
-            }
+    for (i = 0; i < cnt; i++) {
+        if (PyObject_GetBuffer(PySequence_GetItem(seq, i), &(*buf)[i],
+                type) == -1) {
+            PyMem_Del(*iov);
+            for (j = 0; j < i; j++) {
+                PyBuffer_Release(&(*buf)[j]);
+           }
+            PyMem_Del(*buf);
+            return 0;
+        }
+        (*iov)[i].iov_base = (*buf)[i].buf;
+        (*iov)[i].iov_len = (*buf)[i].len;
+    }
+    return 1;
+}
 
-            footerlist = PyTuple_GetItem(headers, 1);
-            if (PyList_Check(footerlist)) {
-                //printf("footer is list\n");
-                footerlist_len = PyList_Size(footerlist);
-            } else if (PyString_Check(footerlist)) {
-                //printf("footer is string\n");
-                footerlist_string = 1;
-                footerlist_len = 1;
-            } else {
-                //printf("footer is invalid\n");
-                footerlist_len = 0;
-            }
-        } else {
-            if (PyTuple_Check(headers)) {
-                headerlist = PyTuple_GetItem(headers, 0);
-            } else {
-                headerlist = headers;
-            }
-            if (PyList_Check(headerlist)) {
-            	headerlist_len = PyList_Size(headerlist);
-            } else if (PyString_Check(headerlist)) {
-                headerlist_string = 1;
-                headerlist_len = 1;
-            } else {
-                headerlist_len = 0;
-            }
+static void
+iov_cleanup(struct iovec *iov, Py_buffer *buf, int cnt)
+{
+    int i;
+    PyMem_Del(iov);
+    for (i = 0; i < cnt; i++) {
+        PyBuffer_Release(&buf[i]);
+    }
+    PyMem_Del(buf);
+}
 
-            footerlist_len = 0;
-            footer_iovs = 0;
+static PyObject *
+method_sendfile(PyObject *self, PyObject *args, PyObject *kwdict)
+{
+    int in, out;
+    Py_ssize_t ret;
+    Py_ssize_t len;
+    off_t offset;
+    PyObject *headers = NULL, *trailers = NULL;
+    Py_buffer *hbuf, *tbuf;
+    off_t sbytes;
+    struct sf_hdtr sf;
+    int flags = 0;
+    sf.headers = NULL;
+    sf.trailers = NULL;
+    static char *keywords[] = {"out", "in",
+                                "offset", "count",
+                                "headers", "trailers", "flags", NULL};
+
+    if (!PyArg_ParseTupleAndKeywords(args, kwdict, "iiO&n|OOi:sendfile",
+        keywords, &out, &in, PyParse_off_t, &offset, &len,
+        &headers, &trailers, &flags))
+            return NULL;
+
+    if (headers != NULL) {
+        if (!PySequence_Check(headers)) {
+            PyErr_SetString(PyExc_TypeError,
+                "sendfile() headers must be a sequence or None");
+            return NULL;
+        } else {
+            sf.hdr_cnt = PySequence_Size(headers);
+            if (sf.hdr_cnt > 0 && !iov_setup(&(sf.headers), &hbuf,
+                    headers, sf.hdr_cnt, PyBUF_SIMPLE))
+                return NULL;
         }
-        //printf("header list size and type: %d, %d\nfooter list size and type: %d, %d\n", headerlist_len, headerlist_string, footerlist_len, footerlist_string);
-
-        header_iovs = (struct iovec*) malloc( sizeof(struct iovec) * headerlist_len );
-        for (i=0; i < headerlist_len; i++) {
-            if (headerlist_string) {
-                //printf("found string\n");
-                string = headerlist;
-            } else {
-                //printf("found list\n");
-                string = PyList_GET_ITEM(headerlist, i);
-            }
-            buf = (char *) PyString_AsString(string);
-            //printf("buf: %s\n", buf);
-            header_iovs[i].iov_base = buf;
-            header_iovs[i].iov_len = PyString_GET_SIZE(string);
+    }
+    if (trailers != NULL) {
+        if (!PySequence_Check(trailers)) {
+            PyErr_SetString(PyExc_TypeError,
+                "sendfile() trailers must be a sequence or None");
+            return NULL;
+        } else {
+            sf.trl_cnt = PySequence_Size(trailers);
+            if (sf.trl_cnt > 0 && !iov_setup(&(sf.trailers), &tbuf,
+                    trailers, sf.trl_cnt, PyBUF_SIMPLE))
+                return NULL;
         }
+    }
 
-        footer_iovs = (struct iovec*) malloc( sizeof(struct iovec) * footerlist_len );
-        for (i=0; i < footerlist_len; i++) {
-            if (footerlist_string) {
-                //printf("found string\n");
-                string = footerlist;
-            } else {
-                //printf("found list\n");
-                string = PyList_GET_ITEM(footerlist, i);
-            }
-            buf = (char *) PyString_AsString(string);
-            //printf("buf: %s\n", buf);
-            footer_iovs[i].iov_base = buf;
-            footer_iovs[i].iov_len = PyString_GET_SIZE(string);
-        }
+    Py_BEGIN_ALLOW_THREADS
+    ret = sendfile(in, out, offset, len, &sf, &sbytes, flags);
+    Py_END_ALLOW_THREADS
 
-        hdtr.headers = header_iovs;
-        hdtr.hdr_cnt = headerlist_len;
-        hdtr.trailers = footer_iovs;
-        hdtr.trl_cnt = footerlist_len;
+    if (sf.headers != NULL)
+        iov_cleanup(sf.headers, hbuf, sf.hdr_cnt);
+    if (sf.trailers != NULL)
+         iov_cleanup(sf.trailers, tbuf, sf.trl_cnt);
 
-	Py_BEGIN_ALLOW_THREADS;
-        sts = sendfile(s, fd, (off_t) offset, (size_t) nbytes, &hdtr, (off_t *) &sbytes, 0);
-	Py_END_ALLOW_THREADS;
-        free(header_iovs);
-        free(footer_iovs);
-    } else {
-	Py_BEGIN_ALLOW_THREADS;
-        sts = sendfile(s, fd, (off_t) offset, (size_t) nbytes, NULL, (off_t *) &sbytes, 0);
-	Py_END_ALLOW_THREADS;
-    }
-    if (sts == -1) {
-        if (errno == EAGAIN || errno == EINTR) {
-            return Py_BuildValue("LL", offset + nbytes, sbytes);
-        } else {
-            PyErr_SetFromErrno(PyExc_OSError);
-            return NULL;
+    if (ret < 0) {
+        if ((errno == EAGAIN) || (errno == EBUSY)) {
+            if (sbytes != 0) {
+                goto done;
+            }
+            else {
+                // upper application is supposed to retry
+                PyErr_SetFromErrno(PyExc_OSError);
+                return NULL;
+            }
         }
-    } else {
-        return Py_BuildValue("LL", offset + nbytes, sbytes);
+        PyErr_SetFromErrno(PyExc_OSError);
+        return NULL;
     }
+    goto done;
+
+done:
+    #if !defined(HAVE_LARGEFILE_SUPPORT)
+        return Py_BuildValue("ll", sbytes, offset + sbytes);
+    #else
+        return Py_BuildValue("LL", sbytes, offset + sbytes);
+    #endif
 }
+/* --- end FreeBSD / Dragonfly --- */
 
+/* --- begin AIX --- */
 #elif defined(_AIX)
 #include <sys/socket.h>
 
@@ -214,7 +230,9 @@ method_sendfile(PyObject *self, PyObject *args)
         return Py_BuildValue("kL", sts, offset);
     }
 }
+/* --- end AIX --- */
 
+/* --- start Linux --- */
 
 #elif defined (__linux__)
 #include <sys/sendfile.h>
@@ -244,7 +262,7 @@ method_sendfile(PyObject *self, PyObject *args)
 #endif
 }
 
-#endif
+#endif  /* --- end Linux --- */
 
 static PyMethodDef SendfileMethods[] = {
     {"sendfile",  method_sendfile, METH_VARARGS,
diff --git a/test_sendfile.py b/test_sendfile.py
index 8590381..5b4168b 100644
--- a/test_sendfile.py
+++ b/test_sendfile.py
@@ -108,7 +108,6 @@ def sendfile_wrapper(sock, file, offset, nbytes):
             else:
                 raise
         else:
-            assert sent <= nbytes
             assert (new_offset - offset) <= sent
             return (sent, new_offset)
 
@@ -134,27 +133,38 @@ class TestSendfile(unittest.TestCase):
 
     def test_send_whole_file(self):
         # normal send
-        sent = 0
+        total_sent = 0
         offset = 0
+        nbytes = 4096
         while 1:
-            sent, offset = sendfile_wrapper(self.sockno, self.fileno, offset, 4096)
+            sent, offset = sendfile_wrapper(self.sockno, self.fileno, offset, nbytes)
             if sent == 0:
                 break
+            total_sent += sent
+            self.assertTrue(sent <= nbytes)
+            self.assertEqual(offset, total_sent)
+
+        self.assertEqual(total_sent, len(DATA))
         time.sleep(.1)
         data = self.server.handler_instance.get_data()
         self.assertEqual(hash(data), hash(DATA))
 
     def test_send_at_certain_offset(self):
         # start sending a file at a certain offset
-        sent = 0
+        total_sent = 0
         offset = len(DATA) / 2
+        nbytes = 4096
         while 1:
-            sent, offset = sendfile_wrapper(self.sockno, self.fileno, offset, 4096)
+            sent, offset = sendfile_wrapper(self.sockno, self.fileno, offset, nbytes)
             if sent == 0:
                 break
+            total_sent += sent
+            self.assertTrue(sent <= nbytes)
+
         time.sleep(.1)
         data = self.server.handler_instance.get_data()
         expected = DATA[len(DATA) / 2:]
+        self.assertEqual(total_sent, len(expected))
         self.assertEqual(hash(data), hash(expected))
 
     def test_offset_overflow(self):
@@ -166,6 +176,7 @@ class TestSendfile(unittest.TestCase):
         data = self.server.handler_instance.get_data()
         self.assertEqual(data, '')
 
+
     def test_invalid_offset(self):
         try:
             sendfile.sendfile(self.sockno, self.fileno, -1, 4096)
@@ -174,14 +185,6 @@ class TestSendfile(unittest.TestCase):
         else:
             self.fail("exception not raised")
 
-    def test_invalid_nbytes(self):
-        try:
-            sendfile.sendfile(self.sockno, self.fileno, 0, -1)
-        except OSError, err:
-            self.assertEqual(err[0], errno.EINVAL)
-        else:
-            self.fail("exception not raised")
-
 
 def test_main():
     tests = [TestSendfile]