socket.cpp

通用网络游戏开发框架

// ============================================================================
// Socket implementation.
//
// (c) 2003 Ken Reed
//
// This is free software. You can redistribute it and/or modify it under the
// terms of the GNU General Public License version 2 as published by the Free
// Software Foundation.
// ============================================================================

#include "stdafx.h"

#include <fstream>
#include <string.h>
#include <winsock2.h>

#include <string>
#include <sstream>

#include "socket.h"


// ============================================================================
// Module global data
// ============================================================================

using namespace std;

namespace {
   bool winsock_not_initialised (true);
   int  socket_count            (0);
}


// ============================================================================
// Socket Stream Buffer
// ============================================================================

class Buffer : public streambuf
{
public:
   Buffer();
   
   void     reset();
   int      underflow();
   int      overflow(int c);
   int      xsputn(const char *s, int n);

   bool     tracing;
   bool     have_listen_socket;
   bool     have_data_socket;
   int      bytes_read;
   int      bytes_sent;
   int      allocated_port;
   ofstream trace;
   char     io_buf[256];
   SOCKET   listen_socket;                   // Used to listen for a connection
   SOCKET   data_socket;                     // Used for data transfer
};

#define TRACE(p) if (tracing) { trace << timestamp() << p; trace.flush();}


// ============================================================================
// Return a timestamp
// ============================================================================

static string timestamp()
{
   char buffer[32];
   GetTimeFormat(0, 0, 0, 0, buffer, sizeof(buffer));
                 
   string result(buffer);
   result += " ";

   return result;
}


// ============================================================================
// Socket Construction.
// ============================================================================

Socket::Socket(int port) : iostream(new Buffer)
{
   if (winsock_not_initialised) {
      WSADATA w;
      int result (WSAStartup(0x0101, &w));
      if (result != 0) {
         Exception e("Unable to initialise Windows sockets (WSAStartup)");
         RAISE(e);
      }
      winsock_not_initialised = false;
   }

   socket_count++;

   if (port >= 0) {
      bind(port);
   }
}


// ============================================================================
// Enable tracing
// ============================================================================

void Socket::set_trace(const char * filename)
{
   Buffer * buffer = (rdbuf());

   buffer->trace.open(filename);
   buffer->tracing = true;
}


// ============================================================================
// Get data socket
// ============================================================================

void * Socket::get_connection()
{
   Buffer * buffer = (rdbuf());
   return reinterpret_cast<void *>(buffer->data_socket);
}


// ============================================================================
// Set data socket
// ============================================================================

void Socket::set_connection(void * handle)
{
   Buffer * buffer = (rdbuf());

   buffer->data_socket      = reinterpret_cast<SOCKET>(handle);
   buffer->have_data_socket = true;
}


// ============================================================================
// Socket Destruction.
// ============================================================================

Socket::~Socket()
{
   Buffer * buffer = rdbuf();

   if (buffer->have_listen_socket) {
      closesocket(buffer->listen_socket);
   }

   if (buffer->have_data_socket) {
      closesocket(buffer->data_socket);
   }

   socket_count--;
   if (socket_count == 0) {
      WSACleanup();
      winsock_not_initialised = true;
   }

   if (buffer->tracing) {
      buffer->trace.close();
      buffer->tracing = false;
   }

   delete buffer;
}


// ============================================================================
// Get a pointer to the stream buffer.
// ============================================================================

Buffer * Socket::rdbuf()
{
   return static_cast<Buffer *>(ios::rdbuf());
}


// ============================================================================
// Get and optionally reset the "bytes read" counter.
// ============================================================================

int Socket::bytes_read(bool reset_count)
{
   if (reset_count) rdbuf()->bytes_read = 0;

   return rdbuf()->bytes_read;
}


// ============================================================================
// Get and optionally reset the "bytes sent" counter.
// ============================================================================

int Socket::bytes_sent(bool reset_count)
{
   if (reset_count) rdbuf()->bytes_sent = 0;

   return rdbuf()->bytes_sent;
}


