Fix #168 - Add !t command to set timeout. (#193)

docs/users/queries.rst

@@ -29,6 +29,7 @@ IRRd style queries
   after a query has been sent. Queries are answered in the order they were
   submitted. Takes no parameters. In deviation from all other queries,
   this query will return no response at all.
+* ``!t<timeout>`` sets the timeout of the connection. The connection is closed when no activity on the connection has occurred for this many seconds. Valid values range from 1 to 1000. The default is 30 seconds. For ``!a`` queries, which can take several minutes, a higher timeout value should be set.
 * ``!a<as-set-name>`` recursively resolves an `as-set`, then resolves all
   combined unique prefixes originating from any of the ASes in the set. Returns
   both IPv4 and IPv6 prefixes. Can be filtered to either IPv4 or IPv6 with
@@ -40,7 +41,7 @@ IRRd style queries
   *Note*: this type of query can take very long to run, due to the amount of
   information it retrieves. Queries may take several minutes to resolve, and
   return up to 10-20 MB of text. Ensure that your client will not time out
-  in this period.
+  in this period, and that a long timeout has been set with ``!t``.
 * ``!gAS<asn>`` finds all IPv4 routes for an origin AS. Only distinct
   prefixes of the routes are returned, seperated by spaces.
 * ``!6AS<asn>`` finds all IPv6 routes for an origin AS. Only distinct

irrd/server/whois/protocol.py

@@ -20,7 +20,6 @@ class WhoisQueryReceiver(TimeoutMixin, LineOnlyReceiver):
     off to a query pipeline thread.
     """
     delimiter = b'\n'
-    time_out = 600
 
     def connectionMade(self) -> None:  # noqa: N802
         """
@@ -47,7 +46,7 @@ def connectionMade(self) -> None:  # noqa: N802
         )
         self.query_pipeline_thread.start()
 
-        self.setTimeout(self.time_out)
+        self.setTimeout(self.query_parser.timeout)
         logger.debug(f'{self.peer_str}: new connection opened')
 
     def lineReceived(self, line_bytes: bytes) -> None:  # noqa: N802
@@ -71,13 +70,14 @@ def return_response(self, response: bytes) -> None:
         connection if needed, and telling the query pipeline thread
         sending the response completed.
         """
-        self.resetTimeout()
         self.transport.write(response)
 
         if not self.query_parser.multiple_command_mode:
             self.transport.loseConnection()
             logger.debug(f'{self.peer_str}: auto-closed connection')
 
+        self.setTimeout(self.query_parser.timeout)
+
         self.query_pipeline_thread.ready_for_next_result()
 
     def lose_connection_callback(self) -> None:

irrd/server/whois/query_parser.py

@@ -48,6 +48,7 @@ def __init__(self, peer, peer_str: str) -> None:
         self.object_classes: List[str] = []
         self.user_agent: Optional[str] = None
         self.multiple_command_mode = False
+        self.timeout = 30
         self.key_fields_only = False
         self.peer = peer
         self.peer_str = peer_str
@@ -114,6 +115,8 @@ def handle_irrd_command(self, full_command: str) -> WhoisQueryResponse:
             self.multiple_command_mode = True
             result = None
             response_type = WhoisQueryResponseType.NO_RESPONSE
+        elif command == 'T':
+            self.handle_irrd_timeout_update(parameter)
         elif command == 'G':
             result = self.handle_irrd_routes_for_origin_v4(parameter)
             if not result:
@@ -159,6 +162,18 @@ def handle_irrd_command(self, full_command: str) -> WhoisQueryResponse:
             result=result,
         )
 
+    def handle_irrd_timeout_update(self, timeout: str) -> None:
+        """!timeout query - update timeout in connection"""
+        try:
+            timeout_value = int(timeout)
+        except ValueError:
+            raise WhoisQueryParserException(f'Invalid value for timeout: {timeout}')
+
+        if timeout_value > 0 and timeout_value <= 1000:
+            self.timeout = timeout_value
+        else:
+            raise WhoisQueryParserException(f'Invalid value for timeout: {timeout}')
+
     def handle_irrd_routes_for_origin_v4(self, origin: str) -> str:
         """!g query - find all originating IPv4 prefixes from an origin, e.g. !gAS65537"""
         return self._routes_for_origin('route', origin)

irrd/server/whois/tests/test_protocol.py

@@ -60,18 +60,21 @@ def test_whois_protocol_no_access_list(self, config_override, mock_pipeline_reac
         receiver.connectionMade()
         assert receiver.peer_str == '[127.0.0.1]:99999'
         assert receiver.query_pipeline_thread.started
+        assert receiver.timeOut == receiver.query_parser.timeout
         mock_transport.reset_mock()
 
         receiver.lineReceived(b' ')
         receiver.lineReceived(b' !v ')
         assert receiver.query_pipeline_thread.pipeline.get(block=False) == b' '
         assert receiver.query_pipeline_thread.pipeline.get(block=False) == b' !v '
 
+        receiver.query_parser.timeout = 5
         receiver.query_pipeline_thread.response_callback(b'response')
         assert mock_transport.mock_calls[0][0] == 'write'
         assert mock_transport.mock_calls[0][1][0] == b'response'
         assert mock_transport.mock_calls[1][0] == 'loseConnection'
         assert len(mock_transport.mock_calls) == 2
+        assert receiver.timeOut == receiver.query_parser.timeout
         mock_transport.reset_mock()
 
         receiver.query_pipeline_thread.lose_connection_callback()

irrd/server/whois/tests/test_query_parser.py

@@ -457,6 +457,24 @@ def test_multiple_command_mode(self, prepare_parser):
         assert not mock_dq.mock_calls
         assert parser.multiple_command_mode
 
+    def test_update_timeout(self, prepare_parser):
+        mock_dq, mock_dh, parser = prepare_parser
+
+        response = parser.handle_query('!t300')
+        assert response.response_type == WhoisQueryResponseType.SUCCESS
+        assert response.mode == WhoisQueryResponseMode.IRRD
+        assert not response.result
+        assert parser.timeout == 300
+
+        for invalid_value in ['foo', '-5', '1001']:
+            response = parser.handle_query(f'!t{invalid_value}')
+            assert response.response_type == WhoisQueryResponseType.ERROR
+            assert response.mode == WhoisQueryResponseMode.IRRD
+            assert response.result == f'Invalid value for timeout: {invalid_value}'
+
+        assert parser.timeout == 300
+        assert not mock_dq.mock_calls
+
     def test_routes_for_origin_v4(self, prepare_parser):
         mock_dq, mock_dh, parser = prepare_parser