libevent multithreaded ssl test code

/*
 * Simple test to verify SSL behavior with libevent. This code is
 * intended to demonstrate a specific behavior- it is not meant for
 * any other use. (i.e. it egregiously uses unprotected globals,
 * albeit in a safe manner, for interthread communication)
 *
 * This code operates in the following approximate phases:
 * 1. Initialize libraries.
 * 2. Synchronously connect to server.
 * 3. Once connected, spawn specified number of transmit threads.
 * 4. Forever: Run threads.
 *
 * Each thread loops:
 * 1. Send one line of data to the server.
 * 2. Sleep 0.1s
 * 3. Repeat.
 *
 * NOTE: Some portions of this code may fall under others' copyrights-
 *    Namely the code for initializing SSL multithreaded was taken
 *    directly from the book referenced in the code below.
 *
 * NOTE: This code explicitly runs SSL in a very insecure
 *    manner. Don't copy this code and assume you have any security at
 *    all.
 *
 * NOTE: If your server isn't running, you might get a SIGPIPE, as I
 *    do not disable SIGPIPE in this code.
 *
 * To build: (I have used Centos 6.3 and OSX 10.8.3)
 *
 * gcc -Wall -Werror test.c -o test -lssl -lcrypto -levent -levent_openssl -levent_pthreads
 *
 * If you have openssl/libevent somewhere, you may need to add
 * -I/your_include_dir and -L/your_library_dir paths to the command
 * line above. I do this for OSX, but not for my linux builds.
 *
 * To make a server that can listen, using openssl:
 *
 * 1. Create a self-signed certificate for your server, run the
 *    following command, and fill out the form: (Generates files
 *    test.key and test.crt)
 *
 * openssl req -x509 -nodes -days 365 -newkey rsa:2048 -keyout test.key -out test.crt
 *
 * 2. Run openssl as a simple server, echoing received traffic to
 *    stdout (listens on port 5000, which is the port this code
 *    connects to):
 *
 * openssl s_server -cert test.crt -key test.key -accept 5000
 *
 * To run:
 *
 * ./test IP_ADDRESS NUMTHREADS
 *
 * This code purposely does NOT do name resolution. To run with one
 * thread, going to localhost:
 *
 * ./test 127.0.0.1 1
 *
 * etc.
 *
 * This code generally writes to stderr.
 */

/* The following protect calls to bufferevents with an external
 * mutex. Use/Don't use at your discretion. It should not be required,
 * but is here for debugging/analysis. */
#define TEST_USE_EXTERNAL_LOCKING

/* If you don't want to use SSL, you can comment this out. If you
 * don't use ssl, you will need to use a different server to connect
 * to rather than the openssl server. */
#define TEST_USE_SSL

/* Behavior is demonstrated with as few as 1 thread. */
#define MAX_THREADS 100

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/socket.h>
#include <sys/time.h>
#include <netinet/in.h>
#include <arpa/inet.h>
#include <pthread.h>
#include <unistd.h>

#include <openssl/ssl.h>

#include <event.h>
#include <event2/thread.h>
#include <event2/bufferevent_ssl.h>




/***************************************************
 *
 * OpenSSL thread initialization routines taken from book:
 *
 * Network Security with OpenSSL
 *
 * Viega, John; Messier, Matt; Chandra, Pravir (2009-02-09)
 */
#if defined(WIN32)
#define MUTEX_TYPE HANDLE
#define MUTEX_SETUP(x) (x) = CreateMutex(NULL, FALSE, NULL)
#define MUTEX_CLEANUP(x) CloseHandle(x)
#define MUTEX_LOCK(x) WaitForSingleObject((x), INFINITE)
#define MUTEX_UNLOCK(x) ReleaseMutex(x)
#define THREAD_ID GetCurrentThreadId( )
#else
/* _POSIX_THREADS is normally defined in unistd.h if pthreads are available
   on your platform. */