// ============================================================================
// Prepare a socket for listening
// ============================================================================

void Socket::bind(int port_number) 
{
   Buffer * buffer = rdbuf();

   buffer->listen_socket = socket(AF_INET, SOCK_STREAM, 0);
   if (buffer->listen_socket == INVALID_SOCKET) {
      Sock_exception e("Unable to bind socket");
      RAISE(e);
   }

   buffer->bytes_read         = 0;
   buffer->bytes_sent         = 0;
   buffer->have_listen_socket = true;

   struct hostent * local_host;
   local_host = gethostbyname("localhost");
   if (local_host == 0) {
      Sock_exception e("Unable to get local host details");
      RAISE(e);
   }

   struct sockaddr_in address;
   memset(reinterpret_cast<void *>(&address), 0, sizeof(address));

   address.sin_family      = AF_INET;
   address.sin_port        = htons(port_number);
   address.sin_addr.s_addr = INADDR_ANY;

   int result = ::bind(buffer->listen_socket, 
                       reinterpret_cast<struct sockaddr *>(&address), 
                       sizeof(address));

   if (result != 0) {
      Sock_exception e("Unable to bind socket to specified port");
      RAISE(e);
   }

   int size = sizeof(address);
   result   = getsockname(buffer->listen_socket, 
                          reinterpret_cast<sockaddr FAR *>(&address), &size);
   if (result != 0) {
      Sock_exception e("Unable to read socket address");
      RAISE(e);
   }

   buffer->allocated_port = ntohs(address.sin_port);
}


// ============================================================================
// Get allocated port
// ============================================================================

int Socket::get_number()
{
   Buffer * buffer = rdbuf();
   return buffer->allocated_port;
}


// ============================================================================
// Listen for a connection
// ============================================================================

void Socket::listen()
{
   Buffer * buffer = rdbuf();

   int result = ::listen(buffer->listen_socket, SOMAXCONN);
   if (result != 0) {
      Sock_exception e("Unable to listen on specified socket");
      RAISE(e);
   }

   struct sockaddr address;
   int length(sizeof(address));

   buffer->data_socket = accept(buffer->listen_socket, &address, &length);
   if (buffer->data_socket == INVALID_SOCKET) {
      Sock_exception e("Unable to accept connection on socket");
      RAISE(e);
   }

   buffer->have_data_socket = true;
}


// ============================================================================
// Connect to an existing port
// ============================================================================

void Socket::connect(const char * host, const int port_number)
{
   Buffer * buffer = rdbuf();

   buffer->data_socket = socket(AF_INET, SOCK_STREAM, 0);
   if (buffer->data_socket == INVALID_SOCKET) {
      Sock_exception e("Unable to create socket");
      RAISE(e);
   }

   buffer->bytes_read       = 0;
   buffer->bytes_sent       = 0;
   buffer->have_data_socket = true;

   struct hostent * remote_host;
   remote_host = gethostbyname(host);
   if (remote_host == 0) {
      Sock_exception e("Cannot find remote host");
      RAISE(e);
   }

   struct sockaddr_in address;
   memset((void *) &address, 0, sizeof(address));

   address.sin_family      = AF_INET;
   address.sin_port        = htons(port_number);
   address.sin_addr.s_addr = *(
                  reinterpret_cast<unsigned long *>(remote_host->h_addr));

   int result = ::connect(buffer->data_socket, 
                          reinterpret_cast<struct sockaddr *>(&address), 
                          sizeof(address));

   if (result == SOCKET_ERROR) {
      Sock_exception e("Unable to make connection to remote system");
      RAISE(e);
   }
}


// ============================================================================
// getline
//
// The microsoft string::getline has a bug in it that will stop it working
// with a blocking socket call. Provide our own instead.
// ============================================================================

void Socket::getline(string & s)
{
   Buffer * buffer = (rdbuf());
   char   * p      = (buffer->io_buf);
   char     c        (buffer->sbumpc());

   for (int i = 0; i < sizeof(buffer->io_buf); i++) {
      if (c == '\n') {
         *p = '\0';
         break;
      }

      *p++ = c;
      c = buffer->sbumpc();
   }

   s = buffer->io_buf;
}


