From c922afbfadae54c55cea9d0b34e060c9686e23b5 Mon Sep 17 00:00:00 2001
From: Shane Nolan <shanenullain@gmail.com>
Date: Tue, 12 Mar 2024 22:18:07 +0000
Subject: [PATCH] fix: tcp_keepalive socket

---
 botocore/args.py        | 33 ++++++++++++++++++++++--
 tests/unit/test_args.py | 57 ++++++++++++++++++++++++++++++++++++++---
 2 files changed, 84 insertions(+), 6 deletions(-)

diff --git a/botocore/args.py b/botocore/args.py
index dbbcbe8a99..ed717e5be2 100644
--- a/botocore/args.py
+++ b/botocore/args.py
@@ -484,8 +484,37 @@ def _compute_socket_options(self, scoped_config, client_config=None):
             scoped_config.get("tcp_keepalive", False)
         )
         # Enables TCP Keepalive if specified in client config object or shared config file.
-        if client_keepalive or scoped_keepalive:
-            socket_options.append((socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1))
+        if not client_keepalive and not scoped_keepalive:
+            return socket_options
+
+        seconds_in_a_minute = 60
+
+        client_read_timeout = client_config and client_config.read_timeout
+        scoped_read_timeout = scoped_config.get("read_timeout", None)
+
+        read_timeout = (
+            scoped_read_timeout if scoped_read_timeout else client_read_timeout
+        ) or seconds_in_a_minute
+
+        maximum_keepalive_probes = int(read_timeout / seconds_in_a_minute) or 1
+        keep_idle = (
+            seconds_in_a_minute
+            if read_timeout > seconds_in_a_minute
+            else read_timeout
+        )
+
+        socket_options.extend(
+            [
+                (socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1),
+                (socket.IPPROTO_TCP, socket.TCP_KEEPIDLE, keep_idle),
+                (socket.IPPROTO_TCP, socket.TCP_KEEPINTVL, keep_idle),
+                (
+                    socket.IPPROTO_TCP,
+                    socket.TCP_KEEPCNT,
+                    maximum_keepalive_probes,
+                ),
+            ],
+        )
         return socket_options
 
     def _compute_retry_config(self, config_kwargs):
diff --git a/tests/unit/test_args.py b/tests/unit/test_args.py
index 8f1c992422..aaf826843c 100644
--- a/tests/unit/test_args.py
+++ b/tests/unit/test_args.py
@@ -222,7 +222,12 @@ def test_tcp_keepalive_enabled_scoped_config(self):
             self.assert_create_endpoint_call(
                 m,
                 socket_options=self.default_socket_options
-                + [(socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1)],
+                + [
+                    (socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1),
+                    (socket.IPPROTO_TCP, socket.TCP_KEEPIDLE, 60),
+                    (socket.IPPROTO_TCP, socket.TCP_KEEPINTVL, 60),
+                    (socket.IPPROTO_TCP, socket.TCP_KEEPCNT, 1),
+                ],
             )
 
     def test_tcp_keepalive_not_specified(self):
@@ -247,7 +252,12 @@ def test_tcp_keepalive_enabled_if_set_anywhere(self):
             self.assert_create_endpoint_call(
                 m,
                 socket_options=self.default_socket_options
-                + [(socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1)],
+                + [
+                    (socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1),
+                    (socket.IPPROTO_TCP, socket.TCP_KEEPIDLE, 60),
+                    (socket.IPPROTO_TCP, socket.TCP_KEEPINTVL, 60),
+                    (socket.IPPROTO_TCP, socket.TCP_KEEPCNT, 1),
+                ],
             )
             self.call_get_client_args(
                 scoped_config={'tcp_keepalive': 'false'},
@@ -256,7 +266,41 @@ def test_tcp_keepalive_enabled_if_set_anywhere(self):
             self.assert_create_endpoint_call(
                 m,
                 socket_options=self.default_socket_options
-                + [(socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1)],
+                + [
+                    (socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1),
+                    (socket.IPPROTO_TCP, socket.TCP_KEEPIDLE, 60),
+                    (socket.IPPROTO_TCP, socket.TCP_KEEPINTVL, 60),
+                    (socket.IPPROTO_TCP, socket.TCP_KEEPCNT, 1),
+                ],
+            )
+            self.call_get_client_args(
+                scoped_config={'tcp_keepalive': 'false'},
+                client_config=Config(tcp_keepalive=True, read_timeout=30),
+            )
+            self.assert_create_endpoint_call(
+                m,
+                socket_options=self.default_socket_options
+                + [
+                    (socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1),
+                    (socket.IPPROTO_TCP, socket.TCP_KEEPIDLE, 30),
+                    (socket.IPPROTO_TCP, socket.TCP_KEEPINTVL, 30),
+                    (socket.IPPROTO_TCP, socket.TCP_KEEPCNT, 1),
+                ],
+                timeout=(60, 30),
+            )
+            self.call_get_client_args(
+                scoped_config={'tcp_keepalive': 'false', 'read_timeout': 900},
+                client_config=Config(tcp_keepalive=True),
+            )
+            self.assert_create_endpoint_call(
+                m,
+                socket_options=self.default_socket_options
+                + [
+                    (socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1),
+                    (socket.IPPROTO_TCP, socket.TCP_KEEPIDLE, 60),
+                    (socket.IPPROTO_TCP, socket.TCP_KEEPINTVL, 60),
+                    (socket.IPPROTO_TCP, socket.TCP_KEEPCNT, 15),
+                ],
             )
 
     def test_tcp_keepalive_explicitly_disabled(self):
@@ -274,7 +318,12 @@ def test_tcp_keepalive_enabled_case_insensitive(self):
             self.assert_create_endpoint_call(
                 m,
                 socket_options=self.default_socket_options
-                + [(socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1)],
+                + [
+                    (socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1),
+                    (socket.IPPROTO_TCP, socket.TCP_KEEPIDLE, 60),
+                    (socket.IPPROTO_TCP, socket.TCP_KEEPINTVL, 60),
+                    (socket.IPPROTO_TCP, socket.TCP_KEEPCNT, 1),
+                ],
             )
 
     def test_client_config_has_use_dualstack_endpoint_flag(self):