From 0fe911899b6e39ba7cd3ba9ee545f321da7bd354 Mon Sep 17 00:00:00 2001
From: Thulinma <jaron@vietors.com>
Date: Fri, 2 Sep 2016 14:33:15 +0200
Subject: [PATCH] Added UDP getDestPort function

---
 lib/socket.cpp | 35 +++++++++++++++++++++++++++--------
 lib/socket.h   |  1 +
 2 files changed, 28 insertions(+), 8 deletions(-)

diff --git a/lib/socket.cpp b/lib/socket.cpp
index 8458c911..54675a3b 100644
--- a/lib/socket.cpp
+++ b/lib/socket.cpp
@@ -1023,8 +1023,12 @@ Socket::UDPConnection::UDPConnection(const UDPConnection & o) {
     destAddr = 0;
     destAddr_size = 0;
   }
-  data = 0;
-  data_size = 0;
+  data = (char*)malloc(1024);
+  if (data){
+    data_size = 1024;
+  }else{
+    data_size = 0;
+  }
   data_len = 0;
 }
 
@@ -1103,6 +1107,19 @@ void Socket::UDPConnection::GetDestination(std::string & destIp, uint32_t & port
   DEBUG_MSG(DLVL_FAIL, "Could not get destination for UDP socket");
 }//Socket::UDPConnection GetDestination
 
+/// Returns the port number of the receiving end of this socket.
+/// Returns 0 on error.
+uint32_t Socket::UDPConnection::getDestPort() const{
+  if (!destAddr || !destAddr_size){return 0;}
+  if (((struct sockaddr_in *)destAddr)->sin_family == AF_INET6) {
+    return ntohs(((struct sockaddr_in6 *)destAddr)->sin6_port);
+  }
+  if (((struct sockaddr_in *)destAddr)->sin_family == AF_INET) {
+    return ntohs(((struct sockaddr_in *)destAddr)->sin_port);
+  }
+  return 0;
+}
+
 /// Sets the socket to be blocking if the parameters is true.
 /// Sets the socket to be non-blocking otherwise.
 void Socket::UDPConnection::setBlocking(bool blocking) {
@@ -1214,20 +1231,22 @@ bool Socket::UDPConnection::Receive() {
     data_size = SOCKETSIZE;
   }
 #endif
-  int r = recvfrom(sock, data, data_size, MSG_PEEK | MSG_TRUNC, 0, 0);
+  int r = recvfrom(sock, data, data_size, MSG_PEEK | MSG_TRUNC | MSG_DONTWAIT, 0, 0);
   if (r == -1) {
     if (errno != EAGAIN) {
-      INFO_MSG("Found an error: %d (%s)", errno, strerror(errno));
+      INFO_MSG("UDP receive: %d (%s)", errno, strerror(errno));
     }
     data_len = 0;
     return false;
   }
   if (data_size < (unsigned int)r) {
-    data = (char *)realloc(data, r);
-    if (data) {
+    char* tmp = (char*)realloc(data, r);
+    if (tmp) {
+      data = tmp;
       data_size = r;
-    } else {
-      data_size = 0;
+    }else{
+      FAIL_MSG("Could not resize socket buffer to %d bytes!", r);
+      return false;
     }
   }
   socklen_t destsize = destAddr_size;
diff --git a/lib/socket.h b/lib/socket.h
index 633d6ab5..78cd13cb 100644
--- a/lib/socket.h
+++ b/lib/socket.h
@@ -141,6 +141,7 @@ namespace Socket {
       void setBlocking(bool blocking);
       void SetDestination(std::string hostname, uint32_t port);
       void GetDestination(std::string & hostname, uint32_t & port);
+      uint32_t getDestPort() const;
       bool Receive();
       void SendNow(const std::string & data);
       void SendNow(const char * data);