# HG changeset patch # User Kristján Valur Jónsson # Date 1332164699 0 # Node ID c5093eb6c15f75e645234678b8a86f9e951c59b4 # Parent 8943649201c57fd357560d5c6aca1d01fc8fa567 imported patch 2012-03-14_14-29-43_r75566+.diff diff --git a/Lib/test/test_socket.py b/Lib/test/test_socket.py --- a/Lib/test/test_socket.py +++ b/Lib/test/test_socket.py @@ -26,6 +26,10 @@ import fcntl except ImportError: fcntl = False +try: + import multiprocessing +except ImportError: + multiprocessing = False HOST = support.HOST MSG = 'Michael Gilfix was here\u1234\r\n'.encode('utf-8') ## test unicode string and carriage return @@ -4641,6 +4645,52 @@ self.checkNonblock(s, False) socket.setdefaulttimeout(t) +def MultiprocessingFunc(q): + #recreate socket + sdata = q.get() + message = q.get() + + s = socket.socket(0, 0, 0, sdata) + s2, c = s.accept() + + s2.sendall(message) + s2.close() + s.close() + +@unittest.skipUnless(os.name == "nt", "Windows specific") +@unittest.skipUnless(multiprocessing, "need multiprocessing") +class TestSocketDuplication(SocketTCPTest): + def testDuplicate(self): + #create process: + q = multiprocessing.Queue() + p = multiprocessing.Process(target=MultiprocessingFunc, args=(q,)) + p.start() + + #create process + data = self.serv.duplicate(p.pid) + + #pass the socket over + addr = self.serv.getsockname() + self.serv.close() + q.put(data) + + message = b"slapmahfro" + q.put(message) + + #connect + s = socket.create_connection(addr) + #listen for the data + m = [] + while True: + data = s.recv(100) + if not data: + break + m.append(data) + s.close() + received = b"".join(m) + self.assertEqual(received, message) + p.join() + def test_main(): tests = [GeneralModuleTests, BasicTCPTest, TCPCloserTest, TCPTimeoutTest, @@ -4698,6 +4748,7 @@ # These are slow when setitimer() is not available InterruptedRecvTimeoutTest, InterruptedSendTimeoutTest, + TestSocketDuplication, ]) thread_info = support.threading_setup() diff --git a/Modules/socketmodule.c b/Modules/socketmodule.c --- a/Modules/socketmodule.c +++ b/Modules/socketmodule.c @@ -381,7 +381,7 @@ return INVALID_SOCKET; return WSASocket(FROM_PROTOCOL_INFO, FROM_PROTOCOL_INFO, - FROM_PROTOCOL_INFO, &info, 0, 0); + FROM_PROTOCOL_INFO, &info, 0, WSA_FLAG_OVERLAPPED); } #define SOCKETCLOSE closesocket #else @@ -3770,6 +3770,35 @@ Control the socket with WSAIoctl syscall. Currently supported 'cmd' values are\n\ SIO_RCVALL: 'option' must be one of the socket.RCVALL_* constants.\n\ SIO_KEEPALIVE_VALS: 'option' is a tuple of (onoff, timeout, interval)."); +#endif + +#if defined(MS_WINDOWS) +static PyObject* +sock_duplicate(PySocketSockObject *s, PyObject *arg) +{ + WSAPROTOCOL_INFO info; + DWORD processId; + int result; + + if (!PyArg_ParseTuple(arg, "I", &processId)) + return NULL; + + Py_BEGIN_ALLOW_THREADS + result = WSADuplicateSocket(s->sock_fd, processId, &info); + Py_END_ALLOW_THREADS + if (result == SOCKET_ERROR) + return set_error(); + return PyBytes_FromStringAndSize((const char*)&info, sizeof(info)); +} +PyDoc_STRVAR(sock_duplicate_doc, +"duplicate(process_id) -> bytes\n\ +\n\ +Duplicate the socket for another process. The target process id\n\ +must be provided and the resulting string passed to the target\n\ +process. The target process can then use this string to recreate\n\ +the socket, by passing it as the fourth argument to the\n\ +socket() constructor."); + #endif @@ -3802,6 +3831,10 @@ {"ioctl", (PyCFunction)sock_ioctl, METH_VARARGS, sock_ioctl_doc}, #endif +#if defined(MS_WINDOWS) + {"duplicate", (PyCFunction)sock_duplicate, METH_VARARGS, + sock_duplicate_doc}, +#endif {"listen", (PyCFunction)sock_listen, METH_O, listen_doc}, {"recv", (PyCFunction)sock_recv, METH_VARARGS, @@ -3929,13 +3962,40 @@ return -1; if (fdobj != NULL && fdobj != Py_None) { - fd = PyLong_AsSocket_t(fdobj); - if (fd == (SOCKET_T)(-1) && PyErr_Occurred()) - return -1; - if (fd == INVALID_SOCKET) { - PyErr_SetString(PyExc_ValueError, - "can't use invalid socket value"); - return -1; +#ifdef MS_WINDOWS + /* recreate a socket that was duplicated */ + if (PyBytes_Check(fdobj)) { + WSAPROTOCOL_INFO info; + if (PyBytes_GET_SIZE(fdobj) != sizeof(info)) { + PyErr_Format(PyExc_ValueError, + "socket descriptor string has wrong size, " + "should be %zu bytes.", sizeof(info)); + return -1; + } + memcpy(&info, PyBytes_AS_STRING(fdobj), sizeof(info)); + Py_BEGIN_ALLOW_THREADS + fd = WSASocket(FROM_PROTOCOL_INFO, FROM_PROTOCOL_INFO, + FROM_PROTOCOL_INFO, &info, 0, WSA_FLAG_OVERLAPPED); + Py_END_ALLOW_THREADS + if (fd == INVALID_SOCKET) { + set_error(); + return -1; + } + family = info.iAddressFamily; + type = info.iSocketType; + proto = info.iProtocol; + } + else +#endif + { + fd = PyLong_AsSocket_t(fdobj); + if (fd == (SOCKET_T)(-1) && PyErr_Occurred()) + return -1; + if (fd == INVALID_SOCKET) { + PyErr_SetString(PyExc_ValueError, + "can't use invalid socket value"); + return -1; + } } } else {