// ============================================================================
// Read binary bytes
// ============================================================================

void Socket::read_binary(void * destination, int buffer_size)
{
   Buffer * buffer = rdbuf();

   int    expected = buffer_size;
   char * p        = static_cast<char *>(destination);

   while (true) {
      int received (recv(buffer->data_socket, p, expected, 0));
      if (received == 0 || received == SOCKET_ERROR) {
         Sock_exception e("Socket binary read failed");
         RAISE(e);
      }

      expected -= received;
      if (expected == 0) break;
      p += received;
   }

   buffer->bytes_read += buffer_size;
}


// ============================================================================
// Write binary bytes
// ============================================================================

void Socket::write_binary(void * destination, int buffer_size)
{
   Buffer * buffer = rdbuf();

   int    togo = buffer_size;
   char * p    = static_cast<char *>(destination);

   while (true) {
      int sent (send(buffer->data_socket, p, togo, 0));
      if (sent == 0 || sent == SOCKET_ERROR) {
         Sock_exception e("Socket binary write failed");
         RAISE(e);
      }

      togo -= sent;
      if (togo == 0) break;
      p += sent;
   }

   buffer->bytes_sent += buffer_size;
}


// ============================================================================
// Close down the data transfer part of a socket
// ============================================================================

void Socket::close()
{
   Buffer * buffer = rdbuf();

   if (buffer->have_data_socket) {
      closesocket(buffer->data_socket);
      buffer->have_data_socket = false;
   }

   buffer->reset();
}


// ============================================================================
// Buffer construction
// ============================================================================

Buffer::Buffer() : have_listen_socket(false), have_data_socket(false),
                   bytes_read(0), bytes_sent(0), tracing(false)
{
}


// ============================================================================
// Read a line from the socket into the input buffer
// ============================================================================

int Buffer::underflow()
{
   char * p;
   char   c;

   for (p = io_buf; p < (io_buf + sizeof(io_buf) - 1); ) {
      int result = recv(data_socket, &c, 1, 0);
      if (result != 1) {
         Sock_exception e("Error reading data from socket");
         TRACE(e.get_error());
         RAISE(e);
      }

      if (c == '\r') {
         continue;
      }

      *p++ = c;
      bytes_read++;

      if (c == '\n') {
         break;
      }
   } 

   *p = '\0';
   setg(io_buf, io_buf, p);

   TRACE("R: " << io_buf );

   return io_buf[0];
}


// ============================================================================
// Write a sequence of characters
//
// If we didn't implement this function characters would be sent one at a time
// using the overflow function.
// ============================================================================

int Buffer::xsputn(const char * buffer, int length)
{
   int          to_go (length);
   const char * p =   (buffer);

   bytes_sent += length;

   TRACE("W: " << buffer << "\n");
   while (true) {
      int sent (send(data_socket, p, to_go, 0));

      if (sent == SOCKET_ERROR) {
         Sock_exception e("Cannot write to socket");
         TRACE(e.get_error());
         RAISE(e);
      }

      to_go -= sent;
      if (to_go <= 0) break;
      p += sent;
   }

   return length;
}


// ============================================================================
// Flush the output stream.
// ============================================================================

int Buffer::overflow(int c)
{
   char buffer(c);
   TRACE("C: ");
   if (tracing) {
      trace.put(c);
      trace.put('\n');
   }

   bytes_sent++;
   return send(data_socket, &buffer, 1, 0);
}


// ============================================================================
// Reset the input buffer
// ============================================================================

void Buffer::reset()
{
   setg(0, 0, 0);
}


// ============================================================================
// Socket Exception Construction
// ============================================================================

Sock_exception::Sock_exception(const string & reason)
   : Exception(reason)
{
   ostringstream s;
   s << "Socket error code: " << WSAGetLastError();
   error_text = s.str();
}