netkite.cpp

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

        return -__LINE__;

    return atoi( respond );
}

// transmit a file to client
int transmit( const char* file, const filemap* fmap, tcpserver* sock )
{
    if( NULL == file || NULL == fmap || NULL == sock )
        return -__LINE__;

    send_order( sock, ORD_UNFINISH );
    int respond = receive_order( sock );
    if( ORD_AFFIRM != respond )
        return -__LINE__;
    
    char info[SIZE_INFO];
    snprintf( info, sizeof( info ),
        "%lli%c%i%c%i%c%i%c%i%c%s%c%i%c%s%c",
        (long long)fmap->st.st_size,    INFO_COMPART,
        (int)fmap->st.st_mode,          INFO_COMPART,
        (int)fmap->st.st_uid,           INFO_COMPART,
        (int)fmap->st.st_gid,           INFO_COMPART,
        (int)fmap->st.st_rdev,          INFO_COMPART,
        fmap->device,                   INFO_COMPART,
        fmap->act,                      INFO_COMPART,
        fmap->md5,                      INFO_COMPART );

    sock->send( info, strlen( info ) );
    respond = receive_order( sock );
    if( ORD_AFFIRM != respond )
        return -__LINE__;

    // the file name maybe contain INFO_COMPART,
    // will make trouble for extmap, so, do it like this
    sock->send( fmap->name, strlen( fmap->name ) );
    respond = receive_order( sock );
    if( respond < 0 )
        return -__LINE__;
    
    int retval = 0;
    if( ORD_UPGRADE == respond )
    {
        if( S_ISREG( fmap->st.st_mode ) )
        {
            if( sock->sendfile( file ) < 0 )
                retval = -__LINE__;
        }
        else
        if( S_ISLNK( fmap->st.st_mode ) )
        {
            char buflink[SIZE_PATH];
            bzero( buflink, sizeof( buflink ) );
            readlink( file, buflink, sizeof( buflink ) - 1 );
            if( sock->send( buflink, strlen( buflink ) ) < 0 )
                retval = -__LINE__;
        }

        log_server( retval, sock->accept_addr(), file );
    }
    
    return retval;
}

void sigclient( int )
{
    return;
}

void siglock( bool lock )
{
    static sigset_t save;
    static sigset_t block;
    if( lock )
    {
        sigfillset( &block );
        sigprocmask( SIG_BLOCK, &block, &save );
    }
    else
    {
        sigprocmask( SIG_SETMASK, &save, NULL );
    }
}

int start_server( tcpserver* sock, const char* root )
{
    if( NULL == sock || NULL == root )
        return -__LINE__;
    
    // send or receive PROTOCOL first;
    char auth[16];
    bzero( auth, sizeof( auth ) );
    sock->recv( auth, sizeof( auth ) - 1 );
    if( 0 != strcmp( auth, PROTOCOL ) )
        return 0;

    send_order( sock, ORD_AFFIRM );
    if( -1 == chroot( root ) || -1 == chdir( "/" ) )
    {
        ERROR();
        return -__LINE__;
    }

    bzero( g_module_name, sizeof( g_module_name ) );
    if( sock->recv( g_module_name, sizeof( g_module_name ) - 1 ) < 0 )
        return -__LINE__;
    
    if( SUCCESS == nftw( "/", search, 512, FTW_PHYS ) )
    {
        // if g_module_path == "/", meet root
        while( 0 != strcmp( g_module_path, "/" ) )
        {
            DIR* dir = opendir( g_module_path );
            if( NULL == dir )
                return -__LINE__;
            
            struct dirent* ent = NULL;
            chdir( g_module_path );
            const filemap* fmap = NULL;
            while( NULL != ( ent = readdir( dir ) ) )
            {
                if( 0 == isdir( ent->d_name ) )
                    continue;
            
                fmap = fillmap( ent->d_name );
                if( NULL == fmap || 0 == *fmap->name )
                    continue;

                if( 0 != transmit( ent->d_name, fmap, sock ) )
                    break;
            }
            
            closedir( dir );
            dirname( g_module_path );
        }
    }
	
    send_order( sock, ORD_FINISH );
    return 0;
}

// if zero can be continue
int hash_check( hash* hs, const struct filemap* fmap )
{
    // overlap mode, shake head...
    if( NULL == hs )
        return 0;
    
    int offset = 0;
    if( '/' == *fmap->name )
        offset = 1;

    char key[strlen( fmap->device ) + strlen( fmap->name ) + 8];
    sprintf( key, "%s:%s", fmap->device, fmap->name + offset );
    ENTRY ent;
    ent.key = key;
    ent.data = NULL;
    int retval = 1;
    if( NULL == hs->search( ent.key ) )
    {
        retval = 0;
        hs->insert( ent, 0 );
    }

    return retval;
}

