Fix clients being to greedy when reading response

Fixes #49

Author: lutostag

tests/unit/client/serial/test_rtu.py

@@ -1,12 +1,17 @@
 import pytest
 from serial import serial_for_url
 
-from umodbus.client.serial.rtu import send_message
+from umodbus.client.serial.rtu import send_message, read_coils
 
 
 def test_send_message_with_timeout():
-    """ Test if TimoutError is raised when serial port doesn't receive data."""
+    """ Test if TimoutError is raised when serial port doesn't receive enough
+    data.
+    """
     s = serial_for_url('loop://', timeout=0)
+    # As we are using a loop, the sent request will be read back as response.
+    # To test timeout use a request that needs more bytes for the response.
+    message = read_coils(slave_id=0, starting_address=1, quantity=40)
 
     with pytest.raises(ValueError):
-        send_message(b'', s)
+        send_message(message, s)

umodbus/client/serial/rtu.py

@@ -45,11 +45,14 @@
 import struct
 
 from umodbus.client.serial.redundancy_check import get_crc, validate_crc
-from umodbus.functions import (create_function_from_response_pdu, ReadCoils,
+from umodbus.functions import (create_function_from_response_pdu,
+                               expected_response_pdu_size_from_request_pdu,
+                               pdu_to_function_code_or_raise_error, ReadCoils,
                                ReadDiscreteInputs, ReadHoldingRegisters,
                                ReadInputRegisters, WriteSingleCoil,
                                WriteSingleRegister, WriteMultipleCoils,
                                WriteMultipleRegisters)
+from umodbus.utils import recv_exactly
 
 
 def _create_request_adu(slave_id, req_pdu):
@@ -189,12 +192,34 @@ def parse_response_adu(resp_adu, req_adu=None):
     return function.data
 
 
+def raise_for_exception_adu(resp_adu):
+    """ Check a response ADU for error
+
+    :param resp_adu: Response ADU.
+    :raises ModbusError: When a response contains an error code.
+    """
+    resp_pdu = resp_adu[1:-2]
+    pdu_to_function_code_or_raise_error(resp_pdu)
+
+
 def send_message(adu, serial_port):
-    """ Send Modbus message over serial port and parse response. """
+    """ Send ADU over serial to to server and return parsed response.
+
+    :param adu: Request ADU.
+    :param sock: Serial port instance.
+    :return: Parsed response from server.
+    """
     serial_port.write(adu)
-    response = serial_port.read(serial_port.in_waiting)
+    serial_port.flush()
+
+    # Check exception ADU (which is shorter than all other responses) first.
+    exception_adu_size = 5
+    response_error_adu = recv_exactly(serial_port.read, exception_adu_size)
+    raise_for_exception_adu(response_error_adu)
 
-    if len(response) == 0:
-        raise ValueError
+    expected_response_size = \
+        expected_response_pdu_size_from_request_pdu(adu[1:-2]) + 3
+    response_remainder = recv_exactly(
+        serial_port.read, expected_response_size - exception_adu_size)
 
-    return parse_response_adu(response, adu)
+    return parse_response_adu(response_error_adu + response_remainder, adu)

umodbus/client/tcp.py

@@ -85,11 +85,14 @@
 import struct
 from random import randint
 
-from umodbus.functions import (create_function_from_response_pdu, ReadCoils,
+from umodbus.functions import (create_function_from_response_pdu,
+                               expected_response_pdu_size_from_request_pdu,
+                               pdu_to_function_code_or_raise_error, ReadCoils,
                                ReadDiscreteInputs, ReadHoldingRegisters,
                                ReadInputRegisters, WriteSingleCoil,
                                WriteSingleRegister, WriteMultipleCoils,
                                WriteMultipleRegisters)
+from umodbus.utils import recv_exactly
 
 
 def _create_request_adu(slave_id, pdu):
@@ -234,13 +237,33 @@ def parse_response_adu(resp_adu, req_adu=None):
     return function.data
 
 
+def raise_for_exception_adu(resp_adu):
+    """ Check a response ADU for error
+
+    :param resp_adu: Response ADU.
+    :raises ModbusError: When a response contains an error code.
+    """
+    resp_pdu = resp_adu[7:]
+    pdu_to_function_code_or_raise_error(resp_pdu)
+
+
 def send_message(adu, sock):
     """ Send ADU over socket to to server and return parsed response.
 
     :param adu: Request ADU.
     :param sock: Socket instance.
     :return: Parsed response from server.
     """
-    sock.send(adu)
-    response = sock.recv(1024)
-    return parse_response_adu(response, adu)
+    sock.sendall(adu)
+
+    # Check exception ADU (which is shorter than all other responses) first.
+    exception_adu_size = 9
+    response_error_adu = recv_exactly(sock.recv, exception_adu_size)
+    raise_for_exception_adu(response_error_adu)
+
+    expected_response_size = \
+        expected_response_pdu_size_from_request_pdu(adu[7:]) + 7
+    response_remainder = recv_exactly(
+        sock.recv, expected_response_size - exception_adu_size)
+
+    return parse_response_adu(response_error_adu + response_remainder, adu)

umodbus/functions.py

@@ -55,8 +55,10 @@
     'b\\x06'
 
 """
+from __future__ import division
 import struct
 import inspect
+import math
 try:
     from functools import reduce
 except ImportError:
@@ -94,14 +96,12 @@
 REPORT_SERVER_ID = 17
 
 
-def create_function_from_response_pdu(resp_pdu, req_pdu=None):
-    """ Parse response PDU and return instance of :class:`ModbusFunction` or
+def pdu_to_function_code_or_raise_error(resp_pdu):
+    """ Parse response PDU and return of :class:`ModbusFunction` or
     raise error.
 
     :param resp_pdu: PDU of response.
-    :param  req_pdu: Request PDU, some functions require more info than in
-        response PDU in order to create instance. Default is None.
-    :return: Number or list with response data.
+    :return: Subclass of :class:`ModbusFunction` matching the response.
     :raises ModbusError: When response contains error code.
     """
     function_code = struct.unpack('>B', resp_pdu[0:1])[0]
@@ -110,6 +110,19 @@ def create_function_from_response_pdu(resp_pdu, req_pdu=None):
         error_code = struct.unpack('>B', resp_pdu[1:2])[0]
         raise error_code_to_exception_map[error_code]
 
+    return function_code
+
+
+def create_function_from_response_pdu(resp_pdu, req_pdu=None):
+    """ Parse response PDU and return instance of :class:`ModbusFunction` or
+    raise error.
+
+    :param resp_pdu: PDU of response.
+    :param  req_pdu: Request PDU, some functions require more info than in
+        response PDU in order to create instance. Default is None.
+    :return: Number or list with response data.
+    """
+    function_code = pdu_to_function_code_or_raise_error(resp_pdu)
     function = function_code_to_function_map[function_code]
 
     if req_pdu is not None and \
@@ -136,6 +149,15 @@ def create_function_from_request_pdu(pdu):
     return function_class.create_from_request_pdu(pdu)
 
 
+def expected_response_pdu_size_from_request_pdu(pdu):
+    """ Return number of bytes expected for response PDU, based on request PDU.
+
+    :param pdu: Array of bytes.
+    :return: number of bytes.
+    """
+    return create_function_from_request_pdu(pdu).expected_response_pdu_size
+
+
 class ModbusFunction(object):
     function_code = None
 
@@ -189,8 +211,8 @@ class ReadCoils(ModbusFunction):
 
     The reponse PDU varies in length, depending on the request. Each 8 coils
     require 1 byte. The amount of bytes needed represent status of the coils to
-    can be calculated with: bytes = round(quantity / 8) + 1. This response
-    contains (3 / 8 + 1) = 1 byte to describe the status of the coils. The
+    can be calculated with: bytes = ceil(quantity / 8). This response
+    contains ceil(3 / 8) = 1 byte to describe the status of the coils. The
     structure of a compleet response PDU looks like this:
 
         ================ ===============
@@ -264,6 +286,14 @@ def create_from_request_pdu(pdu):
 
         return instance
 
+    @property
+    def expected_response_pdu_size(self):
+        """ Return number of bytes expected for response PDU.
+
+        :return: number of bytes.
+        """
+        return 2 + int(math.ceil(self.quantity / 8))
+
     def create_response_pdu(self, data):
         """ Create response pdu.
 
@@ -394,8 +424,8 @@ class ReadDiscreteInputs(ModbusFunction):
 
     The reponse PDU varies in length, depending on the request. 8 inputs
     require 1 byte. The amount of bytes needed represent status of the inputs
-    to can be calculated with: bytes = round(quantity / 8) + 1. This response
-    contains (3 / 8 + 1) = 1 byte to describe the status of the inputs. The
+    to can be calculated with: bytes = ceil(quantity / 8). This response
+    contains ceil(3 / 8) = 1 byte to describe the status of the inputs. The
     structure of a compleet response PDU looks like this:
 
         ================ ===============
@@ -469,6 +499,14 @@ def create_from_request_pdu(pdu):
 
         return instance
 
+    @property
+    def expected_response_pdu_size(self):
+        """ Return number of bytes expected for response PDU.
+
+        :return: number of bytes.
+        """
+        return 2 + int(math.ceil(self.quantity / 8))
+
     def create_response_pdu(self, data):
         """ Create response pdu.
 
@@ -665,6 +703,14 @@ def create_from_request_pdu(pdu):
 
         return instance
 
+    @property
+    def expected_response_pdu_size(self):
+        """ Return number of bytes expected for response PDU.
+
+        :return: number of bytes.
+        """
+        return 2 + self.quantity * 2
+
     def create_response_pdu(self, data):
         """ Create response pdu.
 
@@ -837,6 +883,14 @@ def create_from_request_pdu(pdu):
 
         return instance
 
+    @property
+    def expected_response_pdu_size(self):
+        """ Return number of bytes expected for response PDU.
+
+        :return: number of bytes.
+        """
+        return 2 + self.quantity * 2
+
     def create_response_pdu(self, data):
         """ Create response pdu.
 
@@ -999,6 +1053,14 @@ def create_from_request_pdu(pdu):
 
         return instance
 
+    @property
+    def expected_response_pdu_size(self):
+        """ Return number of bytes expected for response PDU.
+
+        :return: number of bytes.
+        """
+        return 5
+
     def create_response_pdu(self):
         """ Create response pdu.
 
@@ -1141,6 +1203,14 @@ def create_from_request_pdu(pdu):
 
         return instance
 
+    @property
+    def expected_response_pdu_size(self):
+        """ Return number of bytes expected for response PDU.
+
+        :return: number of bytes.
+        """
+        return 5
+
     def create_response_pdu(self):
         fmt = '>BH' + conf.TYPE_CHAR
         return struct.pack(fmt, self.function_code, self.address, self.value)
@@ -1347,6 +1417,14 @@ def create_from_request_pdu(pdu):
 
         return instance
 
+    @property
+    def expected_response_pdu_size(self):
+        """ Return number of bytes expected for response PDU.
+
+        :return: number of bytes.
+        """
+        return 5
+
     def create_response_pdu(self):
         """ Create response pdu.
 
@@ -1491,6 +1569,14 @@ def create_from_request_pdu(pdu):
 
         return instance
 
+    @property
+    def expected_response_pdu_size(self):
+        """ Return number of bytes expected for response PDU.
+
+        :return: number of bytes.
+        """
+        return 5
+
     def create_response_pdu(self):
         """ Create response pdu.
 

umodbus/server/__init__.py

@@ -8,7 +8,7 @@
 from umodbus.functions import create_function_from_request_pdu
 from umodbus.exceptions import ModbusError, ServerDeviceFailureError
 from umodbus.utils import (get_function_code_from_request_pdu,
-                           pack_exception_pdu)
+                           pack_exception_pdu, recv_exactly)
 
 
 def route(self, slave_ids=None, function_codes=None, addresses=None):
@@ -38,13 +38,14 @@ class AbstractRequestHandler(BaseRequestHandler):
     def handle(self):
         try:
             while True:
-                request_adu = self.request.recv(1024)
-
-                # When client terminates connection length of request_adu is 0.
-                if len(request_adu) == 0:
+                try:
+                    mbap_header = recv_exactly(self.request.recv, 7)
+                    remaining = self.get_meta_data(mbap_header)['length'] - 1
+                    request_pdu = recv_exactly(self.request.recv, remaining)
+                except ValueError:
                     return
 
-                response_adu = self.process(request_adu)
+                response_adu = self.process(mbap_header + request_pdu)
                 self.respond(response_adu)
         except:
             import traceback

umodbus/utils.py

@@ -112,3 +112,32 @@ def inner(arg):
             cache[arg] = f(arg)
         return cache[arg]
     return inner
+
+
+def recv_exactly(recv_fn, size):
+    """ Use the function to read and return exactly number of bytes desired.
+
+    https://docs.python.org/3/howto/sockets.html#socket-programming-howto for
+    more information about why this is necessary.
+
+    :param recv_fn: Function that can return up to given bytes
+        (i.e. socket.recv, file.read)
+    :param size: Number of bytes to read.
+    :return: Byte string with length size.
+    :raises ValueError: Could not receive enough data (usually timeout).
+    """
+    recv_bytes = 0
+    chunks = []
+    while recv_bytes < size:
+        chunk = recv_fn(size - recv_bytes)
+        if len(chunk) == 0:  # when closed or empty
+            break
+        recv_bytes += len(chunk)
+        chunks.append(chunk)
+
+    response = b''.join(chunks)
+
+    if len(response) != size:
+        raise ValueError
+
+    return response