/*
 * sslbench.c - TLS handshake benchmark with configurable group/curve
 * Similar to openssl s_time but with explicit group selection for key agreement
 */

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <signal.h>
#include <sys/time.h>

#include <openssl/ssl.h>
#include <openssl/err.h>
#include <openssl/bn.h>

static volatile int running = 1;

static void alarm_handler(int sig)
{
    running = 0;
}

static double time_now(void)
{
    struct timeval tv;
    gettimeofday(&tv, NULL);
    return tv.tv_sec + tv.tv_usec / 1e6;
}

static int next_proto_cb(SSL *s, unsigned char **out, unsigned char *outlen,
                         const unsigned char *in, unsigned int inlen, void *arg)
{
    /* No ALPN selection needed for benchmark */
    return SSL_TLSEXT_ERR_OK;
}

static void usage(const char *prog)
{
    fprintf(stderr,
        "Usage: %s -connect host:port [options]\n"
        "Options:\n"
        "  -connect host:port    Target server (required)\n"
        "  -groups list          Colon-separated group names (e.g., 'X25519:P-256:P-384')\n"
        "  -curves list          Alias for -groups (TLS 1.2 legacy naming)\n"
        "  -cipher list          Cipher list string\n"
        "  -ciphersuites list    TLS 1.3 ciphersuites\n"
        "  -time seconds         Benchmark duration (default: 30)\n"
        "  -new                  Use new session for each handshake\n"
        "  -reuse                Reuse session (default)\n"
        "  -verify depth         Enable certificate verification\n"
        "  -CAfile file          CA certificate file\n"
        "  -cert file            Client certificate\n"
        "  -key file             Client private key\n"
        "  -tls1_3               Force TLS 1.3\n"
        "  -tls1_2               Force TLS 1.2\n"
        "  -nbio                 Use non-blocking IO\n"
        "  -www page             HTTP GET request path\n"
        "  -bytes n              Read n bytes after handshake\n"
        , prog);
    exit(1);
}