#define MUTEX_TYPE pthread_mutex_t
#define MUTEX_SETUP(x) pthread_mutex_init(&(x), NULL)
#define MUTEX_CLEANUP(x) pthread_mutex_destroy(&(x))
#define MUTEX_LOCK(x) pthread_mutex_lock(&(x))
#define MUTEX_UNLOCK(x) pthread_mutex_unlock(&(x))
#define THREAD_ID pthread_self( )
#endif
/* This array will store all of the mutexes available to OpenSSL. */
static MUTEX_TYPE *mutex_buf = NULL;
static void locking_function(int mode, int n, const char * file, int line)
{
    if (mode & CRYPTO_LOCK)
        MUTEX_LOCK(mutex_buf[n]);
    else
        MUTEX_UNLOCK(mutex_buf[n]);
}
static unsigned long id_function(void)
{
    return ((unsigned long)THREAD_ID);
}
struct CRYPTO_dynlock_value
{
    MUTEX_TYPE mutex;
};
static struct CRYPTO_dynlock_value * dyn_create_function(const char *file,
                                                         int line)
{
    struct CRYPTO_dynlock_value *value;
    value = (struct CRYPTO_dynlock_value *)malloc(sizeof(
                                                         struct CRYPTO_dynlock_value));
    if (!value)
        return NULL;
    MUTEX_SETUP(value->mutex);
    return value;
}
static void dyn_lock_function(int mode, struct CRYPTO_dynlock_value *l,
                              const char *file, int line)
{
    if (mode & CRYPTO_LOCK)
        MUTEX_LOCK(l->mutex);
    else
        MUTEX_UNLOCK(l->mutex);
}
static void dyn_destroy_function(struct CRYPTO_dynlock_value *l,
                                 const char *file, int line)
{
    MUTEX_CLEANUP(l->mutex);
    free(l);
}
int THREAD_setup(void)
{
    int i;
    mutex_buf = (MUTEX_TYPE *)malloc(CRYPTO_num_locks( ) * sizeof(MUTEX_TYPE));
    if (!mutex_buf)
        return 0;
    for (i = 0; i < CRYPTO_num_locks( ); i++)
        MUTEX_SETUP(mutex_buf[i]);
    CRYPTO_set_id_callback(id_function);
    CRYPTO_set_locking_callback(locking_function);
    /* The following three CRYPTO_... functions are the OpenSSL functions
       for registering the callbacks we implemented above */
    CRYPTO_set_dynlock_create_callback(dyn_create_function);
    CRYPTO_set_dynlock_lock_callback(dyn_lock_function);
    CRYPTO_set_dynlock_destroy_callback(dyn_destroy_function);
    return 1;
}
int THREAD_cleanup(void)
{
    int i;
    if (!mutex_buf)
        return 0;
    CRYPTO_set_id_callback(NULL);
    CRYPTO_set_locking_callback(NULL);
    CRYPTO_set_dynlock_create_callback(NULL);
    CRYPTO_set_dynlock_lock_callback(NULL);
    CRYPTO_set_dynlock_destroy_callback(NULL);
    for (i = 0; i < CRYPTO_num_locks( ); i++)
        MUTEX_CLEANUP(mutex_buf[i]);
    free(mutex_buf);
    mutex_buf = NULL;
    return 1;
}
/*
 * OpenSSL thread initialization code complete.
 *
 ***************************************************/


#ifdef TEST_USE_EXTERNAL_LOCKING
static MUTEX_TYPE global_lock;
#endif



/* Yes, lovely globals */
struct event_base *ev_base;
struct bufferevent *ev_buf;
struct event *ev_forever;
int num_threads = 0;
int thread_numbers[MAX_THREADS];
pthread_t thread_ids[MAX_THREADS];


/*
 * Simple thread. Just writes one line of data to output, once every 0.1 seconds.
 *
 * arg points at an int containing our (user-defined) thread number.
 */
