connect.cpp

来自「包含客户和服务器的升级程序,在linux下可以运行的.」· C++ 代码 · 共 613 行

/*
 * Copyright (C) 2006, Binary Ma
 * Licence: GNU GPL 1991 - version 2
 * Bug report: [e-mail]
*/

#include "connect.h"
#include <errno.h>
#include <unistd.h>
#include <stdlib.h>
#include <stdio.h>
#include <fcntl.h>
#include <sys/socket.h>
#include <arpa/inet.h>
#include <netdb.h>

static const char* VERSION = "0.6.2";

const int ONTIME                = 1;
static const int BLOCKMODE      = MSG_DONTWAIT | MSG_NOSIGNAL;
static const int NETBUF         = 4096;
static const int INFOSIZE       = 256;
static const int PROSIZE        = 32;
static const char COMPART       = ':';
static const char* VER          = "1.2";

int select_io( int fd, int delay, bool is_in )
{
    if( fd < 0 )
        return -__LINE__;
    
    fd_set set;
    FD_ZERO( &set );
    FD_SET( fd, &set );
    timeval timeout;
    timeout.tv_sec = delay;
    timeout.tv_usec = 0;
    fd_set* read_p = &set;
    fd_set* write_p = NULL;
    if( !is_in )
    {
        read_p = NULL;
        write_p = &set;
    }

    int retval = select( fd + 1, read_p, write_p, NULL, delay < 0 ? NULL : &timeout );
    if( -1 == retval )
        return -__LINE__;
    
    if( 0 != retval )
        retval = FD_ISSET( fd, &set );
    
    return retval;
}

int select_i( int fd, int delay )
{
    return select_io( fd, delay, true );
}

int select_o( int fd, int delay )
{
    return select_io( fd, delay, false );
}

/*
  verify whether the reveive version match with current version
  and return the follow length of data indicate by protocol head */
static long exthead( char* prohead )
{
    if( NULL == prohead )
        return -__LINE__;

    // protocol head format: "1.2:1024:"
    char* point = strchr( prohead, COMPART );
    if( NULL == point )
        return -__LINE__;

    // verify version
    *point = 0;
    if( 0 != strcmp( prohead, VER ) )
        return -__LINE__;

    prohead = point + 1;
    point = strchr( prohead, COMPART );
    if( NULL == point )
        return -__LINE__;
	
	// parse data length
    *point = 0;
    return atol( prohead );
}

static int char2bin( char* info, struct stat* st )
{
    if( NULL == info || NULL == st )
        return -__LINE__;
	
    char* point = NULL;
    int count = 0;
    while( NULL != ( point = strchr( info, COMPART ) ) )
    {
        *point = 0;
        switch( count++ )
        {
        case 0:
            st->st_size = atoll( info );
            break;

        case 1:
            st->st_mode = atoi( info );
            break;

        // reserve extend
        default:
            return 0;
        }
		
        info = point + 1;
    }

    return 0;
}

tcpbase::tcpbase( int delay ): m_delay( delay )
{
    m_sock = -1;
}

tcpbase::tcpbase( int sock, int delay ): m_delay( delay )
{
    m_error = 0;
    m_sock = dup( sock );
    if( -1 == m_sock )
        m_error = -__LINE__;
}

tcpbase::~tcpbase()
{
    close_sock();
}

long tcpbase::recv( void* buf, long size )
{
    if( 0 != m_error || NULL == buf )
        return -__LINE__;

    long retval = select_i( m_sock, m_delay );
    if( retval != ONTIME )
        return -__LINE__;

    char prohead[PROSIZE];
    bzero( prohead, sizeof( prohead ) );
    long count      = 0;
    long tmp        = sizeof( prohead ) - 1;
    char* tmp_p     = prohead;
    
    // first, receive protocol head.
    // second, receive text.
    for( int i = 0; i < 2; i++ )
    {
        if( 1 == i )
        {
            tmp = exthead( prohead );
            if( tmp < 0 || tmp > size )
                return -__LINE__;

            tmp_p = (char*)buf;
        }

        count = 0;
        while( count < tmp )
        {
            retval = select_i( m_sock, m_delay );
            if( retval != ONTIME )
                return -__LINE__;
	
            retval = ::recv( m_sock, tmp_p + count, tmp - count, BLOCKMODE );
            if( -1 == retval )
            {
                if( EWOULDBLOCK == errno )
                    continue;
                else
                    return -__LINE__;
            }
			
            // indicate remote socket closed
            if( 0 == retval )
                return -__LINE__;
            
            count += retval;
        }
    }

    return count;
}

