From bd4c12a6a30c0d98c980971ee68c27fd8f8c2a24 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jeremy=20Lain=C3=A9?= <jeremy.laine@m4x.org>
Date: Sun, 22 Jan 2023 21:03:26 +0100
Subject: [PATCH] Add support for DTLS timeouts

When performing a DTLS handshake, the DTLS state machine may need to be
updated based on the passage of time, for instance in response to packet
loss.

OpenSSL supports this by means of the `DTLSv1_get_timeout` and
`DTLSv1_handle_timeout` methods, both of which are included in
cryptography's bindings. This change adds Python wrappers for these
methods in the `Connection` class.
---
 src/OpenSSL/SSL.py | 31 ++++++++++++++++++++++++++++
 tests/test_ssl.py  | 50 ++++++++++++++++++++++++++++++++++++++++++----
 2 files changed, 77 insertions(+), 4 deletions(-)

diff --git a/src/OpenSSL/SSL.py b/src/OpenSSL/SSL.py
index dfbd1094e..54d71359f 100644
--- a/src/OpenSSL/SSL.py
+++ b/src/OpenSSL/SSL.py
@@ -2160,6 +2160,37 @@ def DTLSv1_listen(self):
         if result < 0:
             self._raise_ssl_error(self._ssl, result)
 
+    def DTLSv1_get_timeout(self):
+        """
+        Determine when the DTLS SSL object next needs to perform internal
+        processing due to the passage of time.
+
+        When the returned number of seconds have passed, the
+        :meth:`DTLSv1_handle_timeout` method needs to be called.
+
+        :return: The time left in seconds before the next timeout or `None`
+            if no timeout is currently active.
+        """
+        ptv_sec = _ffi.new("time_t *")
+        ptv_usec = _ffi.new("long *")
+        if _lib.Cryptography_DTLSv1_get_timeout(self._ssl, ptv_sec, ptv_usec):
+            return ptv_sec[0] + (ptv_usec[0] / 1000000)
+        else:
+            return None
+
+    def DTLSv1_handle_timeout(self):
+        """
+        Handles any timeout events which have become pending on a DTLS SSL
+        object.
+
+        :return: `True` if there was a pending timeout, `False` otherwise.
+        """
+        result = _lib.DTLSv1_handle_timeout(self._ssl)
+        if result < 0:
+            self._raise_ssl_error(self._ssl, result)
+        else:
+            return bool(result)
+
     def bio_shutdown(self):
         """
         If the Connection was created with a memory BIO, this method can be
diff --git a/tests/test_ssl.py b/tests/test_ssl.py
index a3617c70d..9152680ef 100644
--- a/tests/test_ssl.py
+++ b/tests/test_ssl.py
@@ -9,6 +9,7 @@
 import gc
 import select
 import sys
+import time
 import uuid
 from errno import (
     EAFNOSUPPORT,
@@ -31,6 +32,7 @@
 )
 from sys import getfilesystemencoding, platform
 from typing import Union
+from unittest.mock import patch
 from warnings import simplefilter
 from weakref import ref
 
@@ -4370,10 +4372,11 @@ class TestDTLS:
     # new versions of OpenSSL, this is unnecessary, but harmless, because the
     # DTLS state machine treats it like a network hiccup that duplicated a
     # packet, which DTLS is robust against.
-    def test_it_works_at_all(self):
-        # arbitrary number larger than any conceivable handshake volley
-        LARGE_BUFFER = 65536
 
+    # Arbitrary number larger than any conceivable handshake volley.
+    LARGE_BUFFER = 65536
+
+    def test_it_works_at_all(self):
         s_ctx = Context(DTLS_METHOD)
 
         def generate_cookie(ssl):
@@ -4404,7 +4407,7 @@ def verify_cookie(ssl, cookie):
 
         def pump_membio(label, source, sink):
             try:
-                chunk = source.bio_read(LARGE_BUFFER)
+                chunk = source.bio_read(self.LARGE_BUFFER)
             except WantReadError:
                 return False
             # I'm not sure this check is needed, but I'm not sure it's *not*
@@ -4484,3 +4487,42 @@ def pump():
             assert 0 < c.get_cleartext_mtu() < 500
         except NotImplementedError:  # OpenSSL 1.1.0 and earlier
             pass
+
+    def test_timeout(self):
+        c_ctx = Context(DTLS_METHOD)
+        c = Connection(c_ctx)
+
+        # No timeout before the handshake starts.
+        assert c.DTLSv1_get_timeout() is None
+        assert c.DTLSv1_handle_timeout() is False
+
+        # Start handshake and check there is data to send.
+        c.set_connect_state()
+        try:
+            c.do_handshake()
+        except SSL.WantReadError:
+            pass
+        assert c.bio_read(self.LARGE_BUFFER)
+
+        # There should now be an active timeout.
+        seconds = c.DTLSv1_get_timeout()
+        assert seconds is not None
+
+        # Handle the timeout and check there is data to send.
+        time.sleep(seconds)
+        assert c.DTLSv1_handle_timeout() is True
+        assert c.bio_read(self.LARGE_BUFFER)
+
+        # After the maximum number of allowed timeouts is reached,
+        # DTLSv1_handle_timeout will return -1.
+        #
+        # Testing this directly is prohibitively time consuming as the timeout
+        # duration is doubled on each retry, so the best we can do is to mock
+        # this condition.
+        with patch(
+            "OpenSSL._util.lib.DTLSv1_handle_timeout"
+        ) as mock_handle_timeout:
+            mock_handle_timeout.return_value = -1
+
+            with pytest.raises(Error):
+                c.DTLSv1_handle_timeout()