From 23dbae63e41747cb2c6246c10ad6b28c127d8ab3 Mon Sep 17 00:00:00 2001
From: Kiritow <1362050620@qq.com>
Date: Tue, 5 Jun 2018 13:13:06 +0800
Subject: [PATCH] Fix getpeername memory problem

---
 gsock.cpp | 53 ++++++++++++++++++++++++++++++-----------------------
 gsock.h   |  3 ++-
 2 files changed, 32 insertions(+), 24 deletions(-)

diff --git a/gsock.cpp b/gsock.cpp
index c7aca1e..469f42c 100644
--- a/gsock.cpp
+++ b/gsock.cpp
@@ -75,6 +75,19 @@ public:
     }
 } _init_winsock2_2_obj;
 
+static inline const char* get_family_name(int family)
+{
+	switch (family)
+	{
+	case AF_INET:
+		return "AF_INET";
+	case AF_INET6:
+		return "AF_INET6";
+	default:
+		return "Unknown";
+	}
+}
+
 struct vsock::_impl
 {
 	int sfd;
@@ -276,23 +289,30 @@ int sock::setrecvtime(int Second)
 //forgive me, but writing code in hospital is really not a good experience.
 using _sock_getname_callback_t = decltype(getsockname);
 
+union _sock_getname_pack
+{
+	sockaddr saddr;
+	sockaddr_in saddr4;
+	sockaddr_in6 saddr6;
+};
+
 static int _sock_getname_call(int sfd,std::string& ip,int& port,_sock_getname_callback_t fn)
 {
-	struct sockaddr saddr;
-	socklen_t saddrlen=sizeof(saddr);
-	memset(&saddr,0,saddrlen);
-	int ret=fn(sfd,&saddr,&saddrlen);
+	_sock_getname_pack pack;
+	socklen_t saddrlen=sizeof(pack);
+	memset(&pack,0,saddrlen);
+	int ret=fn(sfd,&pack.saddr,&saddrlen);
 	if(ret<0) return ret; //don't bother errno. stop here.
-	if (saddr.sa_family == AF_INET)
+	if (pack.saddr.sa_family == AF_INET)
 	{
-		struct sockaddr_in* paddr = (sockaddr_in*)&saddr;
+		struct sockaddr_in* paddr = &pack.saddr4;
 		char ip_buff[64] = { 0 };
 		const char* pret = inet_ntop(AF_INET, paddr, ip_buff, 64);
 		if (pret)
 		{
 			ip = std::string(ip_buff);
 			port = ntohs(paddr->sin_port);
-			return GSOCK_OK;
+			return 0;
 		}
 		else
 		{
@@ -300,16 +320,16 @@ static int _sock_getname_call(int sfd,std::string& ip,int& port,_sock_getname_ca
 			return GSOCK_ERROR_NTOP;
 		}
 	}
-	else if (saddr.sa_family == AF_INET6)
+	else if (pack.saddr.sa_family == AF_INET6)
 	{
-		struct sockaddr_in6* paddr = (sockaddr_in6*)&saddr;
+		struct sockaddr_in6* paddr = &pack.saddr6;
 		char ip_buff[128] = { 0 };
 		const char* pret = inet_ntop(AF_INET6, paddr, ip_buff, 128);
 		if (pret)
 		{
 			ip = std::string(ip_buff);
 			port = ntohs(paddr->sin6_port);
-			return GSOCK_OK;
+			return 1;
 		}
 		else
 		{
@@ -521,19 +541,6 @@ struct udpsock::_impl
 	}
 };
 
-static inline const char* get_family_name(int family)
-{
-	switch (family)
-	{
-	case AF_INET:
-		return "AF_INET";
-	case AF_INET6:
-		return "AF_INET6";
-	default:
-		return "Unknown";
-	}
-}
-
 udpsock::udpsock(int use_family) : _pp(new _impl)
 {
 	
diff --git a/gsock.h b/gsock.h
index 302ffb6..64cccde 100644
--- a/gsock.h
+++ b/gsock.h
@@ -62,7 +62,8 @@ public:
     int setrecvtime(int Second);
     
     /// Return:
-    /// 0: Success. No Error.
+	/// 0: Success. No Error. IPv4
+	/// 1: Success. No Error. IPv6
     /// -1: getlocalname() or getpeername() call error. See errno.
     /// -2: Socket not created.
     int getlocal(std::string& IPStr,int& Port);