long tcpbase::send( const void* buf, long len )
{
    if( 0 != m_error || NULL == buf )
        return -__LINE__;

    long retval = select_o( m_sock, m_delay );
    if( retval != ONTIME )
        return -__LINE__;
    
    char prohead[PROSIZE];
    bzero( prohead, sizeof( prohead ) );
    snprintf( prohead, sizeof( prohead ), "%s%c%li%c", VER, COMPART, len, COMPART );
    long count          = 0;
    long tmp            = sizeof( prohead ) - 1;
    const char* tmp_p   = prohead;

    // first, send protocol head.
    // second, send text.
    for( int i = 0; i < 2; i++ )
    {
        if( 1 == i )
        {
            tmp = len;
            tmp_p = (const char*)buf;
        }

        count = 0;
        while( count < tmp )
        {
            retval = select_o( m_sock, m_delay );
            if( retval != ONTIME )
                return -__LINE__;
            
            retval = ::send( m_sock, tmp_p + count, tmp - count, BLOCKMODE );
            if( -1 == retval )
            {
                if( EWOULDBLOCK == errno )
                    continue;
                else
                    return -__LINE__;
            }
            
            count += retval;
        }
    }
    
    return count;
}

off_t tcpbase::recvfile( const char* file, bool brkmode )
{
    if( NULL == file )
        return -__LINE__;
    
    char info[INFOSIZE];
    bzero( info, sizeof( info ) );
    if( recv( info, sizeof( info ) - 1 ) < 0 )
        return -__LINE__;

    struct stat st;
    if( 0 != char2bin( info, &st ) )
        return -__LINE__;
    
    umask( 0 );
    int flag = O_RDWR | O_CREAT;
    if( brkmode )
    {
        flag |= O_APPEND;
    }
    else
    {
        // will create new file
        // dont need care file mode
        if( 0 == access( file, F_OK ) )
            remove( file );
    }
	
    int	fd = open( file, flag, st.st_mode );
    if( -1 == fd )
        return -__LINE__;
    
    long count = 0;
    off_t retval = 0;
    while( retval < st.st_size )
    {
        static char buf[NETBUF];
        bzero( buf, sizeof( buf ) );
        count = recv( buf, sizeof( buf ) );
        if( count < 0 )
        {
            close( fd );
            return -__LINE__;
        }

        if( count != write( fd, buf, count ) )
        {
            close( fd );
            return -__LINE__;
        }

        retval += count;
    }

    close( fd );
    return retval;
}

off_t tcpbase::sendfile( const char* file, off_t offset )
{
    if( NULL == file )
        return -__LINE__;

    int fd = open( file, O_RDONLY );
    if( -1 == fd )
        return -__LINE__;

    struct stat st;
    if( -1 == fstat( fd, &st ) )
    {
        close( fd );
        return -__LINE__;
    }

    st.st_size -= offset;
    if( st.st_size < 0 )
    {
        close( fd );
        return -__LINE__;
    }

    if( -1 == lseek( fd, offset, SEEK_SET ) )
    {
        close( fd );
        return -__LINE__;
    }

    char info[INFOSIZE];
    bzero( info, sizeof( info ) );

    long long tmp = st.st_size;
    snprintf( info, sizeof( info ), "%lli%c%i%c", tmp, COMPART, st.st_mode, COMPART );
    if( send( info, strlen( info ) ) < 0 )
    {
        close( fd );
        return -__LINE__;
    }

    off_t retval = 0;
    long count = 0;
    while( retval < st.st_size )
    {
        static char buf[NETBUF];
        bzero( buf, sizeof( buf ) );
        count = read( fd, buf, sizeof( buf ) );
        if( count <= 0 )
        {
            if( -1 == count )
            {
                close( fd );
                return -__LINE__;
            }

            // indicate end of file
            // but, impossible at here
            if( 0 == count )
                break;
        }

        if( send( buf, count ) < 0 )
        {
            close( fd );
            return -__LINE__;
        }

        retval += count;
    }
	
    close( fd );
    return retval;
}

void tcpbase::close_sock()
{
    if( -1 != m_sock )
    {
        close( m_sock );
        m_sock = -1;
    }
}

int tcpbase::dup_sock()
{
    return dup( m_sock );
}

int tcpbase::error()
{
    return m_error;
}

