示例#1
0
/* record network stream as per spec in opt
 */
static int
record()
{
    int rsock = -1, destfd = -1, rc = 0, wtime_sec = 0;
    struct in_addr raddr;
    struct timeval rtv;
    struct dstream_ctx ds;
    ssize_t nmsgs = 0;
    ssize_t nrcv = -1, lrcv = -1, t_delta = 0;
    int64_t n_total = 0;
    ssize_t nwr = -1, lwr = -1;
    sig_atomic_t quit = 0;
    struct rdata_opt ropt;
    int oflags = 0;

    char* data = NULL;

    static const u_short RSOCK_TIMEOUT  = 5;
    extern const char CMD_UDP[];

    /* NOPs to eliminate warnings in lean version */
    (void)&t_delta; (void)&lrcv;
    t_delta = lrcv = lwr = 0; quit=0;

    check_fragments( NULL, 0, 0, 0, 0, g_flog );

    /* init */
    do {
        data = malloc( g_recopt.bufsize );
        if( NULL == data ) {
            mperror(g_flog, errno, "%s: cannot allocate [%ld] bytes",
                    __func__, (long)g_recopt.bufsize );
            rc = ERR_INTERNAL;
            break;
        }

        rc = subscribe( &rsock, &raddr );
        if( 0 != rc ) break;

        rtv.tv_sec = RSOCK_TIMEOUT;
        rtv.tv_usec = 0;

        rc = setsockopt( rsock, SOL_SOCKET, SO_RCVTIMEO, &rtv, sizeof(rtv) );
        if( -1 == rc ) {
            mperror(g_flog, errno, "%s: setsockopt - SO_RCVTIMEO",
                    __func__);
            rc = ERR_INTERNAL;
            break;
        }

        oflags = O_CREAT | O_TRUNC | O_WRONLY |
                 S_IRUSR | S_IWUSR | S_IRGRP | S_IROTH;
        # if defined(O_LARGEFILE)
            /* O_LARGEFILE is not defined under FreeBSD ??-7.1 */
            oflags |= O_LARGEFILE;
        # endif
        destfd = open( g_recopt.dstfile, oflags,
                (mode_t)(S_IRUSR | S_IWUSR | S_IRGRP | S_IROTH));
        if( -1 == destfd ) {
            mperror( g_flog, errno, "%s: cannot create destination file [%s]",
                     __func__, g_recopt.dstfile );
            rc = ERR_INTERNAL;
            break;
        }

        rc = calc_buf_settings( &nmsgs, NULL );
        if (0 != rc) return -1;

        if( nmsgs < (ssize_t)1 ) {
            (void) tmfprintf( g_flog, "Buffer for inbound data is too small [%ld] bytes; "
                    "the minimum size is [%ld] bytes\n",
                    (long)g_recopt.bufsize, (long)ETHERNET_MTU );
            rc = ERR_PARAM;
            break;
        }

        TRACE( (void)tmfprintf( g_flog, "Inbound buffer set to "
                        "[%d] messages\n", nmsgs ) );

        rc = init_dstream_ctx( &ds, CMD_UDP, NULL, nmsgs );
        if( 0 != rc ) return -1;

        (void) set_nice( g_recopt.nice_incr, g_flog );

        /* set up alarm to break main loop */
        if( 0 != g_recopt.end_time ) {
            wtime_sec = (int)difftime( g_recopt.end_time, time(NULL) );
            assert( wtime_sec >= 0 );

            (void) alarm( wtime_sec );

            (void)tmfprintf( g_flog, "Recording will end in [%d] seconds\n",
                    wtime_sec );
        }
    } while(0);

    /* record loop */
    ropt.max_frgs = g_recopt.rbuf_msgs;
    ropt.buf_tmout = -1;

