daemon/OlySocket.cpp

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