static void *run_thread(void *arg)
{
    int id = *((int *) arg);
    char str[100];
    int len;
    int count = 0;

    fprintf(stderr, "Thread %d starting, ID = %ld, %016lx \n", id, id_function(), id_function());
    while(1) {
        len = snprintf(str, sizeof(str), "Thread %d sending data, iteration %d\n", id, count);
        if (len >= sizeof(str)) {
            fprintf(stderr, "Erk, Line got long.\n");
            exit(1);
        }
        fprintf(stderr, "%s", str);

#ifdef TEST_USE_EXTERNAL_LOCKING
        MUTEX_LOCK(global_lock);
#endif
        if (bufferevent_write(ev_buf, str, len)) {
#ifdef TEST_USE_EXTERNAL_LOCKING
            MUTEX_UNLOCK(global_lock);
#endif
            fprintf(stderr, "Error writing, thread %d, at iteration %d\n", id, count);
            exit(1);
        }
#ifdef TEST_USE_EXTERNAL_LOCKING
        MUTEX_UNLOCK(global_lock);
#endif
        usleep(100000);
        count++;
    }
}

int test_ssl_init(void)
{
    THREAD_setup();
    SSL_load_error_strings();
    SSL_library_init();

    printf("SSL version: %s\n", SSLeay_version(SSLEAY_VERSION));
    return 0;
}

int test_libevent_init(void)
{
    printf("Libevent version: %s\n", event_get_version());

    if (evthread_use_pthreads() != 0) {
        fprintf(stderr, "Error: libevent not using threads\n");
        exit(1);
    }

    return 0;
}

/*
 * A simple timer event, to show that we have not deadlocked.
 */
static void timer_cb(evutil_socket_t sock, short flags, void *cookie)
{
    fprintf(stderr, "Time passes\n");
}

/*
 * For completeness: A read callback.
 */
static void read_cb(struct bufferevent *ev_buf, void *cookie)
{
    fprintf(stderr, "Received read callback\n");
}

/*
 * For completeness: A write callback.
 */
static void write_cb(struct bufferevent *ev_buf, void *cookie)
{
    fprintf(stderr, "Received write callback\n");
}

/*
 * All the events we see. On connect, we spawn our threads.
 */
static void event_cb(struct bufferevent *bev, short style, void *cookie)
{
    fprintf(stderr, "Received event :");
    if (style & BEV_EVENT_READING) fprintf (stderr, " BEV_EVENT_READING");
    if (style & BEV_EVENT_WRITING) fprintf (stderr, " BEV_EVENT_WRITING");
    if (style & BEV_EVENT_EOF) fprintf (stderr, " BEV_EVENT_EOF");
    if (style & BEV_EVENT_ERROR) fprintf (stderr, " BEV_EVENT_ERROR");
    if (style & BEV_EVENT_TIMEOUT) fprintf (stderr, " BEV_EVENT_TIMEOUT");
    if (style & BEV_EVENT_CONNECTED) fprintf (stderr, " BEV_EVENT_CONNECTED");
    fprintf(stderr, "\n");

    if (style & BEV_EVENT_CONNECTED) {
        /* Start threads. */
        int i;
        for (i = 0; i < num_threads; i++) {
            fprintf(stderr, "Creating thread %d\n", i);
            thread_numbers[i] = i;
            if (pthread_create(&(thread_ids[i]), NULL, run_thread, &(thread_numbers[i]))) {
                fprintf(stderr, "Err: Thread creation failed for thread %d\n", i);
                exit(1);
            }
        }
    }

    if (style & BEV_EVENT_ERROR) {
        fprintf(stderr, "Is your server running?\n");
    }
}