    for( n_total = 0; (0 == rc) && !(quit = must_quit()); ) {
        nrcv = read_data( &ds, rsock, data, g_recopt.bufsize, &ropt );
        if( -1 == nrcv ) { rc = ERR_INTERNAL; break; }

        if( 0 == n_total ) {
            (void) tmfprintf( g_flog, "Recording to file=[%s] started.\n",
                    g_recopt.dstfile );
        }

        TRACE( check_fragments( "received new", g_recopt.bufsize,
                    lrcv, nrcv, t_delta, g_flog ) );
        lrcv = nrcv;

        if( nrcv > 0 ) {
            if( g_recopt.max_fsize &&
                ((n_total + nrcv) >= g_recopt.max_fsize) ) {
                break;
            }

            nwr = write_data( &ds, data, nrcv, destfd );
            if( -1 == nwr ) { rc = ERR_INTERNAL; break; }

            n_total += (size_t)nwr;
            /*
            TRACE( tmfprintf( g_flog, "Wrote [%ld] to file, total=[%ld]\n",
                        (long)nwr, (long)n_total ) );
            */

            TRACE( check_fragments( "wrote to file",
                    nrcv, lwr, nwr, t_delta, g_flog ) );
            lwr = nwr;
        }

        if( ds.flags & F_SCATTERED ) reset_pkt_registry( &ds );

    } /* record loop */

    (void) tmfprintf( g_flog, "Recording to file=[%s] stopped at filesize=[%lu] bytes\n",
                      g_recopt.dstfile, (u_long)n_total );

    /* CLEANUP
     */
    (void) alarm(0);

    TRACE( (void)tmfprintf( g_flog, "Exited record loop: wrote [%lu] bytes to file [%s], "
                    "rc=[%d], alarm=[%ld], quit=[%ld]\n",
                    (u_long)n_total, g_recopt.dstfile, rc, g_alarm, (long)quit ) );

    free_dstream_ctx( &ds );
    if( data ) free( data );

    close_mcast_listener( rsock, &raddr );
    if( destfd >= 0 ) (void) close( destfd );

    if( quit )
        TRACE( (void)tmfprintf( g_flog, "%s process must quit\n",
                        g_udpxrec_app ) );

    return rc;
}
示例#2
0
文件: udpxy.c 项目: avble/udpxy
/* relay traffic from source to destination socket
 *
 */