int start_client( tcpclient* sock, const char* module, const char* tmppath, hash* hs )
{
    if( NULL == sock || NULL == module || NULL == tmppath )
        return -__LINE__;
    
    struct filemap fmap;
    char info[SIZE_INFO];
    int respond = 0;
    sock->send( PROTOCOL, strlen( PROTOCOL ) );
    respond = receive_order( sock );
    if( ORD_AFFIRM != respond )
    {
        log_finish( respond );
        return -__LINE__;
    }

    sock->send( module, strlen( module ) );
    
    while( 1 )
    {
        respond = receive_order( sock );
        if( respond < 0 )
        {
            ERROR();
            return -__LINE__;
        }

        // not have suit upgrade file
        // unique quit way
        if( ORD_UNFINISH != respond )
        {
            log_finish( respond );
            break;
        }
        
        send_order( sock, ORD_AFFIRM );
        bzero( &fmap, sizeof( fmap ) );
        bzero( info, sizeof( info ) );
        if( sock->recv( info, sizeof( info ) - 1 ) < 0 )
        {
            ERROR();
            return -__LINE__;
        }

        extmap( info, &fmap );
        send_order( sock, ORD_AFFIRM );
        if( sock->recv( fmap.name, sizeof( fmap.name ) - 1 ) < 0 )
        {
            ERROR();
            return -__LINE__;
        }

        if( 0 != hash_check( hs, &fmap ) )
        {
            send_order( sock, ORD_NEXT );
            continue;
        }
        
        const char* destroot = mkroot( tmppath, fmap.device );
        respond = ORD_NEXT;
        int len = 0;
        if( NULL != destroot )
            len = strlen( destroot );

        char destfile[ strlen( fmap.name ) + len + 8];
        if( NULL != destroot )
        {
            strcpy( destfile, destroot );
            if( '/' != *fmap.name )
                strcat( destfile, "/" );
            
            strcat( destfile, fmap.name );
            if( 0 == isrequired( destfile, &fmap ) )
                respond = ORD_UPGRADE;
        }
        else
        {
            // maybe without such device
            log_client( -__LINE__, "error", &fmap );
        }
        
        send_order( sock, respond );
        if( ORD_UPGRADE == respond )
        {
            int retval = 0;
            if( S_ISREG( fmap.st.st_mode ) )
            {
                retval = -__LINE__;
                char tmpfile[strlen( tmppath ) + strlen( XXX ) + 8];
                snprintf( tmpfile, sizeof( tmpfile ), "%s/tmp.%s", tmppath, XXX );
                int tmpfd = mkstemp( tmpfile );
                if( -1 == tmpfd )
                {
                    ERROR();
                    return -__LINE__;
                }
                
                close( tmpfd );
                sock->recvfile( tmpfile );
                const char* cal = calculate_md5( tmpfile );
                if( NULL != cal && 0 == strcmp( cal, fmap.md5 ) )
                {
                    // this time, block SIGTERM at least
                    siglock( 1 );
                    retval = upgrade_file( tmpfile, destfile, &fmap );
                    siglock( 0 );
                }

                if( 0 == access( tmpfile, F_OK ) )
                    remove( tmpfile );
            }
            else
            if( S_ISLNK( fmap.st.st_mode ) )
            {
                retval = -__LINE__;
                char link[SIZE_PATH];
                bzero( link, sizeof( link ) );
                if( sock->recv( link, sizeof( link ) - 1 ) > 0 )
                {
                    retval = -__LINE__;
                    mdmethod mdc( MDC_MD5 );
                    const char* cal = mdc.sum_data( link, strlen( link ) );
                    if( NULL != cal && 0 == strcmp( cal, fmap.md5 ) )
                    {
                        siglock( 1 );
                        retval = upgrade_link( destfile, link, &fmap );
                        siglock( 0 );
                    }
                }
            }
            else
            {
                siglock( 1 );
                retval = upgrade_node( destfile, &fmap );
                siglock( 0 );
            }

            log_client( retval, "upgrade", &fmap );
        }

        if( NULL != destroot && 0 != *destroot )
        {
            umount( destroot );
            rmdir( destroot );
        }
    }

    return 0;
}

int main( int argc, char* argv[] )
{
    struct option opt;
    init( &opt );
    parse_option( argc, argv, &opt );
    parse_config( &opt );

    umask( 0 );
    daemon( 0, 0 );
    signal( SIGCHLD, SIG_IGN );
    if( opt.log_enable )
    {
        char startinfo[512];
        sprintf( startinfo, "netkite version: %s starting", VERSION );
        const char* ident = "client";
        if( opt.server_mode )
            ident = "server";
        
        duplog( ident, opt.log_file );
        startlog( startinfo );
        if( opt.server_mode )
        {
            snprintf( startinfo, sizeof( startinfo ), "server will change root to %s", opt.root );
            startlog( startinfo );
        }
    }

    if( opt.server_mode )
    {
        tcpserver sock( opt.host, opt.port, opt.timeout, opt.backlog );
        if( 0 != sock.error() )
        {
            ERROR();
            return -__LINE__;
        }

        while( 1 )
        {
            if( 0 != sock.accept() )
            {
                if( EINTR == errno )
                {
                    continue;
                }
                else
                {
                    ERROR();
                    return -__LINE__;
                }
            }
    
            pid_t pid = fork();
            if( pid < 0 )
            {
                ERROR();
                return -__LINE__;
            }
    	
            if( 0 == pid )
            {
                setsid();
                sock.close_listen();
                break;
            }
    		
            sock.close_sock();
        }

        start_server( &sock, opt.root );
    }
    else
    {
        sigset_t block;
        sigset_t empty;
        sigemptyset( &empty );
        sigemptyset( &block );
        sigaddset( &block, SIGALRM );
        sigprocmask( SIG_BLOCK, &block, NULL );
        regsignal( SIGALRM, sigclient );

        while( 1 )
        {
            // still have running child process
            // block waiting, release system resource
            if( -1 != waitpid( -1, NULL, 0 ) )
                continue;
            
            pid_t pid = fork();
            if( pid < 0 )
            {
                ERROR();
                return -1;
            }
            
            if( 0 == pid )
                break;
            
            // running once only
            if( 0 == opt.idle )
                _exit( 0 );
            
            alarm( opt.idle );
            sigsuspend( &empty );
        }

        tcpclient sock( opt.host, opt.port, opt.timeout );
        if( 0 != sock.error() )
        {
            ERROR();
            return -1;
        }

        hash* hs = NULL;
        if( !opt.overlap )
            hs = new hash( 1024 * 64 );
        
        start_client( &sock, opt.module, opt.tmp_path, hs );
    }

    return 0;
}