fix(serverless): close pinned SSRF sessions to prevent FD leaks

Address PR #533 review feedback.

- Tie each single-use SyncClientSession's lifecycle to its response via
  _bind_session_to_response, so closing the response (with-block or explicit)
  also closes the session; without it the session's pools/FDs leaked once per
  download in long-lived workers.
- Close the session on redirect hops and on request exceptions; make teardown
  best-effort so a cleanup failure cannot mask the caller's real error.
- Close response and session deterministically in the pinned-IP adapter test.
- Add tests for the request-exception path and the error-masking guarantee.

Author: deanq

runpod/serverless/utils/rp_ssrf.py

@@ -13,6 +13,7 @@
 import ipaddress
 import os
 import socket
+from contextlib import suppress
 from typing import Iterator, List, Optional
 from urllib.parse import urljoin, urlparse
 
@@ -189,6 +190,42 @@ def _build_pinned_session(pinned_ip: str) -> SyncClientSession:
     return session
 
 
+def _close_session_quietly(session: SyncClientSession) -> None:
+    """
+    Close a pinned session best-effort, tolerating teardown errors.
+
+    This runs only on cleanup paths (a failed request, a closing response) where
+    the caller's original exception is the meaningful one. A failure while
+    closing the session's connection pools must not mask that exception, so it is
+    suppressed rather than raised. Session.close() is idempotent and does no
+    actionable I/O, so nothing worth surfacing is lost.
+    """
+    with suppress(Exception):
+        session.close()
+
+
+def _bind_session_to_response(response, session: SyncClientSession) -> None:
+    """
+    Tie `session` cleanup to `response.close()`.
+
+    Each safe_get() call owns a single-use session; the caller only holds the
+    response. Closing a requests.Response releases its connection but not the
+    session (its adapters and connection pools), so without this the session
+    leaks one set of pooled sockets/FDs per download in a long-lived worker.
+    Wrapping close() makes `with safe_get(...) as r:` (and any direct
+    response.close()) tear down both.
+    """
+    original_close = response.close
+
+    def close_both(*args, **kwargs):
+        try:
+            original_close(*args, **kwargs)
+        finally:
+            _close_session_quietly(session)
+
+    response.close = close_both
+
+
 def _enforce_content_length(response, max_bytes: Optional[int]) -> None:
     if max_bytes is None:
         return
