irc/ArchBlackmann/SockUtils.cpp

   1 // SockUtils.cpp - Some basic socket utility functions.
   2 // (C) 2002-2004 Royce Mitchell III
   3 // This file is under the BSD & LGPL licenses
   4
   5 #include <stdio.h>
   6 #include "SockUtils.h"
   7 #ifdef WIN32
   8 #  ifndef SD_SEND // defined in winsock2.h, but not winsock.h
   9 #    define SD_SEND 1
  10 #  endif
  11 #  define snprintf _snprintf
  12 #  ifdef _MSC_VER
  13 #    pragma comment ( lib, "ws2_32.lib" )
  14 #  endif//_MSC_VER
  15 #elif defined(UNIX)
  16 #  include <errno.h>
  17 #  include "string.h" // memset
  18 #  include <netdb.h> // hostent
  19 #  include <arpa/inet.h> //inet_addr
  20 #  include <sys/time.h>
  21 #  define SD_SEND SHUT_WR //bah thou shalt name thy defines the same
  22 #else
  23 #  error unrecognized target
  24 #endif
  25 //// Constants /////////////////////////////////////////////////////////
  26 const int kBufferSize = 1024;
  27 // creates broadcast address
  28 SockAddrIn::SockAddrIn()
  29 {
  30         memset ( this, 0, sizeof(sockaddr_in) );
  31         sin_family = AF_INET;
  32 }
  33 SockAddrIn::SockAddrIn ( const char* szAddr, u_short iPort )
  34 {
  35         memset ( this, 0, sizeof(sockaddr_in) );
  36         sin_family = AF_INET;
  37         sin_addr.s_addr = suLookupAddress(szAddr);
  38         sin_port = htons(iPort);
  39 }
  40 SockAddrIn::SockAddrIn ( in_addr_t iAddr, u_short iPort )
  41 {
  42         memset ( this, 0, sizeof(sockaddr_in) );
  43         sin_family = AF_INET;
  44         sin_addr.s_addr = iAddr;
  45         sin_port = htons(iPort);
  46 }
  47 bool suStartup()
  48 {
  49 #ifdef WIN32
  50         WSADATA wsaData;
  51         if ( WSAStartup ( MAKEWORD(2,0), &wsaData ) )
  52                 return false;
  53         if ( wsaData.wVersion != MAKEWORD(2,0) )
  54         {
  55                 WSACleanup();
  56                 return false;
  57         }
  58         return true;
  59 #elif defined(UNIX)
  60         // nothing special required here
  61         return true;
  62 #else
  63 #  error unrecognized target
  64 #endif
  65 }
  66 //// suTcpSocket ////////////////////////////////////////////////
  67 // Creates a TCP socket.
  68 SOCKET suTcpSocket()
  69 {
  70         SOCKET so = socket ( AF_INET, SOCK_STREAM, 0 );
  71 #if defined(_DEBUG) && defined(WIN32)
  72         if ( so == INVALID_SOCKET && WSANOTINITIALISED == WSAGetLastError() )
  73                 MessageBox ( NULL, "You forgot to call suStartup()!", "SockUtils", MB_OK|MB_ICONEXCLAMATION );
  74 #endif
  75         return so;
  76 }
  77 //// suUdpSocket ////////////////////////////////////////////////
  78 // Creates a UDP socket. Compensates for new "functionality" introduced
  79 // in Win2K with regards to select() calls
  80 // MS Transport Provider IOCTL to control
  81 // reporting PORT_UNREACHABLE messages
  82 // on UDP sockets via recv/WSARecv/etc.
  83 // Path TRUE in input buffer to enable (default if supported),
  84 // FALSE to disable.
  85 #ifndef SIO_UDP_CONNRESET
  86 #define SIO_UDP_CONNRESET _WSAIOW(IOC_VENDOR,12)
  87 #endif//SIO_UDP_CONNRESET
  88 SOCKET suUdpSocket()
  89 {
  90         SOCKET so = socket ( AF_INET, SOCK_DGRAM, 0 );
  91 #if defined(_DEBUG) && defined(WIN32)
  92         if ( so == INVALID_SOCKET && WSANOTINITIALISED == WSAGetLastError() )
  93                 MessageBox ( NULL, "You forgot to call suStartup()!", "SockUtils", MB_OK|MB_ICONEXCLAMATION );
  94 #endif
  95 #ifdef WIN32
  96         // for Windows 2000, disable new behavior...
  97         // see: http://www-pc.uni-regensburg.de/systemsw/W2KPRO/UPDATE/POSTSP1/Q263823.htm
  98         // this code is innocuous on other win32 platforms
  99         DWORD dwBytesReturned = 0;
 100         BOOL bNewBehavior = FALSE;
 101         // disable  new Win2K behavior using
 102         // IOCTL: SIO_UDP_CONNRESET
 103         // we don't care about return value :)
 104         WSAIoctl(so, SIO_UDP_CONNRESET,
 105                 &bNewBehavior, sizeof(bNewBehavior),
 106                 NULL, 0, &dwBytesReturned,
 107                 NULL, NULL);
 108 #endif
 109         return so;
 110 }
 111 //// suShutdownConnection ////////////////////////////////////////////////
 112 // Gracefully shuts the connection sd down.  Returns true if it was able
 113 // to shut it down nicely, false if we had to "slam" it shut.
 114 // (either way, the socket does get closed)
 115 bool suShutdownConnection(SOCKET sd)
 116 {
 117         if ( sd == INVALID_SOCKET )
 118                 return true;
 119         // Disallow any further data sends.  This will tell the other side
 120         // that we want to go away now.  If we skip this step, we don't
 121         // shut the connection down nicely.
 122         if (shutdown(sd, SD_SEND) == SOCKET_ERROR)
 123         {
 124                 closesocket(sd);
 125                 return false;
 126         }
 127         // Receive any extra data still sitting on the socket.  After all
 128         // data is received, this call will block until the remote host
 129         // acknowledges the TCP control packet sent by the shutdown above.
 130         // Then we'll get a 0 back from recv, signalling that the remote
 131         // host has closed its side of the connection.
 132         char acReadBuffer[kBufferSize];
 133         for ( ;; )
 134         {
 135                 int nNewBytes = recv(sd, acReadBuffer, kBufferSize, 0);
 136                 if (nNewBytes == SOCKET_ERROR)
 137                 {
 138                         closesocket(sd);
 139                         return false;
 140                 }
 141                 else if (nNewBytes != 0)
 142                 {
 143                         // FYI, received (nNewBytes) unexpected bytes during shutdown.
 144                 }
 145                 else
 146                 {
 147                         // Okay, we're done!
 148                         break;
 149                 }
 150         }
 151         // Close the socket.
 152         if (closesocket(sd) == SOCKET_ERROR)
 153         {
 154                 return false;
 155         }
 156         return true;
 157 }
 158 //// suLookupAddress ////////////////////////////////////////////////
 159 // Basically converts a name address to an ip address
 160 in_addr_t suLookupAddress ( const char* pcHost )
 161 {
 162         in_addr_t nRemoteAddr = inet_addr(pcHost);
 163         if ( nRemoteAddr == INADDR_NONE )
 164         {
 165                 // pcHost isn't a dotted IP, so resolve it through DNS
 166                 hostent* pHE = gethostbyname(pcHost);
 167                 if ( pHE == 0 )
 168                 {
 169 #if defined(_DEBUG) && defined(WIN32)
 170                         if ( WSANOTINITIALISED == WSAGetLastError() )
 171                                 MessageBox ( NULL, "You forgot to call suStartup()!", "SockUtils", MB_OK|MB_ICONEXCLAMATION );
 172 #endif
 173                         return INADDR_NONE;
 174                 }
 175                 nRemoteAddr = *((in_addr_t*)pHE->h_addr_list[0]);
 176         }
 177         return nRemoteAddr;
 178 }
 179 bool suConnect ( SOCKET so, in_addr_t iAddress, u_short iPort )
 180 {
 181         SockAddrIn sinRemote ( iAddress, iPort );
 182         if ( SOCKET_ERROR == connect(so,sinRemote,sizeof(sinRemote)) )
 183         {
 184 #if defined(_DEBUG) && defined(WIN32)
 185                 if ( WSANOTINITIALISED == WSAGetLastError() )
 186                         MessageBox ( NULL, "You forgot to call suStartup()!", "SockUtils", MB_OK|MB_ICONEXCLAMATION );
 187 #endif
 188                 return false;
 189         }
 190         return true;
 191 }
 192 bool suConnect ( SOCKET so, const char* szAddress, u_short iPort )
 193 {
 194         return suConnect ( so, suLookupAddress(szAddress), iPort );
 195 }
 196 //// suEstablishConnection ////////////////////////////////////////////////
 197 // creates a socket of the specified type, connects to the ip address/port
 198 // requested, and returns the SOCKET created
 199 SOCKET suEstablishConnection ( in_addr_t iAddress, u_short iPort, int type )
 200 {
 201         // Create a socket
 202         if ( type != SOCK_STREAM && type != SOCK_DGRAM )
 203                 return INVALID_SOCKET;
 204         SOCKET so = socket(AF_INET, type, 0);
 205         if ( so == INVALID_SOCKET )
 206                 return so;
 207         if ( !suConnect(so, iAddress, iPort) )
 208         {
 209                 closesocket(so);
 210                 return INVALID_SOCKET;
 211         }
 212         return so;
 213 }
 214 //// suEstablishConnection ////////////////////////////////////////////////
 215 // creates a socket of the specified type, connects to the address/port
 216 // requested, and returns the SOCKET created
 217 SOCKET suEstablishConnection ( const char* szAddress, u_short iPort, int type )
 218 {
 219         return suEstablishConnection ( suLookupAddress ( szAddress ), iPort, type );
 220 }
 221 //// suBroadcast ////////////////////////////////////////////////
 222 // takes a previously created broadcast-enabled UDP socket, and broadcasts
 223 // a message on the local network
 224 bool suBroadcast ( SOCKET so, u_short port, const char* buf, int len /* = -1 */ )
 225 {
 226         if ( len == -1 )
 227                 len = (int)strlen(buf);
 228 #if 1
 229         SockAddrIn to ( INADDR_BROADCAST, port );
 230 #else // some strange MS OS's don't broadcast to localhost...
 231         SockAddrIn to ( "127.0.0.1", port );
 232         if ( SOCKET_ERROR == sendto ( so, buf, len, 0, to, sizeof(to) ) )
 233                 return false;
 234         to.sin_addr.s_addr = INADDR_BROADCAST;
 235 #endif
 236         if ( SOCKET_ERROR == sendto ( so, buf, len, 0, to, sizeof(to) ) )
 237                 return false;
 238         return true;
 239 }
 240 //// suRecv ////////////////////////////////////////////////
 241 // retrieves data sent to our TCP socket. If no data, waits for
 242 // a period of timeout ms.
 243 // returns bytes received
 244 // -1 == SOCKET_ERROR
 245 // -2 == timed out waiting for data
 246 int suRecv ( SOCKET so, char* buf, int buflen, int timeout )
 247 {
 248         struct timeval to;
 249         fd_set rread;
 250         int res;
 251         FD_ZERO(&rread);   // clear the fd_set
 252         FD_SET(so,&rread); // indicate which socket(s) we want to check
 253         memset((char *)&to,0,sizeof(to)); // clear the timeval struct
 254         to.tv_sec = timeout;  // timeout select after (timeout) seconds
 255         // select returns > 0 if there is an event on the socket
 256         res = select((int)so+1, &rread, (fd_set *)0, (fd_set *)0, &to );
 257         if (res < 0)
 258                 return -1; // socket error
 259         // there was an event on the socket
 260         if ( (res>0) && (FD_ISSET(so,&rread)) )
 261                 return recv ( so, buf, buflen, 0 );
 262         return -2;
 263 }
 264 //// suRecvFrom ////////////////////////////////////////////////
 265 // retrieves data sent to our UDP socket. If no data, waits for
 266 // a period of timeout ms.
 267 // returns bytes received
 268 // returns bytes received
 269 // -1 == SOCKET_ERROR
 270 // -2 == timed out waiting for data
 271 int suRecvFrom ( SOCKET so, char* buf, int buflen, int timeout, sockaddr_in* from, socklen_t* fromlen )
 272 {
 273         struct timeval to;
 274         fd_set rread;
 275         int res;
 276         FD_ZERO(&rread);   // clear the fd_set
 277         FD_SET(so,&rread); // indicate which socket(s) we want to check
 278         memset((char *)&to,0,sizeof(to)); // clear the timeval struct
 279         to.tv_sec = timeout;  // timeout select after (timeout) seconds
 280         // select returns > 0 if there is an event on the socket
 281         res = select((int)so+1, &rread, (fd_set *)0, (fd_set *)0, &to );
 282         if (res < 0)
 283                 return -1; // socket error
 284         // there was an event on the socket
 285         if ( (res>0) && (FD_ISSET(so,&rread)) )
 286                 return recvfrom ( so, buf, buflen, 0, (sockaddr*)from, fromlen );
 287         return -2; // timeout
 288 }
 289 //// suBind ////////////////////////////////////////////////
 290 // binds a UDP socket to an interface & port to receive
 291 // data on that port
 292 bool suBind ( SOCKET so, in_addr_t iInterfaceAddress, u_short iListenPort, bool bReuseAddr /* = false */ )
 293 {
 294         SockAddrIn sinInterface ( iInterfaceAddress, iListenPort );
 295         if ( bReuseAddr )
 296         {
 297                 int optval = -1; // true
 298                 if ( SOCKET_ERROR == setsockopt ( so, SOL_SOCKET, SO_REUSEADDR, (const char*)&optval, sizeof(optval) ) )
 299                 {
 300 #if defined(_DEBUG) && defined(WIN32)
 301                         if ( WSANOTINITIALISED == WSAGetLastError() )
 302                                 MessageBox ( NULL, "You forgot to call suStartup()!", "SockUtils", MB_OK|MB_ICONEXCLAMATION );
 303 #endif
 304                         return false;
 305                 }
 306         }
 307         if ( SOCKET_ERROR == bind(so, sinInterface, sizeof(sinInterface)) )
 308         {
 309                 int err = SUERRNO;
 310                 if ( err != EADDRINUSE )
 311                         return false;
 312 #if defined(_DEBUG) && defined(WIN32)
 313                 if ( WSANOTINITIALISED == WSAGetLastError() )
 314                         MessageBox ( NULL, "You forgot to call suStartup()!", "SockUtils", MB_OK|MB_ICONEXCLAMATION );
 315 #endif
 316         }
 317         return true;
 318 }
 319 //// suBind ////////////////////////////////////////////////
 320 // binds a UDP socket to an interface & port to receive
 321 // data on that port
 322 bool suBind ( SOCKET so, const char* szInterfaceAddress, u_short iListenPort, bool bReuseAddr /* = false */ )
 323 {
 324         in_addr_t iInterfaceAddr = inet_addr(szInterfaceAddress);
 325         if ( iInterfaceAddr == INADDR_NONE )
 326                 return false;
 327         return suBind ( so, iInterfaceAddr, iListenPort, bReuseAddr );
 328 }
 329 //// suEnableBroadcast ////////////////////////////////////////////////
 330 // in order to send broadcast messages on a UDP socket, this function
 331 // must be called first
 332 bool suEnableBroadcast ( SOCKET so, bool bEnable /* = true */ )
 333 {
 334         int optval = bEnable ? -1 : 0;
 335         if ( SOCKET_ERROR == setsockopt ( so, SOL_SOCKET, SO_BROADCAST, (const char*)&optval, sizeof(optval) ) )
 336                 return false;
 337         return true;
 338 }
 339 //// suErrDesc ////////////////////////////////////////////////
 340 // returns text description of error code
 341 const char* suErrDesc ( int err )
 342 {
 343         static char errbuf[256];
 344 #ifdef WIN32
 345         switch ( err )
 346         {
 347 #define X(E) case E: return #E;
 348                 X(WSAEINTR)     X(WSAEBADF)
 349                 X(WSAEACCES)    X(WSAEFAULT)
 350                 X(WSAEINVAL)    X(WSAEMFILE)
 351                 X(WSAEWOULDBLOCK)       X(WSAEINPROGRESS)
 352                 X(WSAEALREADY)  X(WSAENOTSOCK)
 353                 X(WSAEDESTADDRREQ)      X(WSAEMSGSIZE)
 354                 X(WSAEPROTOTYPE)        X(WSAENOPROTOOPT)
 355                 X(WSAEPROTONOSUPPORT)   X(WSAESOCKTNOSUPPORT)
 356                 X(WSAEOPNOTSUPP)        X(WSAEPFNOSUPPORT)
 357                 X(WSAEAFNOSUPPORT)      X(WSAEADDRINUSE)
 358                 X(WSAEADDRNOTAVAIL)     X(WSAENETDOWN)
 359                 X(WSAENETUNREACH)       X(WSAENETRESET)
 360                 X(WSAECONNABORTED)      X(WSAECONNRESET)
 361                 X(WSAENOBUFS)   X(WSAEISCONN)
 362                 X(WSAENOTCONN)  X(WSAESHUTDOWN)
 363                 X(WSAETOOMANYREFS)      X(WSAETIMEDOUT)
 364                 X(WSAECONNREFUSED)      X(WSAELOOP)
 365                 X(WSAENAMETOOLONG)      X(WSAEHOSTDOWN)
 366                 X(WSAEHOSTUNREACH)      X(WSAENOTEMPTY)
 367                 X(WSAEPROCLIM)  X(WSAEUSERS)
 368                 X(WSAEDQUOT)    X(WSAESTALE)
 369                 X(WSAEREMOTE)   X(WSASYSNOTREADY)
 370                 X(WSAVERNOTSUPPORTED)   X(WSANOTINITIALISED)
 371                 X(WSAEDISCON)   X(WSAENOMORE)
 372                 X(WSAECANCELLED)        X(WSAEINVALIDPROCTABLE)
 373                 X(WSAEINVALIDPROVIDER)  X(WSAEPROVIDERFAILEDINIT)
 374                 X(WSASYSCALLFAILURE)    X(WSASERVICE_NOT_FOUND)
 375                 X(WSATYPE_NOT_FOUND)    X(WSA_E_NO_MORE)
 376                 X(WSA_E_CANCELLED)      X(WSAEREFUSED)
 377 #undef X
 378         }
 379         snprintf ( errbuf, sizeof(errbuf), "Unknown socket error (%lu)", err );
 380         errbuf[sizeof(errbuf)-1] = '\0';
 381         return errbuf;
 382 #elif defined(UNIX)
 383         perror(errbuf);
 384         return errbuf;
 385 #else
 386 #  error unrecognized target
 387 #endif
 388 }
 389 #if defined(UNICODE) || defined(_UNICODE)
 390 in_addr_t suLookupAddress ( const wchar_t* pcHost )
 391 {
 392         int len = wcslen(pcHost);
 393         char* p = new char[len+1];
 394         wcstombs ( p, pcHost, len );
 395         p[len] = 0;
 396         in_addr_t rc = suLookupAddress ( p );
 397         delete[] p;
 398         return rc;
 399 }
 400 bool suBroadcast ( SOCKET so, u_short port, const wchar_t* buf, int len /* = -1 */ )
 401 {
 402         char* p = new char[len+1];
 403         wcstombs ( p, buf, len );
 404         p[len] = 0;
 405         bool rc = suBroadcast ( so, port, p, len );
 406         delete[] p;
 407         return rc;
 408 }
 409 int suRecv ( SOCKET so, wchar_t* buf, int buflen, int timeout )
 410 {
 411         char* p = new char[buflen+1];
 412         int rc = suRecv ( so, p, buflen, timeout );
 413         p[buflen] = 0;
 414         mbstowcs ( buf, p, buflen );
 415         delete[] p;
 416         return rc;
 417 }
 418 int suRecvFrom ( SOCKET so, wchar_t* buf, int buflen, int timeout, sockaddr_in* from, int* fromlen )
 419 {
 420         char* p = new char[buflen+1];
 421         int rc = suRecvFrom ( so, p, buflen, timeout, from, fromlen );
 422         p[buflen] = 0;
 423         mbs2wcs ( buf, p, buflen );
 424         delete[] p;
 425         return rc;
 426 }
 427 bool suBind ( SOCKET so, const wchar_t* szInterfaceAddress, u_short iListenPort, bool bReuseAddr /* = false */ )
 428 {
 429         int len = wcslen(szInterfaceAddress);
 430         char* p = new char[len+1];
 431         wcstombs ( p, szInterfaceAddress, len );
 432         p[len] = 0;
 433         bool rc = suBind ( so, p, iListenPort, bReuseAddr );
 434         delete[] p;
 435         return rc;
 436 }
 437 #endif//UNICODE
 438
 439 suBufferedRecvSocket::suBufferedRecvSocket ( SOCKET so )
 440         : suSocket ( so ), _off(0), _len(0)
 441 {
 442 }
 443
 444 int suBufferedRecvSocket::recvUntil ( std::string& buf, char until, int timeout )
 445 {
 446         if ( !_len )
 447                 _off = 0;
 448         else if ( _off > (sizeof(_buf)>>1) )
 449         {
 450                 memmove ( _buf, &_buf[_off], _len );
 451                 _off = 0;
 452         }
 453         char* poff = &_buf[_off];
 454         for ( ;; )
 455         {
 456                 char* p = (char*)memchr ( poff, until, _len );
 457                 if ( p /*&& p < &poff[_len]*/ )
 458                 {
 459                         int ret_len = p-poff+1;
 460                         buf.resize ( ret_len );
 461                         memmove ( &buf[0], poff, ret_len );
 462                         _off += ret_len;
 463                         _len -= ret_len;
 464                         return ret_len;
 465                 }
 466                 int rc = suRecv ( *this, &poff[_len], sizeof(_buf)-_len-_off, timeout );
 467                 if ( rc < 0 )
 468                 {
 469                         if ( _len )
 470                         {
 471                                 rc = _len;
 472                                 buf.resize ( rc );
 473                                 memmove ( &buf[0], &_buf[_off], rc );
 474                                 _len = 0;
 475                         }
 476                         return rc;
 477                 }
 478                 _len += rc;
 479         }
 480 }
 481
 482 void suBufferedRecvSocket::recvPending()
 483 {
 484         if ( !_len )
 485                 _off = 0;
 486         else if ( _off > (sizeof(_buf)>>1) )
 487         {
 488                 memmove ( _buf, &_buf[_off], _len );
 489                 _off = 0;
 490         }
 491         char* poff = &_buf[_off];
 492         while ( sizeof(_buf)-_len-_off )
 493         {
 494                 int rc = suRecv ( *this, &poff[_len], sizeof(_buf)-_len-_off, 1 );
 495                 if ( rc <= 0 )
 496                         break;
 497                 _len += rc;
 498         }
 499 }
 500
 501 bool suBufferedRecvSocket::recvInStr ( char c )
 502 {
 503         return NULL != memchr ( &_buf[_off], c, _len );
 504 }