Add tests for socket audit hook leak

Author: aisk

Lib/test/test_socket.py

@@ -2,7 +2,7 @@
 from unittest import mock
 from test import support
 from test.support import (
-    cpython_only, is_apple, os_helper, refleak_helper, socket_helper, threading_helper
+    cpython_only, is_apple, os_helper, refleak_helper, script_helper, socket_helper, threading_helper
 )
 from test.support.import_helper import ensure_lazy_imports
 import _thread as thread
@@ -28,6 +28,7 @@
 import struct
 import sys
 import tempfile
+import textwrap
 import threading
 import time
 import traceback
@@ -7498,6 +7499,57 @@ def close_fds(fds):
             self.assertEqual(data,  str(index).encode())
 
 
+@support.requires_subprocess()
+@unittest.skipUnless(hasattr(sys, "gettotalrefcount"),
+                     "requires sys.gettotalrefcount()")
+class AuditHookLeakTests(unittest.TestCase):
+    # gh-146245: Reference and buffer may leaks in audit hook's failures path.
+
+    def test_getaddrinfo_audit_hook_leak(self):
+        code = textwrap.dedent("""
+            import socket
+            import sys
+            import gc
+            sys.addaudithook(lambda *a: (_ for _ in ()).throw(RuntimeError("audit")))
+            gc.collect()
+            before = sys.gettotalrefcount()
+            for _ in range(100):
+                try:
+                    socket.getaddrinfo(None, 80)
+                except RuntimeError:
+                    pass
+            gc.collect()
+            after = sys.gettotalrefcount()
+            print(after - before)
+        """)
+        rc, out, err = script_helper.assert_python_ok("-c", code)
+        leaked = int(out.strip())
+        self.assertLessEqual(leaked, 2, f"Leaked {leaked} references")
+
+    def test_sendto_audit_hook_leak(self):
+        code = textwrap.dedent("""
+            import socket
+            import sys
+            import gc
+            s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
+            sys.addaudithook(lambda *a: (_ for _ in ()).throw(RuntimeError("audit")))
+            gc.collect()
+            before = sys.gettotalrefcount()
+            for _ in range(100):
+                try:
+                    s.sendto(bytearray(b"x"), ("127.0.0.1", 80))
+                except RuntimeError:
+                    pass
+            gc.collect()
+            after = sys.gettotalrefcount()
+            s.close()
+            print(after - before)
+        """)
+        rc, out, err = script_helper.assert_python_ok("-c", code)
+        leaked = int(out.strip())
+        self.assertLessEqual(leaked, 2, f"Leaked {leaked} references")
+
+
 class FreeThreadingTests(unittest.TestCase):
 
     def test_close_detach_race(self):