int main (int argc, char *argv[])
{
    SSL_CTX *ssl_ctx;
    SSL *ssl;
    struct sockaddr_in sin;
    struct timeval tv;
    int res;

#ifdef TEST_USE_EXTERNAL_LOCKING
    MUTEX_SETUP(global_lock);
#endif

    if (argc != 3) {
        fprintf(stderr, "Use %s <IP-ADDRESS> <NUM-THREADS>\n", argv[0]);
        exit(1);
    }

    num_threads = atoi(argv[2]);
    if ((num_threads > MAX_THREADS) || (num_threads < 0)) {
        fprintf(stderr, "Invalid number of threads. 0 to %d, please.\n", MAX_THREADS);
        exit(1);
    }

    res = test_ssl_init();
    if (res) {
        fprintf(stderr, "Err: test_ssl_init returned %d\n", res);
        exit(1);
    }

    res = test_libevent_init();
    if (res) {
        fprintf(stderr, "Err: test_libevent_init returned %d\n", res);
        exit(1);
    }

    ssl_ctx = SSL_CTX_new(SSLv3_client_method());
    if (!ssl_ctx) {
        fprintf(stderr, "Err: Could not create SSL context\n");
        exit(1);
    }

    /* Don't verify certificates. THIS IS EXTRAORDINARILY BAD for
     * security. It is done here ONLY TO TEST. */
    SSL_CTX_set_verify(ssl_ctx, SSL_VERIFY_NONE, NULL);

    ssl = SSL_new(ssl_ctx);
    if (!ssl) {
        fprintf(stderr, "Err: Could not create SSL\n");
        exit(1);
    }

    ev_base = event_base_new();
    if (!ev_base) {
        fprintf(stderr, "Err: event_base_new failed\n");
        exit(1);
    }

    /* Create a timer that runs forever. Handy to see timing, handy to
     * have something run forever in dispatch. */
    ev_forever = event_new(ev_base, -1, EV_PERSIST, timer_cb, NULL);
    if (!ev_forever) {
        fprintf(stderr, "Err: Could not create forever event.\n");
        exit(1);
    }
    tv.tv_sec = 1;
    tv.tv_usec = 0;
    if (event_add(ev_forever, &tv)) {
        fprintf(stderr, "Err: Could not add forever event\n");
        exit(1);
    }

#ifdef TEST_USE_SSL
    ev_buf = bufferevent_openssl_socket_new(ev_base,
                                            -1,
                                            ssl,
                                            BUFFEREVENT_SSL_CONNECTING,
                                            BEV_OPT_CLOSE_ON_FREE | BEV_OPT_THREADSAFE);
#else
    ev_buf = bufferevent_socket_new(ev_base,
                                    -1,
                                    BEV_OPT_CLOSE_ON_FREE | BEV_OPT_THREADSAFE);
#endif
    if (!ev_buf) {
        fprintf(stderr, "Err: Could not create ev_buf\n");
        exit(1);
    }

    bufferevent_setcb(ev_buf, read_cb, write_cb, event_cb, NULL);

    /* These should not be needed, but I threw them in here for good
     * luck. */
    evbuffer_enable_locking(bufferevent_get_input(ev_buf), NULL);
    evbuffer_enable_locking(bufferevent_get_output(ev_buf), NULL);

    memset(&sin, 0, sizeof(sin));
#ifndef __linux__
    sin.sin_len = sizeof(sin);
#endif
    sin.sin_family = AF_INET;
    sin.sin_port = htons(5000);
    sin.sin_addr.s_addr = inet_addr(argv[1]);

    res = bufferevent_socket_connect(ev_buf,
                                     (struct sockaddr *) &sin,
                                     sizeof(sin));
    if (res != 0) {
        fprintf(stderr,
                "Err: bufferevent_socket_connect_hostname returned %d\n",
                res);
        exit(1);
    }

    bufferevent_enable(ev_buf, EV_READ|EV_WRITE);

    /* Note: No locking really required before this point, as no other
     * threads are running yet. Threads are started once connect event
     * returns. */
    fprintf(stderr, "Starting dispatch loop\n");

    event_base_dispatch(ev_base);

    fprintf(stderr, "Leaving dispatch loop\n");

    exit(0);
}