daemon/OlySocket.cpp

   1 /**
   2  * Copyright (C) ARM Limited 2010-2013. 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 #ifdef WIN32
  13 #include <Winsock2.h>
  14 #else
  15 #include <netinet/in.h>
  16 #include <sys/socket.h>
  17 #include <unistd.h>
  18 #include <netdb.h>
  19 #endif
  20
  21 #include "Logging.h"
  22
  23 #ifdef WIN32
  24 #define CLOSE_SOCKET(x) closesocket(x)
  25 #define SHUTDOWN_RX_TX SD_BOTH
  26 #define snprintf       _snprintf
  27 #else
  28 #define CLOSE_SOCKET(x) close(x)
  29 #define SHUTDOWN_RX_TX SHUT_RDWR
  30 #endif
  31
  32 OlySocket::OlySocket(int port, bool multiple) {
  33 #ifdef WIN32
  34   WSADATA wsaData;
  35   if (WSAStartup(0x0202, &wsaData) != 0) {
  36     logg->logError(__FILE__, __LINE__, "Windows socket initialization failed");
  37     handleException();
  38   }
  39 #endif
  40
  41   if (multiple) {
  42     createServerSocket(port);
  43   } else {
  44     createSingleServerConnection(port);
  45   }
  46 }
  47
  48 OlySocket::OlySocket(int port, char* host) {
  49   mFDServer = 0;
  50   createClientSocket(host, port);
  51 }
  52
  53 OlySocket::~OlySocket() {
  54   if (mSocketID > 0) {
  55     CLOSE_SOCKET(mSocketID);
  56   }
  57 }
  58
  59 void OlySocket::shutdownConnection() {
  60   // Shutdown is primarily used to unblock other threads that are blocking on send/receive functions
  61   shutdown(mSocketID, SHUTDOWN_RX_TX);
  62 }
  63
  64 void OlySocket::closeSocket() {
  65   // Used for closing an accepted socket but keeping the server socket active
  66   if (mSocketID > 0) {
  67     CLOSE_SOCKET(mSocketID);
  68     mSocketID = -1;
  69   }
  70 }
  71
  72 void OlySocket::closeServerSocket() {
  73   if (CLOSE_SOCKET(mFDServer) != 0) {
  74     logg->logError(__FILE__, __LINE__, "Failed to close server socket.");
  75     handleException();
  76   }
  77   mFDServer = 0;
  78 }
  79
  80 void OlySocket::createClientSocket(char* hostname, int portno) {
  81 #ifdef WIN32
  82   // TODO: Implement for Windows
  83 #else
  84   char buf[32];
  85   struct addrinfo hints, *res, *res0;
  86
  87   snprintf(buf, sizeof(buf), "%d", portno);
  88   mSocketID = -1;
  89   memset((void*)&hints, 0, sizeof(hints));
  90   hints.ai_family = PF_UNSPEC;
  91   hints.ai_socktype = SOCK_STREAM;
  92
  93   if (getaddrinfo(hostname, buf, &hints, &res0)) {
  94     logg->logError(__FILE__, __LINE__, "Client socket failed to get address info for %s", hostname);
  95     handleException();
  96   }
  97   for (res=res0; res!=NULL; res = res->ai_next) {
  98     if ( res->ai_family != PF_INET || res->ai_socktype != SOCK_STREAM ) {
  99       continue;
 100     }
 101     mSocketID = socket(res->ai_family, res->ai_socktype, res->ai_protocol);
 102     if (mSocketID < 0) {
 103       continue;
 104     }
 105     if (connect(mSocketID, res->ai_addr, res->ai_addrlen) < 0) {
 106       close(mSocketID);
 107       mSocketID = -1;
 108     }
 109     if (mSocketID > 0) {
 110       break;
 111     }
 112   }
 113   freeaddrinfo(res0);
 114   if (mSocketID <= 0) {
 115     logg->logError(__FILE__, __LINE__, "Could not connect to client socket. Ensure ARM Streamline is running.");
 116     handleException();
 117   }
 118 #endif
 119 }
 120
 121 void OlySocket::createSingleServerConnection(int port) {
 122   createServerSocket(port);
 123
 124   mSocketID = acceptConnection();
 125   closeServerSocket();
 126 }
 127
 128 void OlySocket::createServerSocket(int port) {
 129   // Create socket
 130   mFDServer = socket(PF_INET, SOCK_STREAM, IPPROTO_TCP);
 131   if (mFDServer < 0) {
 132     logg->logError(__FILE__, __LINE__, "Error creating server socket");
 133     handleException();
 134   }
 135
 136   // Enable address reuse, another solution would be to create the server socket once and only close it when the object exits
 137   int on = 1;
 138   if (setsockopt(mFDServer, SOL_SOCKET, SO_REUSEADDR, (const char*)&on, sizeof(on)) != 0) {
 139     logg->logError(__FILE__, __LINE__, "Setting server socket options failed");
 140     handleException();
 141   }
 142
 143   // Create sockaddr_in structure, ensuring non-populated fields are zero
 144   struct sockaddr_in sockaddr;
 145   memset((void*)&sockaddr, 0, sizeof(struct sockaddr_in));
 146   sockaddr.sin_family = AF_INET;
 147   sockaddr.sin_port   = htons(port);
 148   sockaddr.sin_addr.s_addr = INADDR_ANY;
 149
 150   // Bind the socket to an address
 151   if (bind(mFDServer, (const struct sockaddr*)&sockaddr, sizeof(sockaddr)) < 0) {
 152     logg->logError(__FILE__, __LINE__, "Binding of server socket failed.\nIs an instance already running?");
 153     handleException();
 154   }
 155
 156   // Listen for connections on this socket
 157   if (listen(mFDServer, 1) < 0) {
 158     logg->logError(__FILE__, __LINE__, "Listening of server socket failed");
 159     handleException();
 160   }
 161 }
 162
 163 // mSocketID is always set to the most recently accepted connection
 164 // The user of this class should maintain the different socket connections, e.g. by forking the process
 165 int OlySocket::acceptConnection() {
 166   if (mFDServer <= 0) {
 167     logg->logError(__FILE__, __LINE__, "Attempting multiple connections on a single connection server socket or attempting to accept on a client socket");
 168     handleException();
 169   }
 170
 171   // Accept a connection, note that this call blocks until a client connects
 172   mSocketID = accept(mFDServer, NULL, NULL);
 173   if (mSocketID < 0) {
 174     logg->logError(__FILE__, __LINE__, "Socket acceptance failed");
 175     handleException();
 176   }
 177   return mSocketID;
 178 }
 179
 180 void OlySocket::send(char* buffer, int size) {
 181   if (size <= 0 || buffer == NULL) {
 182     return;
 183   }
 184
 185   while (size > 0) {
 186     int n = ::send(mSocketID, buffer, size, 0);
 187     if (n < 0) {
 188       logg->logError(__FILE__, __LINE__, "Socket send error");
 189       handleException();
 190     }
 191     size -= n;
 192     buffer += n;
 193   }
 194 }
 195
 196 // Returns the number of bytes received
 197 int OlySocket::receive(char* buffer, int size) {
 198   if (size <= 0 || buffer == NULL) {
 199     return 0;
 200   }
 201
 202   int bytes = recv(mSocketID, buffer, size, 0);
 203   if (bytes < 0) {
 204     logg->logError(__FILE__, __LINE__, "Socket receive error");
 205     handleException();
 206   } else if (bytes == 0) {
 207     logg->logMessage("Socket disconnected");
 208     return -1;
 209   }
 210   return bytes;
 211 }
 212
 213 // Receive exactly size bytes of data. Note, this function will block until all bytes are received
 214 int OlySocket::receiveNBytes(char* buffer, int size) {
 215   int bytes = 0;
 216   while (size > 0 && buffer != NULL) {
 217     bytes = recv(mSocketID, buffer, size, 0);
 218     if (bytes < 0) {
 219       logg->logError(__FILE__, __LINE__, "Socket receive error");
 220       handleException();
 221     } else if (bytes == 0) {
 222       logg->logMessage("Socket disconnected");
 223       return -1;
 224     }
 225     buffer += bytes;
 226     size -= bytes;
 227   }
 228   return bytes;
 229 }
 230
 231 // Receive data until a carriage return, line feed, or null is encountered, or the buffer fills
 232 int OlySocket::receiveString(char* buffer, int size) {
 233   int bytes_received = 0;
 234   bool found = false;
 235
 236   if (buffer == 0) {
 237     return 0;
 238   }
 239
 240   while (!found && bytes_received < size) {
 241     // Receive a single character
 242     int bytes = recv(mSocketID, &buffer[bytes_received], 1, 0);
 243     if (bytes < 0) {
 244       logg->logError(__FILE__, __LINE__, "Socket receive error");
 245       handleException();
 246     } else if (bytes == 0) {
 247       logg->logMessage("Socket disconnected");
 248       return -1;
 249     }
 250
 251     // Replace carriage returns and line feeds with zero
 252     if (buffer[bytes_received] == '\n' || buffer[bytes_received] == '\r' || buffer[bytes_received] == '\0') {
 253       buffer[bytes_received] = '\0';
 254       found = true;
 255     }
 256
 257     bytes_received++;
 258   }
 259
 260   return bytes_received;
 261 }