From b31622b369618746e54242ebfcc305154bf0ef59 Mon Sep 17 00:00:00 2001
From: Maximilian Hils <git@maximilianhils.com>
Date: Fri, 13 May 2022 05:53:44 +0200
Subject: add `Connection.set_verify`, fix #255 (#1073)

* add `Connection.set_verify`, fix #255

* show that it works with cryptography main

* Revert "show that it works with cryptography main"

This reverts commit fb0136a8e5aa5d2c6e0c16f8f4ecee2f3c72a16b.

* make it black
---
 CHANGELOG.rst      |  4 ++++
 src/OpenSSL/SSL.py | 29 +++++++++++++++++++++++++++++
 tests/test_ssl.py  | 46 ++++++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 79 insertions(+)

diff --git a/CHANGELOG.rst b/CHANGELOG.rst
index d91b1a1..5f2589f 100644
--- a/CHANGELOG.rst
+++ b/CHANGELOG.rst
@@ -19,6 +19,10 @@ Deprecations:
 Changes:
 ^^^^^^^^
 
+- Add ``OpenSSL.SSL.Connection.set_verify`` and ``OpenSSL.SSL.Connection.get_verify_mode``
+  to override the context object's verification flags.
+  `#1073 <https://github.com/pyca/pyopenssl/pull/1073>`_
+
 22.0.0 (2022-01-29)
 -------------------
 
diff --git a/src/OpenSSL/SSL.py b/src/OpenSSL/SSL.py
index d100e6c..f470460 100644
--- a/src/OpenSSL/SSL.py
+++ b/src/OpenSSL/SSL.py
@@ -1745,6 +1745,35 @@ class Connection:
 
         return _ffi.string(name)
 
+    def set_verify(self, mode, callback=None):
+        """
+        Override the Context object's verification flags for this specific
+        connection. See :py:meth:`Context.set_verify` for details.
+        """
+        if not isinstance(mode, int):
+            raise TypeError("mode must be an integer")
+
+        if callback is None:
+            self._verify_helper = None
+            self._verify_callback = None
+            _lib.SSL_set_verify(self._ssl, mode, _ffi.NULL)
+        else:
+            if not callable(callback):
+                raise TypeError("callback must be callable")
+
+            self._verify_helper = _VerifyHelper(callback)
+            self._verify_callback = self._verify_helper.callback
+            _lib.SSL_set_verify(self._ssl, mode, self._verify_callback)
+
+    def get_verify_mode(self):
+        """
+        Retrieve the Connection object's verify mode, as set by
+        :meth:`set_verify`.
+
+        :return: The verify mode
+        """
+        return _lib.SSL_get_verify_mode(self._ssl)
+
     def set_ciphertext_mtu(self, mtu):
         """
         For DTLS, set the maximum UDP payload size (*not* including IP/UDP
diff --git a/tests/test_ssl.py b/tests/test_ssl.py
index 56748fa..5e69ace 100644
--- a/tests/test_ssl.py
+++ b/tests/test_ssl.py
@@ -2630,6 +2630,52 @@ class TestConnection:
         server = Connection(ctx, None)
         assert None is server.get_verified_chain()
 
+    def test_set_verify_overrides_context(self):
+        context = Context(SSLv23_METHOD)
+        context.set_verify(VERIFY_PEER)
+        conn = Connection(context, None)
+        conn.set_verify(VERIFY_NONE)
+
+        assert context.get_verify_mode() == VERIFY_PEER
+        assert conn.get_verify_mode() == VERIFY_NONE
+
+        with pytest.raises(TypeError):
+            conn.set_verify(None)
+
+        with pytest.raises(TypeError):
+            conn.set_verify(VERIFY_PEER, "not a callable")
+
+    def test_set_verify_callback_reference(self):
+        """
+        The callback for certificate verification should only be forgotten if
+        the context and all connections created by it do not use it anymore.
+        """
+
+        def callback(conn, cert, errnum, depth, ok):  # pragma: no cover
+            return ok
+
+        tracker = ref(callback)
+
+        context = Context(SSLv23_METHOD)
+        context.set_verify(VERIFY_PEER, callback)
+        del callback
+
+        conn = Connection(context, None)
+        context.set_verify(VERIFY_NONE)
+
+        collect()
+        collect()
+        assert tracker()
+
+        conn.set_verify(VERIFY_PEER, lambda conn, cert, errnum, depth, ok: ok)
+        collect()
+        collect()
+        callback = tracker()
+        if callback is not None:  # pragma: nocover
+            referrers = get_referrers(callback)
+            if len(referrers) > 1:
+                pytest.fail("Some references remain: %r" % (referrers,))
+
     def test_get_session_unconnected(self):
         """
         `Connection.get_session` returns `None` when used with an object
-- 
cgit v1.2.1