tools/gator/daemon/OlySocket.cpp

   1 /**
   2  * Copyright (C) ARM Limited 2010-2014. All rights reserved.
   3  *
   4  * This program is free software; you can redistribute it and/or modify
   5  * it under the terms of the GNU General Public License version 2 as
   6  * published by the Free Software Foundation.
   7  */
   8
   9 #include "OlySocket.h"
  10
  11 #include <stdio.h>
  12 #include <string.h>
  13 #ifdef WIN32
  14 #include <Winsock2.h>
  15 #include <ws2tcpip.h>
  16 #else
  17 #include <netinet/in.h>
  18 #include <sys/un.h>
  19 #include <unistd.h>
  20 #include <netdb.h>
  21 #include <fcntl.h>
  22 #endif
  23
  24 #include "Logging.h"
  25
  26 #ifdef WIN32
  27 #define CLOSE_SOCKET(x) closesocket(x)
  28 #define SHUTDOWN_RX_TX SD_BOTH
  29 #define snprintf       _snprintf
  30 #else
  31 #define CLOSE_SOCKET(x) close(x)
  32 #define SHUTDOWN_RX_TX SHUT_RDWR
  33 #endif
  34
  35 int socket_cloexec(int domain, int type, int protocol) {
  36 #ifdef SOCK_CLOEXEC
  37   return socket(domain, type | SOCK_CLOEXEC, protocol);
  38 #else
  39   int sock = socket(domain, type, protocol);
  40 #ifdef FD_CLOEXEC
  41   if (sock < 0) {
  42     return -1;
  43   }
  44   int fdf = fcntl(sock, F_GETFD);
  45   if ((fdf == -1) || (fcntl(sock, F_SETFD, fdf | FD_CLOEXEC) != 0)) {
  46     close(sock);
  47     return -1;
  48   }
  49 #endif
  50   return sock;
  51 #endif
  52 }
  53
  54 int accept_cloexec(int sockfd, struct sockaddr *addr, socklen_t *addrlen) {
  55   int sock;
  56 #ifdef SOCK_CLOEXEC
  57   sock = accept4(sockfd, addr, addrlen, SOCK_CLOEXEC);
  58   if (sock >= 0) {
  59     return sock;
  60   }
  61   // accept4 with SOCK_CLOEXEC may not work on all kernels, so fallback
  62 #endif
  63   sock = accept(sockfd, addr, addrlen);
  64 #ifdef FD_CLOEXEC
  65   if (sock < 0) {
  66     return -1;
  67   }
  68   int fdf = fcntl(sock, F_GETFD);
  69   if ((fdf == -1) || (fcntl(sock, F_SETFD, fdf | FD_CLOEXEC) != 0)) {
  70     close(sock);
  71     return -1;
  72   }
  73 #endif
  74   return sock;
  75 }
  76
  77 OlyServerSocket::OlyServerSocket(int port) {
  78 #ifdef WIN32
  79   WSADATA wsaData;
  80   if (WSAStartup(0x0202, &wsaData) != 0) {
  81     logg->logError(__FILE__, __LINE__, "Windows socket initialization failed");
  82     handleException();
  83   }
  84 #endif
  85
  86   createServerSocket(port);
  87 }
  88
  89 OlySocket::OlySocket(int socketID) : mSocketID(socketID) {
  90 }
  91
  92 #ifndef WIN32
  93
  94 #define MIN(A, B) ({ \
  95   const __typeof__(A) __a = A; \
  96   const __typeof__(B) __b = B; \
  97   __a > __b ? __b : __a; \
  98 })
  99
 100 OlyServerSocket::OlyServerSocket(const char* path, const size_t pathSize) {
 101   // Create socket
 102   mFDServer = socket_cloexec(PF_UNIX, SOCK_STREAM, 0);
 103   if (mFDServer < 0) {
 104     logg->logError(__FILE__, __LINE__, "Error creating server socket");
 105     handleException();
 106   }
 107
 108   // Create sockaddr_in structure, ensuring non-populated fields are zero
 109   struct sockaddr_un sockaddr;
 110   memset((void*)&sockaddr, 0, sizeof(sockaddr));
 111   sockaddr.sun_family = AF_UNIX;
 112   memcpy(sockaddr.sun_path, path, MIN(pathSize, sizeof(sockaddr.sun_path)));
 113   sockaddr.sun_path[sizeof(sockaddr.sun_path) - 1] = '\0';
 114
 115   // Bind the socket to an address
 116   if (bind(mFDServer, (const struct sockaddr*)&sockaddr, sizeof(sockaddr)) < 0) {
 117     logg->logError(__FILE__, __LINE__, "Binding of server socket failed.");
 118     handleException();
 119   }
 120
 121   // Listen for connections on this socket
 122   if (listen(mFDServer, 1) < 0) {
 123     logg->logError(__FILE__, __LINE__, "Listening of server socket failed");
 124     handleException();
 125   }
 126 }
 127
 128 int OlySocket::connect(const char* path, const size_t pathSize) {
 129   int fd = socket_cloexec(PF_UNIX, SOCK_STREAM, 0);
 130   if (fd < 0) {
 131     return -1;
 132   }
 133
 134   // Create sockaddr_in structure, ensuring non-populated fields are zero
 135   struct sockaddr_un sockaddr;
 136   memset((void*)&sockaddr, 0, sizeof(sockaddr));
 137   sockaddr.sun_family = AF_UNIX;
 138   memcpy(sockaddr.sun_path, path, MIN(pathSize, sizeof(sockaddr.sun_path)));
 139   sockaddr.sun_path[sizeof(sockaddr.sun_path) - 1] = '\0';
 140
 141   if (::connect(fd, (const struct sockaddr*)&sockaddr, sizeof(sockaddr)) < 0) {
 142     close(fd);
 143     return -1;
 144   }
 145
 146   return fd;
 147 }
 148
 149 #endif
 150
 151 OlySocket::~OlySocket() {
 152   if (mSocketID > 0) {
 153     CLOSE_SOCKET(mSocketID);
 154   }
 155 }
 156
 157 OlyServerSocket::~OlyServerSocket() {
 158   if (mFDServer > 0) {
 159     CLOSE_SOCKET(mFDServer);
 160   }
 161 }
 162
 163 void OlySocket::shutdownConnection() {
 164   // Shutdown is primarily used to unblock other threads that are blocking on send/receive functions
 165   shutdown(mSocketID, SHUTDOWN_RX_TX);
 166 }
 167
 168 void OlySocket::closeSocket() {
 169   // Used for closing an accepted socket but keeping the server socket active
 170   if (mSocketID > 0) {
 171     CLOSE_SOCKET(mSocketID);
 172     mSocketID = -1;
 173   }
 174 }
 175
 176 void OlyServerSocket::closeServerSocket() {
 177   if (CLOSE_SOCKET(mFDServer) != 0) {
 178     logg->logError(__FILE__, __LINE__, "Failed to close server socket.");
 179     handleException();
 180   }
 181   mFDServer = 0;
 182 }
 183
 184 void OlyServerSocket::createServerSocket(int port) {
 185   int family = AF_INET6;
 186
 187   // Create socket
 188   mFDServer = socket_cloexec(PF_INET6, SOCK_STREAM, IPPROTO_TCP);
 189   if (mFDServer < 0) {
 190     family = AF_INET;
 191     mFDServer = socket_cloexec(PF_INET, SOCK_STREAM, IPPROTO_TCP);
 192     if (mFDServer < 0) {
 193       logg->logError(__FILE__, __LINE__, "Error creating server socket");
 194       handleException();
 195     }
 196   }
 197
 198   // Enable address reuse, another solution would be to create the server socket once and only close it when the object exits
 199   int on = 1;
 200   if (setsockopt(mFDServer, SOL_SOCKET, SO_REUSEADDR, (const char*)&on, sizeof(on)) != 0) {
 201     logg->logError(__FILE__, __LINE__, "Setting server socket options failed");
 202     handleException();
 203   }
 204
 205   // Create sockaddr_in structure, ensuring non-populated fields are zero
 206   struct sockaddr_in6 sockaddr;
 207   memset((void*)&sockaddr, 0, sizeof(sockaddr));
 208   sockaddr.sin6_family = family;
 209   sockaddr.sin6_port = htons(port);
 210   sockaddr.sin6_addr = in6addr_any;
 211
 212   // Bind the socket to an address
 213   if (bind(mFDServer, (const struct sockaddr*)&sockaddr, sizeof(sockaddr)) < 0) {
 214     logg->logError(__FILE__, __LINE__, "Binding of server socket failed.\nIs an instance already running?");
 215     handleException();
 216   }
 217
 218   // Listen for connections on this socket
 219   if (listen(mFDServer, 1) < 0) {
 220     logg->logError(__FILE__, __LINE__, "Listening of server socket failed");
 221     handleException();
 222   }
 223 }
 224
 225 // mSocketID is always set to the most recently accepted connection
 226 // The user of this class should maintain the different socket connections, e.g. by forking the process
 227 int OlyServerSocket::acceptConnection() {
 228   int socketID;
 229   if (mFDServer <= 0) {
 230     logg->logError(__FILE__, __LINE__, "Attempting multiple connections on a single connection server socket or attempting to accept on a client socket");
 231     handleException();
 232   }
 233
 234   // Accept a connection, note that this call blocks until a client connects
 235   socketID = accept_cloexec(mFDServer, NULL, NULL);
 236   if (socketID < 0) {
 237     logg->logError(__FILE__, __LINE__, "Socket acceptance failed");
 238     handleException();
 239   }
 240   return socketID;
 241 }
 242
 243 void OlySocket::send(const char* buffer, int size) {
 244   if (size <= 0 || buffer == NULL) {
 245     return;
 246   }
 247
 248   while (size > 0) {
 249     int n = ::send(mSocketID, buffer, size, 0);
 250     if (n < 0) {
 251       logg->logError(__FILE__, __LINE__, "Socket send error");
 252       handleException();
 253     }
 254     size -= n;
 255     buffer += n;
 256   }
 257 }
 258
 259 // Returns the number of bytes received
 260 int OlySocket::receive(char* buffer, int size) {
 261   if (size <= 0 || buffer == NULL) {
 262     return 0;
 263   }
 264
 265   int bytes = recv(mSocketID, buffer, size, 0);
 266   if (bytes < 0) {
 267     logg->logError(__FILE__, __LINE__, "Socket receive error");
 268     handleException();
 269   } else if (bytes == 0) {
 270     logg->logMessage("Socket disconnected");
 271     return -1;
 272   }
 273   return bytes;
 274 }
 275
 276 // Receive exactly size bytes of data. Note, this function will block until all bytes are received
 277 int OlySocket::receiveNBytes(char* buffer, int size) {
 278   int bytes = 0;
 279   while (size > 0 && buffer != NULL) {
 280     bytes = recv(mSocketID, buffer, size, 0);
 281     if (bytes < 0) {
 282       logg->logError(__FILE__, __LINE__, "Socket receive error");
 283       handleException();
 284     } else if (bytes == 0) {
 285       logg->logMessage("Socket disconnected");
 286       return -1;
 287     }
 288     buffer += bytes;
 289     size -= bytes;
 290   }
 291   return bytes;
 292 }
 293
 294 // Receive data until a carriage return, line feed, or null is encountered, or the buffer fills
 295 int OlySocket::receiveString(char* buffer, int size) {
 296   int bytes_received = 0;
 297   bool found = false;
 298
 299   if (buffer == 0) {
 300     return 0;
 301   }
 302
 303   while (!found && bytes_received < size) {
 304     // Receive a single character
 305     int bytes = recv(mSocketID, &buffer[bytes_received], 1, 0);
 306     if (bytes < 0) {
 307       logg->logError(__FILE__, __LINE__, "Socket receive error");
 308       handleException();
 309     } else if (bytes == 0) {
 310       logg->logMessage("Socket disconnected");
 311       return -1;
 312     }
 313
 314     // Replace carriage returns and line feeds with zero
 315     if (buffer[bytes_received] == '\n' || buffer[bytes_received] == '\r' || buffer[bytes_received] == '\0') {
 316       buffer[bytes_received] = '\0';
 317       found = true;
 318     }
 319
 320     bytes_received++;
 321   }
 322
 323   return bytes_received;
 324 }