diff --git a/src/lib/python/isc/cc/session.py b/src/lib/python/isc/cc/session.py index f6b62653be..33a47bdfcb 100644 --- a/src/lib/python/isc/cc/session.py +++ b/src/lib/python/isc/cc/session.py @@ -72,7 +72,7 @@ class Session: self._lname = None self._closed = True - def sendmsg(self, env, msg = None): + def sendmsg(self, env, msg=None): with self._lock: if self._closed: raise SessionError("Session has been closed.") @@ -82,15 +82,24 @@ class Session: raise ProtocolError("Envelope too large") if type(msg) == dict: msg = isc.cc.message.to_wire(msg) - self._socket.setblocking(1) length = 2 + len(env); - if msg: + if msg is not None: length += len(msg) - self._socket.send(struct.pack("!I", length)) - self._socket.send(struct.pack("!H", len(env))) - self._socket.send(env) - if msg: - self._socket.send(msg) + + # Build entire message. + data = struct.pack("!I", length) + data += struct.pack("!H", len(env)) + data += env + if msg is not None: + data += msg + + # Send it in the blocking mode. On some systems send() may + # actually send only part of the data, so we need to repeat it + # until all data have been sent out. + self._socket.setblocking(1) + while len(data) > 0: + cc = self._socket.send(data) + data = data[cc:] def recvmsg(self, nonblock = True, seq = None): """Reads a message. If nonblock is true, and there is no diff --git a/src/lib/python/isc/cc/tests/session_test.py b/src/lib/python/isc/cc/tests/session_test.py index 772ed0c961..e589085189 100644 --- a/src/lib/python/isc/cc/tests/session_test.py +++ b/src/lib/python/isc/cc/tests/session_test.py @@ -29,6 +29,7 @@ class MySocket(): self.recvqueue = bytearray() self.sendqueue = bytearray() self._blocking = True + self.send_limit = None def connect(self, to): pass @@ -40,7 +41,14 @@ class MySocket(): self._blocking = val def send(self, data): - self.sendqueue.extend(data); + # If the upper limit is specified, only "send" up to the specified + # limit + if self.send_limit is not None and len(data) > self.send_limit: + self.sendqueue.extend(data[0:self.send_limit]) + return self.send_limit + else: + self.sendqueue.extend(data) + return len(data) def readsent(self, length): if length > len(self.sendqueue): @@ -101,6 +109,17 @@ class MySocket(): def gettimeout(self): return 0 + def set_send_limit(self, limit): + '''Specify the upper limit of the transmittable data at once. + + By default, the send() method of this class "sends" all given data. + If this method is called and the its parameter is not None, + subsequent calls to send() will only transmit the specified amount + of data. This can be used to emulate the situation where send() + on a real socket object results in partial write. + ''' + self.send_limit = limit + # # We subclass the Session class we're testing here, only # to override the __init__() method, which wants a socket, @@ -157,6 +176,16 @@ class testSession(unittest.TestCase): #print(sent) #self.assertRaises(SessionError, sess.sendmsg, {}, {"hello": "a"}) + def test_session_sendmsg_shortwrite(self): + sess = MySession() + # Specify the upper limit of the size that can be transmitted at + # a single send() call on the faked socket (10 is an arbitrary choice, + # just reasonably small). + sess._socket.set_send_limit(10) + sess.sendmsg({'to': 'someone', 'reply': 1}, {"hello": "a"}) + # The complete message should still have been transmitted in the end. + sent = sess._socket.readsentmsg(); + def recv_and_compare(self, session, bytes, env, msg): """Adds bytes to the recvqueue (which will be read by the session object, and compare the resultinv env and msg to