@@ -254,22 +291,32 @@ def safe_get(
         ips = resolve_and_validate(parsed.hostname, port)
 
         session = _build_pinned_session(ips[0])
-        response = session.get(
-            current,
-            headers=headers,
-            allow_redirects=False,
-            stream=stream,
-            timeout=timeout,
-        )
+        try:
+            response = session.get(
+                current,
+                headers=headers,
+                allow_redirects=False,
+                stream=stream,
+                timeout=timeout,
+            )
+        except BaseException:
+            _close_session_quietly(session)
+            raise
 
         if response.status_code in _REDIRECT_STATUS:
             location = response.headers.get("Location")
-            response.close()
+            try:
+                response.close()
+            finally:
+                _close_session_quietly(session)
             if not location:
                 raise SSRFError(f"redirect without Location header: {current}")
             current = urljoin(current, location)
             continue
 
+        # Bind before the cap check so a Content-Length breach (which closes
+        # the response) also tears down the session.
+        _bind_session_to_response(response, session)
         _enforce_content_length(response, max_bytes)
         return response
 

tests/test_serverless/test_utils/test_ssrf.py

@@ -23,13 +23,16 @@
 class _FakeResponse:
     """Minimal stand-in for requests.Response used by safe_get."""
 
-    def __init__(self, status_code, headers=None):
+    def __init__(self, status_code, headers=None, close_exc=None):
         self.status_code = status_code
         self.headers = headers or {}
         self.closed = False
+        self._close_exc = close_exc
 
     def close(self):
         self.closed = True
+        if self._close_exc is not None:
+            raise self._close_exc
 
 
 def _addrinfo(*ips):
@@ -145,6 +148,8 @@ def resolve(host, _port):
         with self.assertRaises(SSRFError):
             safe_get("https://example.com/image.jpg")
 
+        session.close.assert_called_once()  # redirect hop must not leak its session
+
     @patch("runpod.serverless.utils.rp_ssrf._build_pinned_session")
     @patch("runpod.serverless.utils.rp_ssrf.resolve_and_validate", return_value=["93.184.216.34"])
     def test_too_many_redirects(self, _mock_resolve, mock_session):
@@ -167,6 +172,8 @@ def test_content_length_over_cap_rejected(self, _mock_resolve, mock_session):
         with self.assertRaises(SSRFError):
             safe_get("https://example.com/big.bin", max_bytes=1024)
 
+        session.close.assert_called_once()  # cap breach must close the pinned session
+
     @patch("runpod.serverless.utils.rp_ssrf._build_pinned_session")
     @patch("runpod.serverless.utils.rp_ssrf.resolve_and_validate", return_value=["93.184.216.34"])
     def test_returns_response_on_success(self, _mock_resolve, mock_session):
@@ -177,6 +184,52 @@ def test_returns_response_on_success(self, _mock_resolve, mock_session):
 
         self.assertIs(safe_get("https://example.com/ok.jpg", max_bytes=4096), response)
 
+    @patch("runpod.serverless.utils.rp_ssrf._build_pinned_session")
+    @patch("runpod.serverless.utils.rp_ssrf.resolve_and_validate", return_value=["93.184.216.34"])
+    def test_closing_response_closes_session(self, _mock_resolve, mock_session):
+        """The session outlives safe_get() only until the caller closes the response."""
+        response = _FakeResponse(200, {"Content-Length": "500"})
+        session = MagicMock()
+        session.get.return_value = response
+        mock_session.return_value = session
+
+        result = safe_get("https://example.com/ok.jpg", max_bytes=4096)
+        session.close.assert_not_called()  # still open for the caller to stream
+
+        result.close()
+        self.assertTrue(response.closed)  # original close still runs
+        session.close.assert_called_once()  # and the session is torn down with it
+
+    @patch("runpod.serverless.utils.rp_ssrf._build_pinned_session")
+    @patch("runpod.serverless.utils.rp_ssrf.resolve_and_validate", return_value=["93.184.216.34"])
+    def test_request_error_closes_session(self, _mock_resolve, mock_session):
+        """A failed GET must tear down the session before the error propagates."""
+        session = MagicMock()
+        session.get.side_effect = RequestException("boom")
+        mock_session.return_value = session
+
+        with self.assertRaises(RequestException):
+            safe_get("https://example.com/x")
+
+        session.close.assert_called_once()
+
+    @patch("runpod.serverless.utils.rp_ssrf._build_pinned_session")
+    @patch("runpod.serverless.utils.rp_ssrf.resolve_and_validate", return_value=["93.184.216.34"])
+    def test_session_close_error_does_not_mask_response_error(self, _mock_resolve, mock_session):
+        """Cleanup is best-effort: the caller's close error wins over a session-teardown error."""
+        response = _FakeResponse(200, {"Content-Length": "500"}, close_exc=RuntimeError("primary"))
+        session = MagicMock()
+        session.get.return_value = response
+        session.close.side_effect = RuntimeError("secondary pool teardown")
+        mock_session.return_value = session
+
+        result = safe_get("https://example.com/ok.jpg", max_bytes=4096)
+        with self.assertRaises(RuntimeError) as ctx:
+            result.close()
+
+        self.assertEqual(str(ctx.exception), "primary")  # response error surfaces
+        session.close.assert_called_once()  # teardown was still attempted
+
 
 class TestIterContentCapped(unittest.TestCase):
     """Streaming byte counter aborts past the cap even when headers lie."""
@@ -232,10 +285,11 @@ def test_dials_pinned_ip_and_preserves_host_header(self):
         thread.start()
         try:
             # URL host is example.com (won't route to our server); pin to loopback.
-            session = _build_pinned_session("127.0.0.1")
-            response = session.get(f"http://example.com:{port}/path", timeout=5)
-            self.assertEqual(response.status_code, 200)
-            self.assertEqual(response.content, b"ok")
+            with _build_pinned_session("127.0.0.1") as session, session.get(
+                f"http://example.com:{port}/path", timeout=5
+            ) as response:
+                self.assertEqual(response.status_code, 200)
+                self.assertEqual(response.content, b"ok")
             self.assertEqual(_RecordingHandler.received_host, f"example.com:{port}")
         finally:
             server.shutdown()