int main()

in wb/wb.c [3477:4098]


int main(int argc, const char * const argv[])
{
    int l;
    char tmp[1024];
    apr_status_t status;
    apr_getopt_t *opt;
    const char *opt_arg;
    char c;
#if OPENSSL_VERSION_NUMBER >= 0x10100000L
    int max_prot = TLS1_2_VERSION;
#ifndef OPENSSL_NO_SSL3
    int min_prot = SSL3_VERSION;
#else
    int min_prot = TLS1_VERSION;
#endif
#endif /* #if OPENSSL_VERSION_NUMBER >= 0x10100000L */
#ifdef USE_SSL
    AB_SSL_METHOD_CONST SSL_METHOD *meth = SSLv23_client_method();
#endif

    /* table defaults  */
    tablestring = "";
    trstring = "";
    tdstring = "bgcolor=white";
    cookie = "";
    auth = "";
    proxyhost = "";
    hdrs = "";

    setbuf(stdout, NULL);
    apr_app_initialize(&argc, &argv, NULL);
    atexit(apr_terminate);
    apr_pool_create(&cntxt, NULL);
    apr_pool_abort_set(abort_on_oom, cntxt);

#ifdef NOT_ASCII
    status = apr_xlate_open(&to_ascii, "ISO-8859-1", APR_DEFAULT_CHARSET, cntxt);
    if (status) {
        fprintf(stderr, "apr_xlate_open(to ASCII)->%d\n", status);
        exit(1);
    }
    status = apr_xlate_open(&from_ascii, APR_DEFAULT_CHARSET, "ISO-8859-1", cntxt);
    if (status) {
        fprintf(stderr, "apr_xlate_open(from ASCII)->%d\n", status);
        exit(1);
    }
    status = apr_base64init_ebcdic(to_ascii, from_ascii);
    if (status) {
        fprintf(stderr, "apr_base64init_ebcdic()->%d\n", status);
        exit(1);
    }
#endif

#ifdef _WAF_BENCH_ // loop for parse args
    g_new_header = xmalloc(g_header_len_MAX);
    g_sub_string = (char **)xcalloc(g_sub_string_num_MAX, sizeof(char *));
    if (argv[0][strlen(argv[0]) - 1] == '0')
        g_enable_ini = 1; // if progname = "wb0" then enable ini

    fprintf(stderr,"\033[1;33m"); // Yellow
    // parsed_args_rounds indicates the current round of parsing args
    int parsed_args_rounds = 0;
PARSE_ARGS:
    // 1:First CLI, 2:INI+CLI, no more than 3 rounds
    parsed_args_rounds ++;
#endif // _WAF_BENCH_ loop for parse args
    myhost = NULL; /* 0.0.0.0 or :: */

    apr_getopt_init(&opt, cntxt, argc, argv);
    while ((status = apr_getopt(opt, "n:c:t:s:b:T:p:u:v:lrkVhwiIx:y:z:C:H:P:A:g:X:de:SqB:m:"
#ifdef USE_SSL
            "Z:f:"
#endif
#ifdef _WAF_BENCH_ // adding more options 0-9,aFINRDUYWEGKQ 
            "Y:a:o:F:j:J:O:R:D:U:Y:W:E:G:Q:K012:3456789"
#endif // _WAF_BENCH_, adding more options 0-9,aFINRDUYWEGKQ 
            ,&c, &opt_arg)) == APR_SUCCESS) {
        switch (c) {
#ifdef _WAF_BENCH_ // more new features of waf-bench
            case '0': // enable read/write ini file for wb's options
                g_enable_ini = 1;
                break;
            case '1': // Don't append Host:localhost if it is set
                g_add_localhost = 0;
                break;
            case '2': // Don't append Connection:close if option is 0, 
                      // Append connection:close to those packets without connection attribution if option is 1,
                      // Append or replace connection attribution to close for any packets if option is 2
                g_add_connection_close = atoi(opt_arg);
                if (g_add_connection_close < 0 || g_add_connection_close > 2) {
                    err("Error option to -2\n");
                }
                break;
            case '3': // microsec-granularity in output-results
                // by default, it's not microsec-granularity
                // but you can change the default value and use -3 to toggle it
                g_us_granularity = !g_us_granularity; 
                break;
            case 'K': // save body or not?
                g_save_body = 1;
                break;
            case 'Q': // max number of packets in pkt file
                g_MAX_PKT_COUNT= atoi(opt_arg) ;
                if (g_MAX_PKT_COUNT < 0) 
                    err("Invalid max packet count\n");
                break;
            case 'R': // RPS number for rate limiting
                {
                int rps_scale = 1;
                if (opt_arg[strlen(opt_arg) - 1] == 'k' || opt_arg[strlen(opt_arg) - 1] == 'K')
                   rps_scale = 1000;
                else if (opt_arg[strlen(opt_arg) - 1] == 'm' || opt_arg[strlen(opt_arg) - 1] == 'M')
                   rps_scale = 1000000;
                else if (opt_arg[strlen(opt_arg) - 1] == 'g' || opt_arg[strlen(opt_arg) - 1] == 'G')
                   rps_scale = 1000000000;
                if (rps_scale == 1)                
                    g_RPS_NUMBER= atoi(opt_arg) ;
                else 
                    g_RPS_NUMBER= rps_scale * atof(opt_arg) ;
                if (g_MAX_PKT_COUNT < 0) 
                    err("Invalid number for RPS (Request per second)\n");
                }
                break;
            case 'G': // max log file size
                g_MAX_FILE_SIZE = atoi(opt_arg) * MB;
                if (g_MAX_FILE_SIZE < 0) 
                    err("Invalid max file size\n");
                break;
            case 'Y': 
            case 'U': // Add prefix("/prefix<seq#>/" to URL
                g_opt_prefix = xstrdup(opt_arg);
                break;
            case 'J': // subsistute sub_string with seq# inside header
                g_sub_string[g_sub_string_num++] = xstrdup(opt_arg);
                if (g_sub_string_num == g_sub_string_num_MAX) {
                    g_sub_string_num_MAX <<= 1;
                    char **sub_string_new;
                    sub_string_new = xcalloc(g_sub_string_num_MAX, sizeof (char *));
                    memcpy(sub_string_new,g_sub_string,g_sub_string_num*sizeof(char *));
                    free(g_sub_string);
                    g_sub_string = sub_string_new;
                }
                break;
            case 'o': // file to save received http messages
                g_save_filename = xstrdup(opt_arg);
                break;
            case 'a':
            case 'F': // a packet file to be sent
                g_pkt_filename = xstrdup(opt_arg);
                break;
            case 'E': // input ini file to be Executed
                opt_file_in = xstrdup(opt_arg);
                break;
            case 'O': // Output ini file to be saved
                opt_file_out = xstrdup(opt_arg);
                break;
            case 'j': // Interval (in secs) of progress report
                g_interval_print = atoi(opt_arg);
                if (g_interval_print < 0) 
                    err("Invalid print interval\n");
                break;
            case 'W': // stats size
                g_stats_window= atoi(opt_arg);
                if (g_stats_window <= 0) 
                    err("Invalid stats size\n");
                break;
            case '4': // keepalive_for_real_traffic
                g_keepalive_for_real_traffic  = 1;
                break;
            case '5': // print out additional progress info
                g_extended_progress = 1;
                break;
#endif // _WAF_BENCH_ end of new arguments processing

            case 'n':
                requests = atoi(opt_arg);
                if (requests <= 0) {
                    err("Invalid number of requests\n");
                }
#ifdef _WAF_BENCH_ // request is specified
                g_set_requests = 1;
#endif //_WAF_BENCH_, request is specified
                break;
            case 'k':
                keepalive = 1;
                g_add_connection_close  = 0;
                break;
            case 'q':
                heartbeatres = 0;
                break;
            case 'c':
                concurrency = atoi(opt_arg);
                break;
            case 'b':
                windowsize = atoi(opt_arg);
                break;
            case 'i':
                if (method != NO_METH)
                    err("Cannot mix HEAD with other methods\n");
                method = HEAD;
                break;
            case 'g':
                gnuplot = xstrdup(opt_arg);
                break;
            case 'd':
                percentile = 0;
                break;
            case 'e':
                csvperc = xstrdup(opt_arg);
                break;
            case 'S':
                confidence = 0;
                break;
            case 's':
                aprtimeout = apr_time_from_sec(atoi(opt_arg)); /* timeout value */
                break;
            case 'p':
                if (method != NO_METH)
                    err("Cannot mix POST with other methods\n");
#ifdef _WAF_BENCH_ // open those file after parsing all arguments
                g_post_filename = xstrdup(opt_arg); 
#else // original code goes here
                if (open_postfile(opt_arg) != APR_SUCCESS) {
                    exit(1);
                }
#endif //_WAF_BENCH_ // open post file after parsing all arguments
                
                method = POST;
                send_body = 1;
                break;
            case 'u':
                if (method != NO_METH)
                    err("Cannot mix PUT with other methods\n");
#ifdef _WAF_BENCH_ // open put file after parsing all arguments
                g_put_filename = xstrdup(opt_arg);  
#else // original code goes here
                if (open_postfile(opt_arg) != APR_SUCCESS) {
                    exit(1);
                }
#endif //_WAF_BENCH_ // open put file after parsing all arguments
                method = PUT;
                send_body = 1;
                break;
            case 'l':
                nolength = 1;
                break;
            case 'r':
                recverrok = 1;
                break;
            case 'v':
                verbosity = atoi(opt_arg);
                break;
            case 't':
                tlimit = atoi(opt_arg);
#ifdef _WAF_BENCH_ // set requests to max only when request is specified
                if (!g_set_requests)
#endif //_WAF_BENCH_, set requests to max only when request is specified
                requests = MAX_REQUESTS;    /* need to size data array on
                                             * something */
                break;
            case 'T':
                content_type = apr_pstrdup(cntxt, opt_arg);
                break;
            case 'C':
                cookie = apr_pstrcat(cntxt, "Cookie: ", opt_arg, "\r\n", NULL);
                break;
            case 'A':
                /*
                 * assume username passwd already to be in colon separated form.
                 * Ready to be uu-encoded.
                 */
                while (apr_isspace(*opt_arg))
                    opt_arg++;
                if (apr_base64_encode_len(strlen(opt_arg)) > sizeof(tmp)) {
                    err("Authentication credentials too long\n");
                }
                l = apr_base64_encode(tmp, opt_arg, strlen(opt_arg));
                tmp[l] = '\0';

                auth = apr_pstrcat(cntxt, auth, "Authorization: Basic ", tmp,
                                       "\r\n", NULL);
                break;
            case 'P':
                /*
                 * assume username passwd already to be in colon separated form.
                 */
                while (apr_isspace(*opt_arg))
                opt_arg++;
                if (apr_base64_encode_len(strlen(opt_arg)) > sizeof(tmp)) {
                    err("Proxy credentials too long\n");
                }
                l = apr_base64_encode(tmp, opt_arg, strlen(opt_arg));
                tmp[l] = '\0';

                auth = apr_pstrcat(cntxt, auth, "Proxy-Authorization: Basic ",
                                       tmp, "\r\n", NULL);
                break;
            case 'H':
                hdrs = apr_pstrcat(cntxt, hdrs, opt_arg, "\r\n", NULL);
                /*
                 * allow override of some of the common headers that ab adds
                 */
                if (strncasecmp(opt_arg, "Host:", 5) == 0) {
                    char *host;
                    apr_size_t len;
                    opt_arg += 5;
                    while (apr_isspace(*opt_arg))
                        opt_arg++;
                    len = strlen(opt_arg);
                    host = strdup(opt_arg);
                    while (len && apr_isspace(host[len-1]))
                        host[--len] = '\0';
                    opt_host = host;
                } else if (strncasecmp(opt_arg, "Accept:", 7) == 0) {
                    opt_accept = 1;
                    
#ifdef _WAF_BENCH_ // Connection:close header
                } else if (strncasecmp(opt_arg, "Connection:", 11) == 0) {
                    opt_connection = 1;
#endif // _WAF_BENCH_ // Connection:close header                
                } else if (strncasecmp(opt_arg, "User-Agent:", 11) == 0) {
                    opt_useragent = 1;
                }
                break;
            case 'w':
                use_html = 1;
                break;
                /*
                 * if any of the following three are used, turn on html output
                 * automatically
                 */
            case 'x':
                use_html = 1;
                tablestring = opt_arg;
                break;
            case 'X':
                {
                    char *p;
                    /*
                     * assume proxy-name[:port]
                     */
                    if ((p = strchr(opt_arg, ':'))) {
                        *p = '\0';
                        p++;
                        proxyport = atoi(p);
                    }
                    proxyhost = apr_pstrdup(cntxt, opt_arg);
                    isproxy = 1;
                }
                break;
            case 'y':
                use_html = 1;
                trstring = opt_arg;
                break;
            case 'z':
                use_html = 1;
                tdstring = opt_arg;
                break;
            case 'h':
                usage(argv[0]);
                break;
            case 'V':
                copyright();
                return 0;
            case 'B':
                myhost = apr_pstrdup(cntxt, opt_arg);
                break;
            case 'm':
                method = CUSTOM_METHOD;
                method_str[CUSTOM_METHOD] = strdup(opt_arg);
                break;
#ifdef USE_SSL
            case 'Z':
                ssl_cipher = strdup(opt_arg);
                break;
            case 'f':
#if OPENSSL_VERSION_NUMBER < 0x10100000L
                if (strncasecmp(opt_arg, "ALL", 3) == 0) {
                    meth = SSLv23_client_method();
#ifndef OPENSSL_NO_SSL2
                } else if (strncasecmp(opt_arg, "SSL2", 4) == 0) {
                    meth = SSLv2_client_method();
#ifdef HAVE_TLSEXT
                    tls_use_sni = 0;
#endif
#endif
#ifndef OPENSSL_NO_SSL3
                } else if (strncasecmp(opt_arg, "SSL3", 4) == 0) {
                    meth = SSLv3_client_method();
#ifdef HAVE_TLSEXT
                    tls_use_sni = 0;
#endif
#endif
#ifdef HAVE_TLSV1_X
                } else if (strncasecmp(opt_arg, "TLS1.1", 6) == 0) {
                    meth = TLSv1_1_client_method();
                } else if (strncasecmp(opt_arg, "TLS1.2", 6) == 0) {
                    meth = TLSv1_2_client_method();
#endif
                } else if (strncasecmp(opt_arg, "TLS1", 4) == 0) {
                    meth = TLSv1_client_method();
                }
#else /* #if OPENSSL_VERSION_NUMBER < 0x10100000L */
                meth = TLS_client_method();
                if (strncasecmp(opt_arg, "ALL", 3) == 0) {
                    max_prot = TLS1_2_VERSION;
#ifndef OPENSSL_NO_SSL3
                    min_prot = SSL3_VERSION;
#else
                    min_prot = TLS1_VERSION;
#endif
#ifndef OPENSSL_NO_SSL3
                } else if (strncasecmp(opt_arg, "SSL3", 4) == 0) {
                    max_prot = SSL3_VERSION;
                    min_prot = SSL3_VERSION;
#endif
                } else if (strncasecmp(opt_arg, "TLS1.1", 6) == 0) {
                    max_prot = TLS1_1_VERSION;
                    min_prot = TLS1_1_VERSION;
                } else if (strncasecmp(opt_arg, "TLS1.2", 6) == 0) {
                    max_prot = TLS1_2_VERSION;
                    min_prot = TLS1_2_VERSION;
                } else if (strncasecmp(opt_arg, "TLS1", 4) == 0) {
                    max_prot = TLS1_VERSION;
                    min_prot = TLS1_VERSION;
                }
#endif /* #if OPENSSL_VERSION_NUMBER < 0x10100000L */
                break;
#ifdef HAVE_TLSEXT
            case 'I':
                tls_use_sni = 0;
                break;
#endif
#endif
        }
    }

#ifdef _WAF_BENCH_ // open those files after parsing all arguments
    // enable ini if ini_file_in/out are set
    if ((g_enable_ini || opt_file_in) && parsed_args_rounds == 1) { 
        // read wb.ini, if ini file is valid, go back to parse ini
        // otherwise, go directly to process args
        opt_string = (char *)xmalloc(g_opt_string_len);
        g_argv_ini = (char **)xcalloc(g_opt_string_len, sizeof (char*));
        if (read_inifile(opt_file_in, opt_string) == APR_SUCCESS) {
            g_argc_ini = parse_opt_string(opt_string, g_argv_ini);
            if (g_argc_ini) { // parse the options again
                int i, j;
                int cli_url = 1;

                // append CLI options at the end of INI options
                if (opt->ind != argc - 1) { //CLI options don't have URL
                    cli_url = 0; 
                    // save the ini URL at the end of argv
                    g_argv_ini[argc+g_argc_ini-1] = g_argv_ini[g_argc_ini];
                }
                for (i = 1, j = g_argc_ini; i < argc; i ++, j ++)
                    g_argv_ini[j] = (char *)argv[i];
                    
                g_argv_ini[0] = (char *)argv[0]; // save the prog name
                argc += g_argc_ini - cli_url; 
                argv = (const char **)g_argv_ini;
                goto PARSE_ARGS;
            }
        }
    }
    
    // open those files after parsing all arguments
    if ((g_save_filename && open_file_for_write(g_save_filename, &g_save_file_fd) != APR_SUCCESS) // -R option 
        || (g_put_filename && open_postfile(g_put_filename) != APR_SUCCESS) // -u option 
        || (g_pkt_filename && open_pktfile(g_pkt_filename) != APR_SUCCESS) // -F option 
        || (g_post_filename && open_postfile(g_post_filename) != APR_SUCCESS)) // -p option
        exit(1);

    /*
     * if -n not specified, set requests to MAX
     * furthermore, if -t is not specified, set it to default
     * -n and -t might be specified at the same time
     * if -n is specified, we should not set -t to default
     * namely, -n 50000 != -n 50000 -t 10
     * because 50000 requests may take time > 10 seconds
     * so need to put "-t" processing logic inside "-n".
     */
    if (!g_set_requests) {
        requests = INT_MAX;
        if (tlimit == 0) {
            tlimit = DEFAULT_TEST_TIME;
        }
    } else {
        // we treat "-n" as the total times to be sent for the whole pkt file
        if (g_pkt_count)
            requests = requests * g_pkt_count;
    }

    // if we need save response to g_save_filename, only 1 connection can be used
    if (g_save_filename) {
        if (concurrency > 1) {
            fprintf(stderr, "WARNING: To save response to %s, only 1 connection can be used, so ignore -c %d!\n", 
                g_save_filename, concurrency);
            concurrency = 1;
        }
    }
    // disable previous req# based heartbeat, use time-based print interval
    heartbeatres = 0; 
    g_interval_print = g_interval_print * APR_USEC_PER_SEC; 

    // use milli-second as the timeout unit when finer granularity is chosen
    //if (g_us_granularity)
    //  aprtimeout = aprtimeout / 1000; // original value is "sec", now use "ms"
        
    if (verbosity >= 4) {
        printf("INFO: ARGS[%d] =", argc);
        int i = 0;
        for (i = 0; i < argc; i ++)
            printf(" %s",argv[i]);
        printf("\n");
    }
    // save the options to ini
    if (opt_file_out || g_enable_ini)
        write_inifile(opt_file_out, argc, (char **)argv);

    if (argc == 1) {
        usage(argv[0]);
    }
#endif // _WAF_BENCH_ // open those files after parsing all arguments

    if (opt->ind != argc - 1) {
        fprintf(stderr, "%s: wrong number of arguments\n", argv[0]);
        usage(argv[0]);
    }

    if (method == NO_METH) {
        method = GET;
    }

    if (parse_url(apr_pstrdup(cntxt, opt->argv[opt->ind++]))) {
        fprintf(stderr, "%s: invalid URL\n", argv[0]);
        usage(argv[0]);
    }

    if ((concurrency < 0) || (concurrency > MAX_CONCURRENCY)) {
        fprintf(stderr, "%s: Invalid Concurrency [Range 0..%d]\n",
                argv[0], MAX_CONCURRENCY);
        usage(argv[0]);
    }

#ifndef _WAF_BENCH_ // concurrency can be greater than requests
    if (concurrency > requests) {
        fprintf(stderr, "%s: Cannot use concurrency level greater than "
                "total number of requests\n", argv[0]);
        usage(argv[0]);
    }
#endif // _WAF_BENCH_ // concurrency can be greater than requests

    if ((heartbeatres) && (requests > 150)) {
        heartbeatres = requests / 10;   /* Print line every 10% of requests */
        if (heartbeatres < 100)
            heartbeatres = 100; /* but never more often than once every 100
                                 * connections. */
    }
    else
        heartbeatres = 0;

#ifdef USE_SSL
#ifdef RSAREF
    R_malloc_init();
#else
#if OPENSSL_VERSION_NUMBER < 0x10100000L
    CRYPTO_malloc_init();
#endif
#endif
    SSL_load_error_strings();
    SSL_library_init();
    bio_out=BIO_new_fp(stdout,BIO_NOCLOSE);
    bio_err=BIO_new_fp(stderr,BIO_NOCLOSE);

    if (!(ssl_ctx = SSL_CTX_new(meth))) {
        BIO_printf(bio_err, "Could not initialize SSL Context.\n");
        ERR_print_errors(bio_err);
        exit(1);
    }
    SSL_CTX_set_options(ssl_ctx, SSL_OP_ALL);
#if OPENSSL_VERSION_NUMBER >= 0x10100000L
    SSL_CTX_set_max_proto_version(ssl_ctx, max_prot);
    SSL_CTX_set_min_proto_version(ssl_ctx, min_prot);
#endif
#ifdef SSL_MODE_RELEASE_BUFFERS
    /* Keep memory usage as low as possible */
    SSL_CTX_set_mode (ssl_ctx, SSL_MODE_RELEASE_BUFFERS);
#endif
    if (ssl_cipher != NULL) {
        if (!SSL_CTX_set_cipher_list(ssl_ctx, ssl_cipher)) {
            fprintf(stderr, "error setting cipher list [%s]\n", ssl_cipher);
        ERR_print_errors_fp(stderr);
        exit(1);
    }
    }
    if (verbosity >= 3) {
        SSL_CTX_set_info_callback(ssl_ctx, ssl_state_cb);
    }
#endif
#ifdef SIGPIPE
    apr_signal(SIGPIPE, SIG_IGN);       /* Ignore writes to connections that
                                         * have been closed at the other end. */
#endif
#ifdef _WAF_BENCH_ // restore to no color mode 
    fprintf(stderr,"\033[0m"); // no color
#endif // _WAF_BENCH_ // restore to no color mode 
    copyright();
    test();
#ifdef _WAF_BENCH_ // close the file handles if they're still opened
    // close file for storing received http messages
    if (g_save_file_fd) {
        save_logfile(NULL, 0);
        apr_file_close(g_save_file_fd);
    }
#endif //_WAF_BENCH_ , close the file handles if they're still opened

    apr_pool_destroy(cntxt);

    return 0;
}