daemon/OlySocket.cpp

   1 /**
   2  * Copyright (C) ARM Limited 2010-2011. 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 extern void handleException();
  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         fdServer = 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(fdServer) != 0) {
  74                 logg->logError(__FILE__, __LINE__, "Failed to close server socket.");
  75                 handleException();
  76         }
  77         fdServer = 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                 mSocketID = socket(res->ai_family, res->ai_socktype, res->ai_protocol);
 101                 if (mSocketID < 0) {
 102                         continue;
 103                 }
 104                 if (connect(mSocketID, res->ai_addr, res->ai_addrlen) < 0) {
 105                         close(mSocketID);
 106                         mSocketID = -1;
 107                 }
 108                 if (mSocketID > 0) break;
 109         }
 110         freeaddrinfo(res0);
 111         if (mSocketID <= 0) {
 112                 logg->logError(__FILE__, __LINE__, "Could not connect to client socket. Ensure ARM Streamline is running.");
 113                 handleException();
 114         }
 115 #endif
 116 }
 117
 118 void OlySocket::createSingleServerConnection(int port) {
 119         createServerSocket(port);
 120
 121         mSocketID = acceptConnection();
 122         closeServerSocket();
 123 }
 124
 125 void OlySocket::createServerSocket(int port) {
 126         // Create socket
 127         fdServer = socket(PF_INET, SOCK_STREAM, IPPROTO_TCP);
 128         if (fdServer < 0) {
 129                 logg->logError(__FILE__, __LINE__, "Error creating server socket");
 130                 handleException();
 131         }
 132
 133         // Enable address reuse, another solution would be to create the server socket once and only close it when the object exits
 134         int on = 1;
 135         if (setsockopt(fdServer, SOL_SOCKET, SO_REUSEADDR, (const char *)&on, sizeof(on)) != 0) {
 136                 logg->logError(__FILE__, __LINE__, "Setting server socket options failed");
 137                 handleException();
 138         }
 139
 140         // Create sockaddr_in structure, ensuring non-populated fields are zero
 141         struct sockaddr_in      sockaddr;
 142         memset((void*)&sockaddr, 0, sizeof(struct sockaddr_in));
 143         sockaddr.sin_family = AF_INET;
 144         sockaddr.sin_port   = htons(port);
 145         sockaddr.sin_addr.s_addr = INADDR_ANY;
 146
 147         // Bind the socket to an address
 148         if (bind(fdServer, (const struct sockaddr*)&sockaddr, sizeof(sockaddr)) < 0) {
 149                 logg->logError(__FILE__, __LINE__, "Binding of server socket failed.\nIs an instance already running?");
 150                 handleException();
 151         }
 152
 153         // Listen for connections on this socket
 154         if (listen(fdServer, 1) < 0) {
 155                 logg->logError(__FILE__, __LINE__, "Listening of server socket failed");
 156                 handleException();
 157         }
 158 }
 159
 160 // mSocketID is always set to the most recently accepted connection
 161 // The user of this class should maintain the different socket connections, e.g. by forking the process
 162 int OlySocket::acceptConnection() {
 163         if (fdServer <= 0) {
 164                 logg->logError(__FILE__, __LINE__, "Attempting multiple connections on a single connection server socket or attempting to accept on a client socket");
 165                 handleException();
 166         }
 167
 168         // Accept a connection, note that this call blocks until a client connects
 169         mSocketID = accept(fdServer, NULL, NULL);
 170         if (mSocketID < 0) {
 171                 logg->logError(__FILE__, __LINE__, "Socket acceptance failed");
 172                 handleException();
 173         }
 174         return mSocketID;
 175 }
 176
 177 void OlySocket::send(char* buffer, int size) {
 178         if (size <= 0 || buffer == NULL) {
 179                 return;
 180         }
 181
 182         while (size > 0) {
 183                 int n = ::send(mSocketID, buffer, size, 0);
 184                 if (n < 0) {
 185                         logg->logError(__FILE__, __LINE__, "Socket send error");
 186                         handleException();
 187                 }
 188                 size -= n;
 189                 buffer += n;
 190         }
 191 }
 192
 193 // Returns the number of bytes received
 194 int OlySocket::receive(char* buffer, int size) {
 195         if (size <= 0 || buffer == NULL) {
 196                 return 0;
 197         }
 198
 199         int bytes = recv(mSocketID, buffer, size, 0);
 200         if (bytes < 0) {
 201                 logg->logError(__FILE__, __LINE__, "Socket receive error");
 202                 handleException();
 203         } else if (bytes == 0) {
 204                 logg->logMessage("Socket disconnected");
 205                 return -1;
 206         }
 207         return bytes;
 208 }
 209
 210 // Receive exactly size bytes of data. Note, this function will block until all bytes are received
 211 int OlySocket::receiveNBytes(char* buffer, int size) {
 212         int bytes = 0;
 213         while (size > 0 && buffer != NULL) {
 214                 bytes = recv(mSocketID, buffer, size, 0);
 215                 if (bytes < 0) {
 216                         logg->logError(__FILE__, __LINE__, "Socket receive error");
 217                         handleException();
 218                 } else if (bytes == 0) {
 219                         logg->logMessage("Socket disconnected");
 220                         return -1;
 221                 }
 222                 buffer += bytes;
 223                 size -= bytes;
 224         }
 225         return bytes;
 226 }
 227
 228 // Receive data until a carriage return, line feed, or null is encountered, or the buffer fills
 229 int OlySocket::receiveString(char* buffer, int size) {
 230         int bytes_received = 0;
 231         bool found = false;
 232
 233         if (buffer == 0)
 234                 return 0;
 235
 236         while (!found && bytes_received < size) {
 237                 // Receive a single character
 238                 int bytes = recv(mSocketID, &buffer[bytes_received], 1, 0);
 239                 if (bytes < 0) {
 240                         logg->logError(__FILE__, __LINE__, "Socket receive error");
 241                         handleException();
 242                 } else if (bytes == 0) {
 243                         logg->logMessage("Socket disconnected");
 244                         return -1;
 245                 }
 246
 247                 // Replace carriage returns and line feeds with zero
 248                 if (buffer[bytes_received] == '\n' || buffer[bytes_received] == '\r' || buffer[bytes_received] == '\0') {
 249                         buffer[bytes_received] = '\0';
 250                         found = true;
 251                 }
 252
 253                 bytes_received++;
 254         }
 255
 256         return bytes_received;
 257 }