int main(int argc, char *argv[])
{
    const char *connect_str = NULL;
    const char *groups = NULL;
    const char *cipher = NULL;
    const char *ciphersuites = NULL;
    const char *cafile = NULL;
    const char *cert = NULL;
    const char *key = NULL;
    const char *www = NULL;
    int verify_depth = -1;
    int time_limit = 30;
    int new_session = 0;
    int max_version = 0;
    int nbio = 0;
    int read_bytes = 0;

    SSL_CTX *ctx = NULL;
    SSL *ssl = NULL;
    BIO *bio = NULL;
    int i;

    /* Parse arguments */
    for (i = 1; i < argc; i++) {
        if (strcmp(argv[i], "-connect") == 0 && i + 1 < argc)
            connect_str = argv[++i];
        else if ((strcmp(argv[i], "-groups") == 0 || strcmp(argv[i], "-curves") == 0) && i + 1 < argc)
            groups = argv[++i];
        else if (strcmp(argv[i], "-cipher") == 0 && i + 1 < argc)
            cipher = argv[++i];
        else if (strcmp(argv[i], "-ciphersuites") == 0 && i + 1 < argc)
            ciphersuites = argv[++i];
        else if (strcmp(argv[i], "-time") == 0 && i + 1 < argc)
            time_limit = atoi(argv[++i]);
        else if (strcmp(argv[i], "-new") == 0)
            new_session = 1;
        else if (strcmp(argv[i], "-reuse") == 0)
            new_session = 0;
        else if (strcmp(argv[i], "-verify") == 0 && i + 1 < argc)
            verify_depth = atoi(argv[++i]);
        else if (strcmp(argv[i], "-CAfile") == 0 && i + 1 < argc)
            cafile = argv[++i];
        else if (strcmp(argv[i], "-cert") == 0 && i + 1 < argc)
            cert = argv[++i];
        else if (strcmp(argv[i], "-key") == 0 && i + 1 < argc)
            key = argv[++i];
        else if (strcmp(argv[i], "-tls1_3") == 0)
            max_version = TLS1_3_VERSION;
        else if (strcmp(argv[i], "-tls1_2") == 0)
            max_version = TLS1_2_VERSION;
        else if (strcmp(argv[i], "-nbio") == 0)
            nbio = 1;
        else if (strcmp(argv[i], "-www") == 0 && i + 1 < argc)
            www = argv[++i];
        else if (strcmp(argv[i], "-bytes") == 0 && i + 1 < argc)
            read_bytes = atoi(argv[++i]);
        else
            usage(argv[0]);
    }

    if (!connect_str) {
        fprintf(stderr, "Error: -connect required\n");
        usage(argv[0]);
    }

    /* Initialize OpenSSL */
    SSL_load_error_strings();
    SSL_library_init();

    /* Create context */
    const SSL_METHOD *meth = TLS_client_method();
    ctx = SSL_CTX_new(meth);
    if (!ctx) {
        ERR_print_errors_fp(stderr);
        return 1;
    }

    /* Configure group/curve selection for key agreement
     * This is the primary feature beyond standard s_time */
    if (groups) {
        /*
         * SSL_CTX_set1_groups_list() configures supported groups for key exchange.
         * For TLS 1.3: controls (EC)DHE group selection (X25519, X448, P-256, P-384, P-521, FFDHE)
         * For TLS 1.2: controls ECDHE curve selection via supported_groups extension
         *
         * Order matters: first match with server preference wins.
         * OpenSSL 3.0+ unified API; older versions use SSL_CTX_set1_curves_list()
         */
#if OPENSSL_VERSION_MAJOR >= 3
        if (!SSL_CTX_set1_groups_list(ctx, groups)) {
#else
        if (!SSL_CTX_set1_curves_list(ctx, groups)) {
#endif
            fprintf(stderr, "Error: Invalid group/curve list: %s\n", groups);
            ERR_print_errors_fp(stderr);
            goto err;
        }
        printf("Configured groups: %s\n", groups);
    }

    /* TLS version constraints */
    if (max_version == TLS1_3_VERSION) {
        SSL_CTX_set_min_proto_version(ctx, TLS1_3_VERSION);
        SSL_CTX_set_max_proto_version(ctx, TLS1_3_VERSION);
    } else if (max_version == TLS1_2_VERSION) {
        SSL_CTX_set_max_proto_version(ctx, TLS1_2_VERSION);
    }

    /* Cipher configuration */
    if (cipher && !SSL_CTX_set_cipher_list(ctx, cipher)) {
        fprintf(stderr, "Error: Invalid cipher list\n");
        goto err;
    }

#if TLS1_3_VERSION
    if (ciphersuites && !SSL_CTX_set_ciphersuites(ctx, ciphersuites)) {
        fprintf(stderr, "Error: Invalid ciphersuites\n");
        goto err;
    }
#endif

    /* Certificate verification */
    if (verify_depth >= 0) {
        SSL_CTX_set_verify(ctx, SSL_VERIFY_PEER, NULL);
        SSL_CTX_set_verify_depth(ctx, verify_depth);
    } else {
        SSL_CTX_set_verify(ctx, SSL_VERIFY_NONE, NULL);
    }

    if (cafile && !SSL_CTX_load_verify_locations(ctx, cafile, NULL)) {
        fprintf(stderr, "Error loading CA file\n");
        goto err;
    }

    if (cert && !SSL_CTX_use_certificate_file(ctx, cert, SSL_FILETYPE_PEM)) {
        fprintf(stderr, "Error loading certificate\n");
        goto err;
    }

    if (key && !SSL_CTX_use_PrivateKey_file(ctx, key, SSL_FILETYPE_PEM)) {
        fprintf(stderr, "Error loading private key\n");
        goto err;
    }

    /* Session caching for reuse mode */
    if (!new_session)
        SSL_CTX_set_session_cache_mode(ctx, SSL_SESS_CACHE_CLIENT);

    /* Build connect BIO */
    bio = BIO_new_ssl_connect(ctx);
    if (!bio) {
        ERR_print_errors_fp(stderr);
        goto err;
    }

    BIO_get_ssl(bio, &ssl);
    if (!ssl) {
        fprintf(stderr, "Error getting SSL from BIO\n");
        goto err;
    }

    SSL_set_mode(ssl, SSL_MODE_AUTO_RETRY);
    if (nbio)
        BIO_set_nbio(bio, 1);

    BIO_set_conn_hostname(bio, connect_str);

    /* Optional: configure ALPN if needed */
    SSL_set_alpn_protos(ssl, (const unsigned char *)"\x08http/1.1", 9);

    /* Benchmark */
    printf("Collecting connection statistics for %d seconds\n", time_limit);
    printf("Connecting to %s\n", connect_str);
    if (groups)
        printf("Groups: %s\n", groups);
    printf("...\n");

    signal(SIGALRM, alarm_handler);
    alarm(time_limit);

    long handshake_count = 0;
    long error_count = 0;
    double total_handshake_time = 0.0;
    double start_time = time_now();

    SSL_SESSION *sess = NULL;

    while (running) {
        double t0 = time_now();

        /* Establish connection */
        if (BIO_do_connect(bio) <= 0) {
            if (!running) /* interrupted by alarm */
                break;
            error_count++;
            ERR_clear_error();
            
            /* Reset connection for retry */
            BIO_reset(bio);
            BIO_get_ssl(bio, &ssl);
            continue;
        }

        /* Complete handshake */
        if (SSL_get_verify_result(ssl) != X509_V_OK && verify_depth >= 0) {
            /* verification failed but we continue for benchmark */
        }

        double t1 = time_now();
        total_handshake_time += (t1 - t0);
        handshake_count++;

        /* Session handling */
        if (!new_session) {
            if (!sess)
                sess = SSL_get1_session(ssl);
            else
                SSL_set_session(ssl, sess);
        }

        /* Optional HTTP request */
        if (www) {
            char req[256];
            snprintf(req, sizeof(req), "GET %s HTTP/1.0\r\nHost: %s\r\n\r\n", 
                     www, connect_str);
            BIO_write(bio, req, strlen(req));
        }

        /* Optional read */
        if (read_bytes > 0 || www) {
            char buf[4096];
            int toread = read_bytes > 0 ? read_bytes : 8192;
            int total = 0;
            while (total < toread) {
                int n = BIO_read(bio, buf, sizeof(buf));
                if (n <= 0) break;
                total += n;
            }
        }

        /* Prepare for next iteration */
        if (new_session || !sess) {
            BIO_ssl_shutdown(bio);
            BIO_reset(bio);
            BIO_get_ssl(bio, &ssl);
            if (groups) {
                /* Re-apply groups after reset if needed */
#if OPENSSL_VERSION_MAJOR >= 3
                SSL_set1_groups_list(ssl, groups);
#else
                SSL_set1_curves_list(ssl, groups);
#endif
            }
        } else {
            /* Session reuse: just reconnect with same SSL object settings */
            BIO_ssl_shutdown(bio);
            BIO_reset(bio);
            BIO_get_ssl(bio, &ssl);
            SSL_set_session(ssl, sess);
        }
    }

    double elapsed = time_now() - start_time;

    /* Results */
    printf("\n");
    printf("%ld connections in %.2fs; %.2f connections/user sec, %.3f ms avg\n",
           handshake_count, elapsed,
           handshake_count / elapsed,
           handshake_count > 0 ? (total_handshake_time / handshake_count) * 1000.0 : 0);

    if (error_count > 0)
        printf("%ld errors, %.2f error rate\n", error_count, 
               (double)error_count / (handshake_count + error_count));

    /* Report negotiated parameters from last successful connection */
    if (ssl && SSL_is_init_finished(ssl)) {
        printf("\nLast connection:\n");
        printf("  Version: %s\n", SSL_get_version(ssl));
        printf("  Cipher: %s\n", SSL_get_cipher(ssl));
        
        /* Get negotiated group for key agreement */
        int nid = SSL_get_shared_curve(ssl, 0);
        if (nid > 0) {
            const char *gname = OBJ_nid2sn(nid);
            printf("  Key agreement group: %s (NID %d)\n", 
                   gname ? gname : "unknown", nid);
        }
#if OPENSSL_VERSION_MAJOR >= 3
        /* TLS 1.3 specific: get group used for key exchange */
        else {
            nid = SSL_get_negotiated_group(ssl);
            if (nid > 0) {
                const char *gname = OBJ_nid2sn(nid);
                printf("  Negotiated group: %s\n", 
                       gname ? gname : "unknown");
            }
        }
#endif
    }

    /* Cleanup */
    SSL_SESSION_free(sess);
    BIO_free_all(bio);
    SSL_CTX_free(ctx);
    EVP_cleanup();
    ERR_free_strings();

    return 0;

err:
    SSL_CTX_free(ctx);
    ERR_print_errors_fp(stderr);
    return 1;
}