From 2584799415449cc0be738dc7c7bf6a1a6d4bf486 Mon Sep 17 00:00:00 2001
From: SChernykh <sergey.v.chernykh@gmail.com>
Date: Sun, 3 Sep 2023 17:13:22 +0200
Subject: [PATCH] TCPServer: fixed ASAN code

---
 src/common.h       |  1 +
 src/tcp_server.cpp | 19 +++++++++++++++----
 src/tcp_server.h   |  2 ++
 3 files changed, 18 insertions(+), 4 deletions(-)

diff --git a/src/common.h b/src/common.h
index f52a02e..72d4dfe 100644
--- a/src/common.h
+++ b/src/common.h
@@ -97,6 +97,7 @@
 #endif
 
 #if __has_feature(address_sanitizer) || defined(__SANITIZE_ADDRESS__)
+#define P2POOL_ASAN
 #define ASAN_POISON_MEMORY_REGION(addr, size) __asan_poison_memory_region((addr), (size))
 #define ASAN_UNPOISON_MEMORY_REGION(addr, size) __asan_unpoison_memory_region((addr), (size))
 extern "C" void __asan_poison_memory_region(void const volatile* addr, size_t size);
diff --git a/src/tcp_server.cpp b/src/tcp_server.cpp
index af5f9e8..882b649 100644
--- a/src/tcp_server.cpp
+++ b/src/tcp_server.cpp
@@ -597,8 +597,7 @@ void TCPServer::loop(void* data)
 		Client* c = server->m_allocateNewClient();
 
 		ASAN_POISON_MEMORY_REGION(wb, sizeof(WriteBuf));
-		ASAN_POISON_MEMORY_REGION(c, c->size());
-		ASAN_UNPOISON_MEMORY_REGION(&c->m_resetCounter, sizeof(c->m_resetCounter));
+		c->asan_poison_this();
 
 		server->m_writeBuffers.emplace(capacity, wb);
 		server->m_preallocatedClients.emplace_back(c);
@@ -967,8 +966,7 @@ TCPServer::Client* TCPServer::get_client()
 
 void TCPServer::return_client(Client* c)
 {
-	ASAN_POISON_MEMORY_REGION(c, c->size());
-	ASAN_UNPOISON_MEMORY_REGION(&c->m_resetCounter, sizeof(c->m_resetCounter));
+	c->asan_poison_this();
 	m_preallocatedClients.push_back(c);
 }
 
@@ -1267,4 +1265,17 @@ void TCPServer::Client::init_addr_string()
 	}
 }
 
+void TCPServer::Client::asan_poison_this() const
+{
+#ifdef P2POOL_ASAN
+	const uint8_t* begin = reinterpret_cast<const uint8_t*>(this);
+	const uint8_t* counter_begin = reinterpret_cast<const uint8_t*>(&m_resetCounter);
+	const uint8_t* counter_end = counter_begin + sizeof(m_resetCounter);
+	const uint8_t* end = begin + size();
+
+	ASAN_POISON_MEMORY_REGION(begin, counter_begin - begin);
+	ASAN_POISON_MEMORY_REGION(counter_end, end - counter_end);
+#endif
+}
+
 } // namespace p2pool
diff --git a/src/tcp_server.h b/src/tcp_server.h
index 489bc85..0eb7986 100644
--- a/src/tcp_server.h
+++ b/src/tcp_server.h
@@ -70,6 +70,8 @@ public:
 
 		void init_addr_string();
 
+		void asan_poison_this() const;
+
 		char* m_readBuf;
 		uint32_t m_readBufSize;