From e3b82a9f004a726d967cc7ad48118aca595275e2 Mon Sep 17 00:00:00 2001
From: Ryan Hall <ryan.hall@datadoghq.com>
Date: Wed, 4 Dec 2024 09:55:20 -0500
Subject: [PATCH] add additional testing for the stream socket addition

---
 datadog/dogstatsd/base.py                     | 18 ++++++++++-----
 .../dogstatsd/test_statsd_sender.py           | 21 ++++++++++++++++++
 tests/unit/dogstatsd/test_statsd.py           | 22 +++++++++++++++++++
 3 files changed, 55 insertions(+), 6 deletions(-)

diff --git a/datadog/dogstatsd/base.py b/datadog/dogstatsd/base.py
index 9689d8370..76b41be4c 100644
--- a/datadog/dogstatsd/base.py
+++ b/datadog/dogstatsd/base.py
@@ -503,9 +503,12 @@ def socket(self):
     def socket(self, new_socket):
         self._socket = new_socket
         if new_socket:
-            self._socket_kind = new_socket.getsockopt(socket.SOL_SOCKET, socket.SO_TYPE)
-        else:
-            self._socket_kind = None
+            try:
+                self._socket_kind = new_socket.getsockopt(socket.SOL_SOCKET, socket.SO_TYPE)
+                return
+            except AttributeError: # _socket can't have a type if it doesn't have sockopts
+                pass
+        self._socket_kind = None
 
     @property
     def telemetry_socket(self):
@@ -515,9 +518,12 @@ def telemetry_socket(self):
     def telemetry_socket(self, t_socket):
         self._telemetry_socket = t_socket
         if t_socket:
-            self._telemetry_socket_kind = t_socket.getsockopt(socket.SOL_SOCKET, socket.SO_TYPE)
-        else:
-            self._telemetry_socket_kind = None
+            try:
+                self._telemetry_socket_kind = t_socket.getsockopt(socket.SOL_SOCKET, socket.SO_TYPE)
+                return
+            except AttributeError:# _telemetry_socket can't have a kind if it doesn't have sockopts
+                pass
+        self._telemetry_socket_kind = None
 
     def enable_background_sender(self, sender_queue_size=0, sender_queue_timeout=0):
         """
diff --git a/tests/integration/dogstatsd/test_statsd_sender.py b/tests/integration/dogstatsd/test_statsd_sender.py
index 23d5c075b..910be81bd 100644
--- a/tests/integration/dogstatsd/test_statsd_sender.py
+++ b/tests/integration/dogstatsd/test_statsd_sender.py
@@ -54,6 +54,27 @@ def test_set_socket_timeout():
     statsd.close_socket()
     assert statsd.get_socket().gettimeout() == 1
 
+def test_stream_cleanup():
+    foo, _ = socket.socketpair(socket.AF_UNIX, socket.SOCK_STREAM, 0)
+
+    foo.settimeout(0)
+    statsd = DogStatsd(disable_buffering=True)
+    statsd.socket = foo
+    statsd.increment("test", 1)
+    statsd.increment("test", 1)
+    statsd.increment("test", 1)
+    assert statsd.socket is not None
+
+    foo.setsockopt(socket.SOL_SOCKET, socket.SO_SNDBUF, 1) # different os's have different mins, e.g. this sets the buffer size to 2304 on certain linux variants
+
+    try:
+        foo.sendall(os.urandom(5000)) # pre-emptively clog the buffer
+    except BlockingIOError:
+        pass
+
+    statsd.increment("test", 1)
+
+    assert statsd.socket is None
 
 @pytest.mark.parametrize(
     "disable_background_sender, disable_buffering",
diff --git a/tests/unit/dogstatsd/test_statsd.py b/tests/unit/dogstatsd/test_statsd.py
index afca54c9b..dafe0eda2 100644
--- a/tests/unit/dogstatsd/test_statsd.py
+++ b/tests/unit/dogstatsd/test_statsd.py
@@ -2116,3 +2116,25 @@ def inner():
         t.start()
         t.join(timeout=5)
         self.assertFalse(t.is_alive())
+
+    def test_fake_sockets(self):
+        """
+        To support legacy behavior wherein customers were able to set sockets directly as long as they supported a .send interface, 
+        ensure that arbitrary values passed to these properties are allowed and are handled correctly
+        """
+        statsd = DogStatsd(disable_buffering=True)
+
+        class fake_sock:
+            def __init__(self, id):
+                self.id = id
+            def send(self, _):
+                pass
+        statsd.socket = fake_sock(5)
+        statsd.telemetry_socket = fake_sock(10)
+
+        assert statsd.socket.id == 5
+        assert statsd.telemetry_socket.id == 10
+
+        statsd.increment("test", 1)
+
+        assert statsd.socket is not None