/*
// Obsolete function, see Listen
static int Listen_Obsolete( int iPort, int iBackLog )
{
	int iSock = socket( AF_INET, SOCK_STREAM, 0 );
	if( -1 == iSock )
		return -__LINE__;

	struct sockaddr_in stAddr;
	memset( &stAddr, 0, sizeof( stAddr ) );
	stAddr.sin_family = AF_INET;
	stAddr.sin_port = htons( iPort );
	stAddr.sin_addr.s_addr = htonl( INADDR_ANY );

	const int iReuse = 1;
	if( -1 == setsockopt( iSock, SOL_SOCKET, SO_REUSEADDR, &iReuse, sizeof( iReuse ) ) )
	{
	    close( iSock );
		return -__LINE__;
	}

	if( -1 == bind( iSock, (struct sockaddr*)&stAddr, sizeof( stAddr ) ) )
	{
	    close( iSock );
		return -__LINE__;
	}
	
	if( -1 == listen( iSock, iBackLog ) )
	{
	    close( iSock );
		return -__LINE__;
	}

	return iSock;
}

// Obsolete function, see Connect
static int Connect_Obsolete( int iPort, const char* pcHost )
{
	if( NULL == pcHost )
		return -__LINE__;

	int iSock = socket( AF_INET, SOCK_STREAM, 0 );
	if( -1 == iSock )
		return -__LINE__;
	
	struct sockaddr_in stAddr;
	memset( &stAddr, 0, sizeof( stAddr ) );
	stAddr.sin_family = AF_INET;
	stAddr.sin_port = htons( iPort );
	
	// The obsolete function is:
	// stAddr.sin_addr.s_addr = inet_addr( pcHost );
	if( inet_pton( AF_INET, pcHost, &stAddr.sin_addr ) <= 0 )
	{
	    close( iSock );
		return -__LINE__;
	}
		
	if( -1 == connect( iSock, ( struct sockaddr* )( &stAddr ), ( socklen_t )( sizeof( stAddr ) ) ) )
	{
	    close( iSock );
		return -__LINE__;
	}
	
	return iSock;
}
*/

static int s_listen( const char* host, const char* port, int backlog )
{
    if( NULL == port )
        return -__LINE__;
    
    struct addrinfo* res = NULL;
    struct addrinfo* res_head = NULL;
    struct addrinfo hint;
    bzero( &hint, sizeof( hint ) );
    
	hint.ai_flags = AI_PASSIVE;
	hint.ai_family = AF_UNSPEC;
	hint.ai_socktype = SOCK_STREAM;
	if( 0 != getaddrinfo( host, port, &hint, &res ) )
	    return -__LINE__;

    int sock = -1;
    const int reuse = 1;
    for( res_head = res; NULL != res; res = res->ai_next )
    {
        sock = socket( res->ai_family, res->ai_socktype, res->ai_protocol );
        if( -1 == sock )
            continue;

        setsockopt( sock, SOL_SOCKET, SO_REUSEADDR, &reuse, sizeof( reuse ) );
        if( 0 == bind( sock, res->ai_addr, res->ai_addrlen ) )
            break;
        
        close( sock );
    }

    freeaddrinfo( res_head );
	
    // the final socket() or bind() error
    if( NULL == res )
        return -__LINE__;
	
    if( -1 == listen( sock, backlog ) )
        return -__LINE__;
	
    return sock;
}

static int s_connect( const char* host, const char* port )
{
    if( NULL == port )
        return -__LINE__;
    
    struct addrinfo* res = NULL;
    struct addrinfo* res_head = NULL;
    struct addrinfo hint;
    bzero( &hint, sizeof( hint ) );

    hint.ai_family = AF_UNSPEC;
    hint.ai_socktype = SOCK_STREAM;
    if( 0 != getaddrinfo( host, port, &hint, &res ) )
        return -__LINE__;

    int sock = -1;
    for( res_head = res; NULL != res; res = res->ai_next )
    {
        sock = socket( res->ai_family, res->ai_socktype, res->ai_protocol );
        if( -1 == sock )
            continue;

        if( 0 == connect( sock, res->ai_addr, res->ai_addrlen ) )
            break;

        close( sock );
    }

    freeaddrinfo( res_head );
    if( NULL == res )
        return -__LINE__;

    return sock;
}

tcpserver::tcpserver( const char* host, const char* port, int delay, int backlog ): tcpbase( delay )
{
    m_error = 0;
    bzero( m_accept, sizeof( m_accept ) );
    if( backlog < 0 )
        backlog = 1024;
    	
    m_sock_listen = s_listen( host, port, backlog );
    if( m_sock_listen < 0 )
    {
        m_error = m_sock_listen;
        m_sock_listen = -1;
    }
}

tcpserver::~tcpserver()
{
	close_sock();
	close_listen();
}

int tcpserver::accept( int delay )
{
    struct sockaddr_in addr;
    socklen_t len = sizeof( addr );
    bzero( &addr, len );
    if( ONTIME != select_i( m_sock_listen, delay ) )
        return 1;
	
    // dont set m_error error, usually, the errno == EINTR
    m_sock = ::accept( m_sock_listen, (struct sockaddr*)&addr, &len );
    if( -1 == m_sock )
        return -__LINE__;

    bzero( m_accept, sizeof( m_accept ) );
    inet_ntop( AF_INET, &addr.sin_addr, m_accept, sizeof( m_accept ) );
    return 0;
}

const char* tcpserver::accept_addr()
{
    return m_accept;
}

void tcpserver::close_listen()
{
    if( -1 != m_sock_listen )
    {
        close( m_sock_listen );
        m_sock_listen = -1;
    }
}

tcpclient::tcpclient( const char* host, const char* port, int delay ): tcpbase( delay )
{
    m_error = 0;
    m_sock = s_connect( host, port );
    if( m_sock < 0 )
    {
        m_error = m_sock;
        m_sock = -1;
    }
}

tcpclient::~tcpclient()
{
    close_sock();
}