devel/libmdw/trunk/libmdw/src/TcpServer.cpp

/***************************************************************************
 *   Copyright (C) 2010 by Jean-Baptiste Wons   *
 *   wonsjb@gmail.com   *
 *                                                                         *
 *   This program is free software; you can redistribute it and/or modify  *
 *   it under the terms of the GNU General Public License as published by  *
 *   the Free Software Foundation; either version 2 of the License, or     *
 *   (at your option) any later version.                                   *
 *                                                                         *
 *   This program is distributed in the hope that it will be useful,       *
 *   but WITHOUT ANY WARRANTY; without even the implied warranty of        *
 *   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the         *
 *   GNU General Public License for more details.                          *
 *                                                                         *
 *   You should have received a copy of the GNU General Public License     *
 *   along with this program; if not, write to the                         *
 *   Free Software Foundation, Inc.,                                       *
 *   59 Temple Place - Suite 330, Boston, MA  02111-1307, USA.             *
 ***************************************************************************/

#include <sys/socket.h> /* for socket(), bind(), and connect() */
#include <arpa/inet.h>  /* for sockaddr_in and inet_ntoa() */
#include <string.h>     /* for memset() */
#include <unistd.h>     /* for close() */
#include <errno.h>
#include <fcntl.h>

#include <vector>

#include "TcpServer.h"
#include "Error.h"
#include "Logger.h"

#include "IoStream.h"

using namespace std;

namespace mdw
{
namespace
{
const int MaxBackLog = 20;
const int MaxConnection = 1000;
const int MaxReadSize = 65536;
}

TcpServer::TcpServer (short port, ObjectPool<vector<char> > &pool)
{
    _pool = &pool;
    init (port);
}

TcpServer::TcpServer (short port)
{
    _pool = 0;
    init (port);
}

void TcpServer::init (short port)
{
    _connectedSocket.resize (MaxConnection);
    struct sockaddr_in localAddr; /* Local address */

    /* Create socket for incoming connections */
    if ( (_servSock = socket (PF_INET, SOCK_STREAM, IPPROTO_TCP)) < 0)
    {
        ERROR_ERRNO ("socket() failed", errno);
    }
    fcntl (_servSock, F_SETFL, O_NONBLOCK);

    int yes = 1;
    if (setsockopt (_servSock, SOL_SOCKET, SO_REUSEADDR, (char *) &yes, sizeof (yes)) < 0)
    {
        ::close (_servSock);
        ERROR_ERRNO ("setsockopt() failed", errno);
    }

    /* Construct local address structure */
    memset (&localAddr, 0, sizeof (localAddr)); /* Zero out structure */
    localAddr.sin_family = AF_INET;                /* Internet address family */
    localAddr.sin_addr.s_addr = htonl (INADDR_ANY); /* Any incoming interface */
    localAddr.sin_port = htons (port);             /* Local port */

    /* Bind to the local address */
    if (bind (_servSock, (struct sockaddr *) &localAddr, sizeof (localAddr)) < 0)
    {
        ::close (_servSock);
        ERROR_ERRNO ("bind() failed", errno);
    }

    /* Mark the socket so it will listen for incoming connections */
    if (listen (_servSock, MaxBackLog) < 0)
    {
        close (_servSock);
        ERROR_ERRNO ("listen() failed", errno);
    }

    _maxSelectFd = _servSock + 1;
}

Stream &TcpServer::outStream (int streamId)
{
    return _connectedSocket[streamId].stream;
}

TcpServer::~TcpServer()
{
    close (_servSock);
    for (int i = 0; i < _maxSelectFd; ++i)
    {
        if (_connectedSocket[i].connected)
        {
            close (i);
        }
    }
}

void TcpServer::accept (StreamServerVisitor &visitor, int timeOutSec, int timeOutMicro)
{
    struct timeval tv;
    tv.tv_sec = timeOutSec;
    tv.tv_usec = timeOutMicro;
    fd_set rfds;
    fd_set efds;
    FD_ZERO (&rfds);
    FD_ZERO (&efds);
    FD_SET (_servSock, &rfds);

    for (int i = 0; i < _maxSelectFd; ++i)
    {
        if (_connectedSocket[i].connected)
        {
            FD_SET (i, &rfds);
            FD_SET (i, &efds);
        }
    }

    int retval = select (_maxSelectFd, &rfds, NULL, &efds, &tv);
 
    if (retval == -1)
    {
        if (errno != EINTR)
        {
            ERROR_ERRNO ("select()", errno);
        }
        else
        {
            // EINTR, just return
            return;
        }
    }
    else if (retval == 0)
    {
        // timeout
        return;
    }

    if (FD_ISSET (_servSock, &rfds))
    {
        // Hum, new connection !!!
        struct sockaddr_in clientAddr;
        socklen_t size = sizeof (clientAddr);
        int newConnection = ::accept (_servSock, (struct sockaddr *) & clientAddr, &size);
        fcntl (newConnection, F_SETFL, O_NONBLOCK);
        if (newConnection >= 0)
        {
            // Yeah !!! new guy !!!
            string ip (inet_ntoa (clientAddr.sin_addr));
            _connectedSocket[newConnection].connected = true;
            _connectedSocket[newConnection].stream.reset (*this, newConnection);
            _connectedSocket[newConnection].info.setIp (ip);
            _connectedSocket[newConnection].info.setStreamId (newConnection);
            _maxSelectFd = max (_maxSelectFd, newConnection + 1);
            visitor.visitConnect (_connectedSocket[newConnection].info);
        }
    }

    for (int i = 0; i < _maxSelectFd; ++i)
    {
        if (_connectedSocket[i].connected && (FD_ISSET (i, &rfds) || FD_ISSET (i, &efds)))
        {
            // Hum, something to read
            static vector<char> buffer;

            vector<char> * bufferToUse;
            if (_pool == 0)
            {
                bufferToUse = &buffer;
            }
            else
            {
                bufferToUse = & _pool->obtain();
            }
            bufferToUse->reserve (MaxReadSize);
            int res = read (i, & (*bufferToUse->begin()), MaxReadSize);
            if (res < 0)
            {
                if (errno == EAGAIN)
                {
                    // spurious, ignore
                }
                else
                {
                    // hum, there was an error.
                    LOG_ERRNO ("Read error", errno);
                    // We then close the connection, and notify the visitor
                    _connectedSocket[i].connected = false;
                    close (i);
                    visitor.visitDisconnect (_connectedSocket[i].info);
                }
            }
            else if (res == 0)
            {
                // EOF ...
                // lets close it then
                _connectedSocket[i].connected = false;
                close (i);
                visitor.visitDisconnect (_connectedSocket[i].info);
            }
            else
            {
                // Yeah !!! data !!!
                visitor.visitRead (&*bufferToUse->begin(), res, _connectedSocket[i].info);
            }
        }
    }
}

void TcpServer::write (const char* buffer, int size, int socketId)
{
    if (_connectedSocket[socketId].connected)
    {
        int res = ::write (socketId, buffer, size);
        if (res < 0)
        {
            LOG_ERRNO ("Cannot write to socket " << socketId, errno);
        }
        else if (res != size)
        {
            LOG_WARN ("Could not write the whole thing (only " << res << " instead of " << size << ")");
        }
    }
}

void TcpServer::close (int sockedId)
{
    _connectedSocket[sockedId].connected = false;
    ::close (sockedId);
}

void TcpServer::TcpStream::reset (TcpServer &t, int streamId)
{
    _t = &t;
    _id = streamId;
}
void TcpServer::TcpStream::output (const char *buffer, int size)
{
    _t->write (buffer, size, _id);
}
}

---