check self connect

IronsDu · Feb 12, 2019 · 53aedcc · 53aedcc
1 parent 587569d
commit 53aedcc
Show file tree

Hide file tree

Showing 3 changed files with 73 additions and 14 deletions.
diff --git a/src/brynet/net/Connector.cpp b/src/brynet/net/Connector.cpp
@@ -153,7 +153,7 @@ namespace brynet { namespace net {
 
             const sock fd = *(sock*)p;
             totalFds.insert(fd);
-            if (isConnectSuccess(fd, false))
+            if (isConnectSuccess(fd, false) && !brynet::net::base::IsSelfConnect(fd))
             {
                 successFds.insert(fd);
             }
@@ -240,9 +240,9 @@ namespace brynet { namespace net {
         sock clientfd = INVALID_SOCKET;
 
 #if defined PLATFORM_WINDOWS
-        const int check_error = WSAEWOULDBLOCK;
+        const int ExpectedError = WSAEWOULDBLOCK;
 #else
-        const int check_error = EINPROGRESS;
+        const int ExpectedError = EINPROGRESS;
 #endif
         int n = 0;
 
@@ -262,16 +262,20 @@ namespace brynet { namespace net {
         n = connect(clientfd, (struct sockaddr*)&server_addr, sizeof(struct sockaddr));
         if (n == 0)
         {
-            goto SUCCESS;
+            if (brynet::net::base::IsSelfConnect(clientfd))
+            {
+                goto FAILED;
+            }
+            else
+            {
+                goto SUCCESS;
+            }
         }
-
-        if (check_error != sErrno)
+        else if (sErrno != ExpectedError)
         {
-            brynet::net::base::SocketClose(clientfd);
-            clientfd = INVALID_SOCKET;
             goto FAILED;
         }
-
+        else
         {
             ConnectingInfo ci;
             ci.startConnectTime = std::chrono::steady_clock::now();
@@ -281,8 +285,9 @@ namespace brynet { namespace net {
 
             mConnectingInfos[clientfd] = ci;
             ox_fdset_add(mFDSet.get(), clientfd, WriteCheck);
+
+            return;
         }
-        return;
 
     SUCCESS:
         if (addr.getSuccessCB() != nullptr)
@@ -292,6 +297,11 @@ namespace brynet { namespace net {
         return;
 
     FAILED:
+        if (clientfd != INVALID_SOCKET)
+        {
+            brynet::net::base::SocketClose(clientfd);
+            clientfd = INVALID_SOCKET;
+        }
         if (addr.getFailedCB() != nullptr)
         {
             addr.getFailedCB()();

diff --git a/src/brynet/net/SocketLibFunction.cpp b/src/brynet/net/SocketLibFunction.cpp
@@ -139,7 +139,7 @@ namespace brynet { namespace net { namespace base {
     {
         InitSocket();
 
-        struct  sockaddr_in ip4Addr = { 0 };
+        struct sockaddr_in ip4Addr = { 0 };
         struct sockaddr_in6 ip6Addr = { 0 };
         struct sockaddr_in* paddr = &ip4Addr;
         int addrLen = sizeof(ip4Addr);
@@ -202,7 +202,7 @@ namespace brynet { namespace net { namespace base {
 #endif
     }
 
-    static std::string get_ip_str(const struct sockaddr *sa)
+    static std::string getIPString(const struct sockaddr *sa)
     {
         char tmp[INET6_ADDRSTRLEN] = { 0 };
         switch (sa->sa_family)
@@ -229,14 +229,14 @@ namespace brynet { namespace net { namespace base {
         int namelen = sizeof(name);
         if (getpeername(fd, (struct sockaddr*)&name, &namelen) == 0)
         {
-            return get_ip_str(&name);
+            return getIPString(&name);
         }
 #else
         struct sockaddr_in name;
         socklen_t namelen = sizeof(name);
         if (getpeername(fd, (struct sockaddr*)&name, &namelen) == 0)
         {
-            return get_ip_str((const struct sockaddr*)&name);
+            return getIPString((const struct sockaddr*)&name);
         }
 #endif
 
@@ -266,4 +266,52 @@ namespace brynet { namespace net { namespace base {
         return socket(af, type, protocol);
     }
 
+    static struct sockaddr_in6 getLocalAddr(sock sockfd)
+    {
+        struct sockaddr_in6 localaddr;
+        memset(&localaddr, 0, sizeof localaddr);
+        socklen_t addrlen = static_cast<socklen_t>(sizeof localaddr);
+        if (::getsockname(sockfd, (struct sockaddr*)(&localaddr), &addrlen) < 0)
+        {
+            return localaddr;
+        }
+        return localaddr;
+    }
+
+    static struct sockaddr_in6 getPeerAddr(sock sockfd)
+    {
+        struct sockaddr_in6 peeraddr;
+        memset(&peeraddr, 0, sizeof peeraddr);
+        socklen_t addrlen = static_cast<socklen_t>(sizeof peeraddr);
+        if (::getpeername(sockfd, (struct sockaddr*)(&peeraddr), &addrlen) < 0)
+        {
+            return peeraddr;
+        }
+        return peeraddr;
+    }
+
+    /* copy from muduo */
+    bool IsSelfConnect(sock fd)
+    {
+        struct sockaddr_in6 localaddr = getLocalAddr(fd);
+        struct sockaddr_in6 peeraddr = getPeerAddr(fd);
+
+        if (localaddr.sin6_family == AF_INET)
+        {
+            const struct sockaddr_in* laddr4 = reinterpret_cast<struct sockaddr_in*>(&localaddr);
+            const struct sockaddr_in* raddr4 = reinterpret_cast<struct sockaddr_in*>(&peeraddr);
+            return laddr4->sin_port == raddr4->sin_port
+                && laddr4->sin_addr.s_addr == raddr4->sin_addr.s_addr;
+        }
+        else if (localaddr.sin6_family == AF_INET6)
+        {
+            return localaddr.sin6_port == peeraddr.sin6_port
+                && memcmp(&localaddr.sin6_addr, &peeraddr.sin6_addr, sizeof localaddr.sin6_addr) == 0;
+        }
+        else
+        {
+            return false;
+        }
+    }
+
 } } }
diff --git a/src/brynet/net/SocketLibFunction.h b/src/brynet/net/SocketLibFunction.h
@@ -22,5 +22,6 @@ namespace brynet { namespace net { namespace base {
     int SocketSend(sock fd, const char* buffer, int len);
     sock Accept(sock listenSocket, struct sockaddr* addr, socklen_t* addrLen);
     sock SocketCreate(int af, int type, int protocol);
+	bool IsSelfConnect(sock fd);
 
 } } }