From 13768dd4285609a486ba5efa145fea4a2a2f4119 Mon Sep 17 00:00:00 2001
From: wangPeak <37855278+wangPeak@users.noreply.github.com>
Date: Wed, 3 Apr 2024 11:26:39 +0800
Subject: [PATCH] Update SessionInfo.cpp

Unreleased Socket
---
 src/SessionInfo.cpp | 23 +++++++++++++++++++++--
 1 file changed, 21 insertions(+), 2 deletions(-)

diff --git a/src/SessionInfo.cpp b/src/SessionInfo.cpp
index cbfb179..2ab4215 100644
--- a/src/SessionInfo.cpp
+++ b/src/SessionInfo.cpp
@@ -8,6 +8,11 @@
 #include "utils/Logger.h"
 #include "eip/EncapsPacket.h"
 #include "eip/EncapsPacketFactory.h"
+#if defined (__unix__) || defined(__APPLE__)
+#include <unistd.h>
+#elif defined(_WIN32) || defined(WIN32) || defined(_WIN64)
+#include <winsock2.h>
+#endif
 
 namespace eipScanner {
 
@@ -17,7 +22,7 @@ namespace eipScanner {
 	using eip::EncapsPacketFactory;
 	using eip::EncapsStatusCodes;
 
-	SessionInfo::SessionInfo(const std::string &host, int port, const std::chrono::milliseconds &timeout)
+	SessionInfo::SessionInfo(const std::string &host, int port, const std::chrono::milliseconds &timeout) try
 			: _socket{sockets::EndPoint(host, port), timeout}
 			, _sessionHandle{0} {
 		_socket.setRecvTimeout(timeout);
@@ -33,6 +38,20 @@ namespace eipScanner {
 		_sessionHandle = packet.getSessionHandle();
 		Logger(LogLevel::INFO) << "Registered session " << _sessionHandle;
 	}
+	catch (std::exception& e) {
+		int sockedFd = this->_socket.getSocketFd();
+		if (sockedFd > 0)
+		{
+#if defined(_WIN32) || defined(WIN32) || defined(_WIN64)
+			shutdown(sockedFd, SD_BOTH);
+			closesocket(sockedFd);
+#else
+			shutdown(sockedFd, SHUT_RDWR);
+			close(sockedFd);
+#endif
+		}
+		throw e;
+	}
 
 	SessionInfo::SessionInfo(const std::string &host, int port)
 			: SessionInfo(host, port, std::chrono::milliseconds(1000)) {
@@ -76,4 +95,4 @@ namespace eipScanner {
 		return _socket.getRemoteEndPoint();
 	}
 
-}
\ No newline at end of file
+}