mudgangster

Tiny, scriptable MUD client
Log | Files | Refs | README

TracySocket.cpp (11684B)


      1 #include <assert.h>
      2 #include <new>
      3 #include <stdio.h>
      4 #include <stdlib.h>
      5 #include <string.h>
      6 #include <sys/types.h>
      7 
      8 #include "TracyAlloc.hpp"
      9 #include "TracySocket.hpp"
     10 
     11 #ifdef _WIN32
     12 #  ifndef NOMINMAX
     13 #    define NOMINMAX
     14 #  endif
     15 #  include <winsock2.h>
     16 #  include <ws2tcpip.h>
     17 #  ifdef _MSC_VER
     18 #    pragma warning(disable:4244)
     19 #    pragma warning(disable:4267)
     20 #  endif
     21 #  define poll WSAPoll
     22 #else
     23 #  include <arpa/inet.h>
     24 #  include <sys/socket.h>
     25 #  include <sys/param.h>
     26 #  include <netinet/in.h>
     27 #  include <netdb.h>
     28 #  include <unistd.h>
     29 #  include <poll.h>
     30 #endif
     31 
     32 #ifndef MSG_NOSIGNAL
     33 #  define MSG_NOSIGNAL 0
     34 #endif
     35 
     36 namespace tracy
     37 {
     38 
     39 #ifdef _WIN32
     40 typedef SOCKET socket_t;
     41 #else
     42 typedef int socket_t;
     43 #endif
     44 
     45 #ifdef _WIN32
     46 struct __wsinit
     47 {
     48     __wsinit()
     49     {
     50         WSADATA wsaData;
     51         if( WSAStartup( MAKEWORD( 2, 2 ), &wsaData ) != 0 )
     52         {
     53             fprintf( stderr, "Cannot init winsock.\n" );
     54             exit( 1 );
     55         }
     56     }
     57 };
     58 
     59 void InitWinSock()
     60 {
     61     static __wsinit init;
     62 }
     63 #endif
     64 
     65 Socket::Socket()
     66     : m_buf( (char*)tracy_malloc( BufSize ) )
     67     , m_bufPtr( nullptr )
     68     , m_sock( -1 )
     69     , m_bufLeft( 0 )
     70 {
     71 #ifdef _WIN32
     72     InitWinSock();
     73 #endif
     74 }
     75 
     76 Socket::Socket( int sock )
     77     : m_buf( (char*)tracy_malloc( BufSize ) )
     78     , m_bufPtr( nullptr )
     79     , m_sock( sock )
     80     , m_bufLeft( 0 )
     81 {
     82 }
     83 
     84 Socket::~Socket()
     85 {
     86     tracy_free( m_buf );
     87     if( m_sock != -1 )
     88     {
     89         Close();
     90     }
     91 }
     92 
     93 bool Socket::Connect( const char* addr, int port )
     94 {
     95     assert( m_sock == -1 );
     96 
     97     struct addrinfo hints;
     98     struct addrinfo *res, *ptr;
     99 
    100     memset( &hints, 0, sizeof( hints ) );
    101     hints.ai_family = AF_UNSPEC;
    102     hints.ai_socktype = SOCK_STREAM;
    103 
    104     char portbuf[32];
    105     sprintf( portbuf, "%i", port );
    106 
    107     if( getaddrinfo( addr, portbuf, &hints, &res ) != 0 ) return false;
    108     int sock = 0;
    109     for( ptr = res; ptr; ptr = ptr->ai_next )
    110     {
    111         if( ( sock = socket( ptr->ai_family, ptr->ai_socktype, ptr->ai_protocol ) ) == -1 ) continue;
    112 #if defined __APPLE__
    113         int val = 1;
    114         setsockopt( sock, SOL_SOCKET, SO_NOSIGPIPE, &val, sizeof( val ) );
    115 #endif
    116         if( connect( sock, ptr->ai_addr, ptr->ai_addrlen ) == -1 )
    117         {
    118 #ifdef _WIN32
    119             closesocket( sock );
    120 #else
    121             close( sock );
    122 #endif
    123             continue;
    124         }
    125         break;
    126     }
    127     freeaddrinfo( res );
    128     if( !ptr ) return false;
    129 
    130     m_sock = sock;
    131     return true;
    132 }
    133 
    134 void Socket::Close()
    135 {
    136     assert( m_sock != -1 );
    137 #ifdef _WIN32
    138     closesocket( m_sock );
    139 #else
    140     close( m_sock );
    141 #endif
    142     m_sock = -1;
    143 }
    144 
    145 int Socket::Send( const void* _buf, int len )
    146 {
    147     auto buf = (const char*)_buf;
    148     assert( m_sock != -1 );
    149     auto start = buf;
    150     while( len > 0 )
    151     {
    152         auto ret = send( m_sock, buf, len, MSG_NOSIGNAL );
    153         if( ret == -1 ) return -1;
    154         len -= ret;
    155         buf += ret;
    156     }
    157     return int( buf - start );
    158 }
    159 
    160 int Socket::GetSendBufSize()
    161 {
    162     int bufSize;
    163 #if defined _WIN32 || defined __CYGWIN__
    164     int sz = sizeof( bufSize );
    165     getsockopt( m_sock, SOL_SOCKET, SO_SNDBUF, (char*)&bufSize, &sz );
    166 #else
    167     socklen_t sz = sizeof( bufSize );
    168     getsockopt( m_sock, SOL_SOCKET, SO_SNDBUF, &bufSize, &sz );
    169 #endif
    170     return bufSize;
    171 }
    172 
    173 int Socket::RecvBuffered( void* buf, int len, int timeout )
    174 {
    175     if( len <= m_bufLeft )
    176     {
    177         memcpy( buf, m_bufPtr, len );
    178         m_bufPtr += len;
    179         m_bufLeft -= len;
    180         return len;
    181     }
    182 
    183     if( m_bufLeft > 0 )
    184     {
    185         memcpy( buf, m_bufPtr, m_bufLeft );
    186         const auto ret = m_bufLeft;
    187         m_bufLeft = 0;
    188         return ret;
    189     }
    190 
    191     if( len >= BufSize ) return Recv( buf, len, timeout );
    192 
    193     m_bufLeft = Recv( m_buf, BufSize, timeout );
    194     if( m_bufLeft <= 0 ) return m_bufLeft;
    195 
    196     const auto sz = len < m_bufLeft ? len : m_bufLeft;
    197     memcpy( buf, m_buf, sz );
    198     m_bufPtr = m_buf + sz;
    199     m_bufLeft -= sz;
    200     return sz;
    201 }
    202 
    203 int Socket::Recv( void* _buf, int len, int timeout )
    204 {
    205     auto buf = (char*)_buf;
    206 
    207     struct pollfd fd;
    208     fd.fd = (socket_t)m_sock;
    209     fd.events = POLLIN;
    210 
    211     if( poll( &fd, 1, timeout ) > 0 )
    212     {
    213         return recv( m_sock, buf, len, 0 );
    214     }
    215     else
    216     {
    217         return -1;
    218     }
    219 }
    220 
    221 bool Socket::Read( void* _buf, int len, int timeout, std::function<bool()> exitCb )
    222 {
    223     auto buf = (char*)_buf;
    224 
    225     while( len > 0 )
    226     {
    227         if( exitCb() ) return false;
    228         const auto sz = RecvBuffered( buf, len, timeout );
    229         switch( sz )
    230         {
    231         case 0:
    232             return false;
    233         case -1:
    234 #ifdef _WIN32
    235         {
    236             auto err = WSAGetLastError();
    237             if( err == WSAECONNABORTED || err == WSAECONNRESET ) return false;
    238         }
    239 #endif
    240             break;
    241         default:
    242             len -= sz;
    243             buf += sz;
    244             break;
    245         }
    246     }
    247 
    248     return true;
    249 }
    250 
    251 bool Socket::ReadRaw( void* _buf, int len, int timeout )
    252 {
    253     auto buf = (char*)_buf;
    254     while( len > 0 )
    255     {
    256         const auto sz = Recv( buf, len, timeout );
    257         if( sz <= 0 ) return false;
    258         len -= sz;
    259         buf += sz;
    260     }
    261     return true;
    262 }
    263 
    264 bool Socket::HasData()
    265 {
    266     if( m_bufLeft > 0 ) return true;
    267 
    268     struct pollfd fd;
    269     fd.fd = (socket_t)m_sock;
    270     fd.events = POLLIN;
    271 
    272     return poll( &fd, 1, 0 ) > 0;
    273 }
    274 
    275 
    276 ListenSocket::ListenSocket()
    277     : m_sock( -1 )
    278 {
    279 #ifdef _WIN32
    280     InitWinSock();
    281 #endif
    282 }
    283 
    284 ListenSocket::~ListenSocket()
    285 {
    286     if( m_sock != -1 ) Close();
    287 }
    288 
    289 bool ListenSocket::Listen( int port, int backlog )
    290 {
    291     assert( m_sock == -1 );
    292 
    293     struct addrinfo* res;
    294     struct addrinfo hints;
    295 
    296     memset( &hints, 0, sizeof( hints ) );
    297     hints.ai_family = AF_INET6;
    298     hints.ai_socktype = SOCK_STREAM;
    299     hints.ai_flags = AI_PASSIVE;
    300 
    301     char portbuf[32];
    302     sprintf( portbuf, "%i", port );
    303 
    304     if( getaddrinfo( nullptr, portbuf, &hints, &res ) != 0 ) return false;
    305 
    306     m_sock = socket( res->ai_family, res->ai_socktype, res->ai_protocol );
    307 #if defined _WIN32 || defined __CYGWIN__
    308     unsigned long val = 0;
    309     setsockopt( m_sock, IPPROTO_IPV6, IPV6_V6ONLY, (const char*)&val, sizeof( val ) );
    310 #elif defined BSD
    311     int val = 0;
    312     setsockopt( m_sock, IPPROTO_IPV6, IPV6_V6ONLY, (const char*)&val, sizeof( val ) );
    313     val = 1;
    314     setsockopt( m_sock, SOL_SOCKET, SO_REUSEADDR, &val, sizeof( val ) );
    315 #else
    316     int val = 1;
    317     setsockopt( m_sock, SOL_SOCKET, SO_REUSEADDR, &val, sizeof( val ) );
    318 #endif
    319     if( bind( m_sock, res->ai_addr, res->ai_addrlen ) == -1 ) { freeaddrinfo( res ); return false; }
    320     if( listen( m_sock, backlog ) == -1 ) { freeaddrinfo( res ); return false; }
    321     freeaddrinfo( res );
    322     return true;
    323 }
    324 
    325 Socket* ListenSocket::Accept()
    326 {
    327     struct sockaddr_storage remote;
    328     socklen_t sz = sizeof( remote );
    329 
    330     struct pollfd fd;
    331     fd.fd = (socket_t)m_sock;
    332     fd.events = POLLIN;
    333 
    334     if( poll( &fd, 1, 10 ) > 0 )
    335     {
    336         int sock = accept( m_sock, (sockaddr*)&remote, &sz);
    337         if( sock == -1 ) return nullptr;
    338 
    339 #if defined __APPLE__
    340         int val = 1;
    341         setsockopt( sock, SOL_SOCKET, SO_NOSIGPIPE, &val, sizeof( val ) );
    342 #endif
    343 
    344         auto ptr = (Socket*)tracy_malloc( sizeof( Socket ) );
    345         new(ptr) Socket( sock );
    346         return ptr;
    347     }
    348     else
    349     {
    350         return nullptr;
    351     }
    352 }
    353 
    354 void ListenSocket::Close()
    355 {
    356     assert( m_sock != -1 );
    357 #ifdef _WIN32
    358     closesocket( m_sock );
    359 #else
    360     close( m_sock );
    361 #endif
    362     m_sock = -1;
    363 }
    364 
    365 UdpBroadcast::UdpBroadcast()
    366     : m_sock( -1 )
    367 {
    368 #ifdef _WIN32
    369     InitWinSock();
    370 #endif
    371 }
    372 
    373 UdpBroadcast::~UdpBroadcast()
    374 {
    375     if( m_sock != -1 ) Close();
    376 }
    377 
    378 bool UdpBroadcast::Open( const char* addr, int port )
    379 {
    380     assert( m_sock == -1 );
    381 
    382     struct addrinfo hints;
    383     struct addrinfo *res, *ptr;
    384 
    385     memset( &hints, 0, sizeof( hints ) );
    386     hints.ai_family = AF_INET;
    387     hints.ai_socktype = SOCK_DGRAM;
    388 
    389     char portbuf[32];
    390     sprintf( portbuf, "%i", port );
    391 
    392     if( getaddrinfo( addr, portbuf, &hints, &res ) != 0 ) return false;
    393     int sock = 0;
    394     for( ptr = res; ptr; ptr = ptr->ai_next )
    395     {
    396         if( ( sock = socket( ptr->ai_family, ptr->ai_socktype, ptr->ai_protocol ) ) == -1 ) continue;
    397 #if defined __APPLE__
    398         int val = 1;
    399         setsockopt( sock, SOL_SOCKET, SO_NOSIGPIPE, &val, sizeof( val ) );
    400 #endif
    401 #if defined _WIN32 || defined __CYGWIN__
    402         unsigned long broadcast = 1;
    403         if( setsockopt( sock, SOL_SOCKET, SO_BROADCAST, (const char*)&broadcast, sizeof( broadcast ) ) == -1 )
    404 #else
    405         int broadcast = 1;
    406         if( setsockopt( sock, SOL_SOCKET, SO_BROADCAST, &broadcast, sizeof( broadcast ) ) == -1 )
    407 #endif
    408         {
    409 #ifdef _WIN32
    410             closesocket( sock );
    411 #else
    412             close( sock );
    413 #endif
    414             continue;
    415         }
    416         break;
    417     }
    418     freeaddrinfo( res );
    419     if( !ptr ) return false;
    420 
    421     m_sock = sock;
    422     return true;
    423 }
    424 
    425 void UdpBroadcast::Close()
    426 {
    427     assert( m_sock != -1 );
    428 #ifdef _WIN32
    429     closesocket( m_sock );
    430 #else
    431     close( m_sock );
    432 #endif
    433     m_sock = -1;
    434 }
    435 
    436 int UdpBroadcast::Send( int port, const void* data, int len )
    437 {
    438     assert( m_sock != -1 );
    439     struct sockaddr_in addr;
    440     addr.sin_family = AF_INET;
    441     addr.sin_port = htons( port );
    442     addr.sin_addr.s_addr = INADDR_BROADCAST;
    443     return sendto( m_sock, (const char*)data, len, MSG_NOSIGNAL, (sockaddr*)&addr, sizeof( addr ) );
    444 }
    445 
    446 IpAddress::IpAddress()
    447     : m_number( 0 )
    448 {
    449     *m_text = '\0';
    450 }
    451 
    452 IpAddress::~IpAddress()
    453 {
    454 }
    455 
    456 void IpAddress::Set( const struct sockaddr& addr )
    457 {
    458 #if __MINGW32__
    459     auto ai = (struct sockaddr_in*)&addr;
    460 #else
    461     auto ai = (const struct sockaddr_in*)&addr;
    462 #endif
    463     inet_ntop( AF_INET, &ai->sin_addr, m_text, 17 );
    464     m_number = ai->sin_addr.s_addr;
    465 }
    466 
    467 UdpListen::UdpListen()
    468     : m_sock( -1 )
    469 {
    470 #ifdef _WIN32
    471     InitWinSock();
    472 #endif
    473 }
    474 
    475 UdpListen::~UdpListen()
    476 {
    477     if( m_sock != -1 ) Close();
    478 }
    479 
    480 bool UdpListen::Listen( int port )
    481 {
    482     assert( m_sock == -1 );
    483 
    484     int sock;
    485     if( ( sock = socket( AF_INET, SOCK_DGRAM, 0 ) ) == -1 ) return false;
    486 
    487 #if defined __APPLE__
    488     int val = 1;
    489     setsockopt( sock, SOL_SOCKET, SO_NOSIGPIPE, &val, sizeof( val ) );
    490 #endif
    491 #if defined _WIN32 || defined __CYGWIN__
    492     unsigned long reuse = 1;
    493     setsockopt( m_sock, SOL_SOCKET, SO_REUSEADDR, (const char*)&reuse, sizeof( reuse ) );
    494 #else
    495     int reuse = 1;
    496     setsockopt( m_sock, SOL_SOCKET, SO_REUSEADDR, &reuse, sizeof( reuse ) );
    497 #endif
    498 #if defined _WIN32 || defined __CYGWIN__
    499     unsigned long broadcast = 1;
    500     if( setsockopt( sock, SOL_SOCKET, SO_BROADCAST, (const char*)&broadcast, sizeof( broadcast ) ) == -1 )
    501 #else
    502     int broadcast = 1;
    503     if( setsockopt( sock, SOL_SOCKET, SO_BROADCAST, &broadcast, sizeof( broadcast ) ) == -1 )
    504 #endif
    505     {
    506 #ifdef _WIN32
    507         closesocket( sock );
    508 #else
    509         close( sock );
    510 #endif
    511         return false;
    512     }
    513 
    514     struct sockaddr_in addr;
    515     addr.sin_family = AF_INET;
    516     addr.sin_port = htons( port );
    517     addr.sin_addr.s_addr = INADDR_ANY;
    518 
    519     if( bind( sock, (sockaddr*)&addr, sizeof( addr ) ) == -1 )
    520     {
    521 #ifdef _WIN32
    522         closesocket( sock );
    523 #else
    524         close( sock );
    525 #endif
    526         return false;
    527     }
    528 
    529     m_sock = sock;
    530     return true;
    531 }
    532 
    533 void UdpListen::Close()
    534 {
    535     assert( m_sock != -1 );
    536 #ifdef _WIN32
    537     closesocket( m_sock );
    538 #else
    539     close( m_sock );
    540 #endif
    541     m_sock = -1;
    542 }
    543 
    544 const char* UdpListen::Read( size_t& len, IpAddress& addr )
    545 {
    546     static char buf[2048];
    547 
    548     struct pollfd fd;
    549     fd.fd = (socket_t)m_sock;
    550     fd.events = POLLIN;
    551     if( poll( &fd, 1, 10 ) <= 0 ) return nullptr;
    552 
    553     sockaddr sa;
    554     socklen_t salen = sizeof( struct sockaddr );
    555     len = (size_t)recvfrom( m_sock, buf, 2048, 0, &sa, &salen );
    556     addr.Set( sa );
    557 
    558     return buf;
    559 }
    560 
    561 }