static int
relay_traffic( int ssockfd, int dsockfd, struct server_ctx* ctx,
               int dfilefd, const struct in_addr* mifaddr )
{
    volatile sig_atomic_t quit = 0;

    int rc = 0;
    ssize_t nmsgs = -1;
    ssize_t nrcv = 0, nsent = 0, nwr = 0,
            lrcv = 0, lsent = 0;
    char*  data = NULL;
    size_t data_len = g_uopt.rbuf_len;
    struct rdata_opt ropt;
    time_t pause_time = 0, rfr_tm = time(NULL);
    sigset_t ubset;

    const int ALLOW_PAUSES = get_flagval( "UDPXY_ALLOW_PAUSES", 0 );
    const ssize_t MAX_PAUSE_MSEC =
        get_sizeval( "UDPXY_PAUSE_MSEC", 1000);

    /* permissible variation in data-packet size */
    static const ssize_t t_delta = 0x20;

    struct dstream_ctx ds;

    static const int SET_PID = 1;
    struct tps_data tps;

    assert( ctx && mifaddr && MAX_PAUSE_MSEC > 0 );

    (void) sigemptyset (&ubset);
    sigaddset (&ubset, SIGINT);
    sigaddset (&ubset, SIGQUIT);
    sigaddset (&ubset, SIGTERM);

    /* restore the ability to receive *quit* signals */
    rc = sigprocmask (SIG_UNBLOCK, &ubset, NULL);
    if (0 != rc) {
        mperror (g_flog, errno, "%s: sigprocmask", __func__);
        return -1;
    }

    /* NOPs to eliminate warnings in lean version */
    (void)&lrcv; (void)&lsent; (void)&t_delta;

    check_fragments( NULL, 0, 0, 0, 0, g_flog );

    /* INIT
     */

    rc = calc_buf_settings( &nmsgs, NULL );
    if (0 != rc) return -1;

    TRACE( (void)tmfprintf( g_flog, "Data buffer will hold up to "
                        "[%d] messages\n", nmsgs ) );

    rc = init_dstream_ctx( &ds, ctx->cmd, g_uopt.srcfile, nmsgs );
    if( 0 != rc ) return -1;

    (void) set_nice( g_uopt.nice_incr, g_flog );

    do {
        if( NULL == g_uopt.srcfile ) {
            rc = set_timeouts( ssockfd, dsockfd,
                               ctx->rcv_tmout, 0,
                               ctx->snd_tmout, 0 );
            if( 0 != rc ) break;
        }

        if( dsockfd > 0 ) {
            rc = sync_dsockbuf_len( ssockfd, dsockfd );
            if( 0 != rc ) break;

            rc = send_http_response( dsockfd, 200, "OK" );
            if( 0 != rc ) break;

            /* timeshift: to detect PAUSE make destination
            * socket non-blocking, otherwise make it blocking
            * (since it might have been set unblocking earlier)
            */
            rc = set_nblock( dsockfd, (ALLOW_PAUSES ? 1 : 0) );
            if( 0 != rc ) break;
        }

        data = malloc(data_len);
        if( NULL == data ) {
            mperror( g_flog, errno, "%s: malloc", __func__ );
            break;
        }

        if( g_uopt.cl_tpstat )
            tpstat_init( &tps, SET_PID );
    } while(0);

    TRACE( (void)tmfprintf( g_flog, "Relaying traffic from socket[%d] "
            "to socket[%d], buffer size=[%d], Rmsgs=[%d], pauses=[%d]\n",
            ssockfd, dsockfd, data_len, g_uopt.rbuf_msgs, ALLOW_PAUSES) );

    /* RELAY LOOP
     */
    ropt.max_frgs = g_uopt.rbuf_msgs;
    ropt.buf_tmout = g_uopt.dhold_tmout;

    pause_time = 0;

    while( (0 == rc) && !(quit = must_quit()) ) {
        if( g_uopt.mcast_refresh > 0 ) {
            check_mcast_refresh( ssockfd, &rfr_tm, mifaddr );
        }

        nrcv = read_data( &ds, ssockfd, data, data_len, &ropt );
        if( -1 == nrcv ) break;

        TRACE( check_fragments( "received new", data_len,
                    lrcv, nrcv, t_delta, g_flog ) );
        lrcv = nrcv;

        if( dsockfd && (nrcv > 0) ) {
            nsent = write_data( &ds, data, nrcv, dsockfd );
            if( -1 == nsent ) break;

            if ( nsent < 0 ) {
                if ( !ALLOW_PAUSES ) break;
                if ( 0 != pause_detect( nsent, MAX_PAUSE_MSEC, &pause_time ) )
                    break;
            }

            TRACE( check_fragments("sent", nrcv,
                        lsent, nsent, t_delta, g_flog) );
            lsent = nsent;
        }

        if( (dfilefd > 0) && (nrcv > 0) ) {
            nwr = write_data( &ds, data, nrcv, dfilefd );
            if( -1 == nwr )
                break;
            TRACE( check_fragments( "wrote to file",
                    nrcv, lsent, nwr, t_delta, g_flog ) );
            lsent = nwr;
        }

        if( ds.flags & F_SCATTERED ) reset_pkt_registry( &ds );

        if( uf_TRUE == g_uopt.cl_tpstat )
            tpstat_update( ctx, &tps, nsent );

    } /* end of RELAY LOOP */

    /* CLEANUP
     */
    TRACE( (void)tmfprintf( g_flog, "Exited relay loop: received=[%ld], "
        "sent=[%ld], quit=[%ld]\n", (long)nrcv, (long)nsent, (long)quit ) );

    free_dstream_ctx( &ds );
    if( NULL != data ) free( data );

    if( 0 != (quit = must_quit()) ) {
        TRACE( (void)tmfprintf( g_flog, "Child process=[%d] must quit\n",
                    getpid()) );
    }

    return rc;
}