From 256b4756d3c0ba5c8fb00921733198e6be9fb8d6 Mon Sep 17 00:00:00 2001 From: Will Date: Sat, 22 Jun 2013 15:48:48 +0900 Subject: [PATCH] finalize and integrate SSL support --- Makefile | 4 +- src/main.h | 55 +++++++++++++++++++++++++ src/net.c | 32 +++++++++++++++ src/net.h | 27 ++++++++++++ src/ssl.c | 95 ++++++++++++++++++++++++++++++++++++++++++ src/ssl.h | 13 ++++++ src/wrk.c | 118 ++++++++++++++++++++++++++++++++--------------------- src/wrk.h | 31 +++----------- 8 files changed, 300 insertions(+), 75 deletions(-) create mode 100644 src/main.h create mode 100644 src/net.c create mode 100644 src/net.h create mode 100644 src/ssl.c create mode 100644 src/ssl.h diff --git a/Makefile b/Makefile index 19066f9..3f1ff9e 100644 --- a/Makefile +++ b/Makefile @@ -1,5 +1,5 @@ CFLAGS := -std=c99 -Wall -O2 -D_REENTRANT -LIBS := -lpthread -lm +LIBS := -lpthread -lm -lcrypto -lssl TARGET := $(shell uname -s | tr [A-Z] [a-z] 2>/dev/null || echo unknown) @@ -8,7 +8,7 @@ ifeq ($(TARGET), sunos) LIBS += -lsocket endif -SRC := wrk.c aprintf.c stats.c units.c ae.c zmalloc.c http_parser.c tinymt64.c +SRC := wrk.c net.c ssl.c aprintf.c stats.c units.c ae.c zmalloc.c http_parser.c tinymt64.c BIN := wrk ODIR := obj diff --git a/src/main.h b/src/main.h new file mode 100644 index 0000000..1ad8ce1 --- /dev/null +++ b/src/main.h @@ -0,0 +1,55 @@ +#ifndef MAIN_H +#define MAIN_H + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "ssl.h" +#include "aprintf.h" +#include "stats.h" +#include "units.h" +#include "zmalloc.h" + +struct config; + +static void *thread_main(void *); +static int connect_socket(thread *, connection *); +static int reconnect_socket(thread *, connection *); + +static int calibrate(aeEventLoop *, long long, void *); +static int sample_rate(aeEventLoop *, long long, void *); +static int check_timeouts(aeEventLoop *, long long, void *); + +static void socket_connected(aeEventLoop *, int, void *, int); +static void socket_writeable(aeEventLoop *, int, void *, int); +static void socket_readable(aeEventLoop *, int, void *, int); +static int request_complete(http_parser *); + +static uint64_t time_us(); + +static char *extract_url_part(char *, struct http_parser_url *, enum http_parser_url_fields); +static char *format_request(char *, char *, char *, char **); + +static int parse_args(struct config *, char **, char **, int, char **); +static void print_stats_header(); +static void print_stats(char *, stats *, char *(*)(long double)); +static void print_stats_latency(stats *); + +#endif /* MAIN_H */ diff --git a/src/net.c b/src/net.c new file mode 100644 index 0000000..6c73153 --- /dev/null +++ b/src/net.c @@ -0,0 +1,32 @@ +// Copyright (C) 2013 - Will Glozer. All rights reserved. + +#include +#include + +#include "net.h" + +status sock_connect(connection *c) { + return OK; +} + +status sock_close(connection *c) { + return OK; +} + +status sock_read(connection *c, size_t *n) { + ssize_t r = read(c->fd, c->buf, sizeof(c->buf)); + *n = (size_t) r; + return r > 0 ? OK : ERROR; +} + +status sock_write(connection *c, char *buf, size_t len, size_t *n) { + ssize_t r; + if ((r = write(c->fd, buf, len)) == -1) { + switch (errno) { + case EAGAIN: return RETRY; + default: return ERROR; + } + } + *n = (size_t) r; + return OK; +} diff --git a/src/net.h b/src/net.h new file mode 100644 index 0000000..dce5715 --- /dev/null +++ b/src/net.h @@ -0,0 +1,27 @@ +#ifndef NET_H +#define NET_H + +#include +#include + +#include "wrk.h" + +typedef enum { + OK, + ERROR, + RETRY +} status; + +struct sock { + status (*connect)(connection *); + status ( *close)(connection *); + status ( *read)(connection *, size_t *); + status ( *write)(connection *, char *, size_t, size_t *); +}; + +status sock_connect(connection *); +status sock_close(connection *); +status sock_read(connection *, size_t *); +status sock_write(connection *, char *, size_t, size_t *); + +#endif /* NET_H */ diff --git a/src/ssl.c b/src/ssl.c new file mode 100644 index 0000000..42d43aa --- /dev/null +++ b/src/ssl.c @@ -0,0 +1,95 @@ +// Copyright (C) 2013 - Will Glozer. All rights reserved. + +#include + +#include +#include +#include + +#include "ssl.h" + +static pthread_mutex_t *locks; + +static void ssl_lock(int mode, int n, const char *file, int line) { + pthread_mutex_t *lock = &locks[n]; + if (mode & CRYPTO_LOCK) { + pthread_mutex_lock(lock); + } else { + pthread_mutex_unlock(lock); + } +} + +static unsigned long ssl_id() { + return (unsigned long) pthread_self(); +} + +SSL_CTX *ssl_init() { + SSL_CTX *ctx = NULL; + + SSL_load_error_strings(); + SSL_library_init(); + OpenSSL_add_all_algorithms(); + + if ((locks = calloc(CRYPTO_num_locks(), sizeof(pthread_mutex_t)))) { + for (int i = 0; i < CRYPTO_num_locks(); i++) { + pthread_mutex_init(&locks[i], NULL); + } + + CRYPTO_set_locking_callback(ssl_lock); + CRYPTO_set_id_callback(ssl_id); + + if ((ctx = SSL_CTX_new(TLSv1_client_method()))) { + SSL_CTX_set_verify(ctx, SSL_VERIFY_NONE, NULL); + SSL_CTX_set_verify_depth(ctx, 0); + SSL_CTX_set_mode(ctx, SSL_MODE_AUTO_RETRY); + SSL_CTX_set_session_cache_mode(ctx, SSL_SESS_CACHE_CLIENT); + } + } + + return ctx; +} + +status ssl_connect(connection *c) { + int r; + SSL_set_fd(c->ssl, c->fd); + if ((r = SSL_connect(c->ssl)) != 1) { + switch (SSL_get_error(c->ssl, r)) { + case SSL_ERROR_WANT_READ: return RETRY; + case SSL_ERROR_WANT_WRITE: return RETRY; + default: return ERROR; + } + } + return OK; +} + +status ssl_close(connection *c) { + SSL_shutdown(c->ssl); + SSL_clear(c->ssl); + return OK; +} + +status ssl_read(connection *c, size_t *n) { + int r; + if ((r = SSL_read(c->ssl, c->buf, sizeof(c->buf))) <= 0) { + switch (SSL_get_error(c->ssl, r)) { + case SSL_ERROR_WANT_READ: return RETRY; + case SSL_ERROR_WANT_WRITE: return RETRY; + default: return ERROR; + } + } + *n = (size_t) r; + return OK; +} + +status ssl_write(connection *c, char *buf, size_t len, size_t *n) { + int r; + if ((r = SSL_write(c->ssl, buf, len)) <= 0) { + switch (SSL_get_error(c->ssl, r)) { + case SSL_ERROR_WANT_READ: return RETRY; + case SSL_ERROR_WANT_WRITE: return RETRY; + default: return ERROR; + } + } + *n = (size_t) r; + return OK; +} diff --git a/src/ssl.h b/src/ssl.h new file mode 100644 index 0000000..6eb13c0 --- /dev/null +++ b/src/ssl.h @@ -0,0 +1,13 @@ +#ifndef SSL_H +#define SSL_H + +#include "net.h" + +SSL_CTX *ssl_init(); + +status ssl_connect(connection *); +status ssl_close(connection *); +status ssl_read(connection *, size_t *); +status ssl_write(connection *, char *, size_t, size_t *); + +#endif /* SSL_H */ diff --git a/src/wrk.c b/src/wrk.c index 75da007..f7f5fb0 100644 --- a/src/wrk.c +++ b/src/wrk.c @@ -1,32 +1,7 @@ // Copyright (C) 2012 - Will Glozer. All rights reserved. #include "wrk.h" - -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include - -#include "aprintf.h" -#include "stats.h" -#include "units.h" -#include "zmalloc.h" -#include "tinymt64.h" +#include "main.h" static struct config { struct addrinfo addr; @@ -35,6 +10,7 @@ static struct config { uint64_t duration; uint64_t timeout; bool latency; + SSL_CTX *ctx; } cfg; static struct { @@ -50,6 +26,13 @@ static struct { pthread_mutex_t mutex; } statistics; +static struct sock sock = { + .connect = sock_connect, + .close = sock_close, + .read = sock_read, + .write = sock_write +}; + static const struct http_parser_settings parser_settings = { .on_message_complete = request_complete }; @@ -96,10 +79,11 @@ int main(int argc, char **argv) { exit(1); } - char *host = extract_url_part(url, &parser_url, UF_HOST); - char *port = extract_url_part(url, &parser_url, UF_PORT); - char *service = port ? port : extract_url_part(url, &parser_url, UF_SCHEMA); - char *path = "/"; + char *schema = extract_url_part(url, &parser_url, UF_SCHEMA); + char *host = extract_url_part(url, &parser_url, UF_HOST); + char *port = extract_url_part(url, &parser_url, UF_PORT); + char *service = port ? port : schema; + char *path = "/"; if (parser_url.field_set & (1 << UF_PATH)) { path = &url[parser_url.field_data[UF_PATH].off]; @@ -130,6 +114,18 @@ int main(int argc, char **argv) { exit(1); } + if (!strncmp("https", schema, 5)) { + if ((cfg.ctx = ssl_init()) == NULL) { + fprintf(stderr, "unable to initialize SSL\n"); + ERR_print_errors_fp(stderr); + exit(1); + } + sock.connect = ssl_connect; + sock.close = ssl_close; + sock.read = ssl_read; + sock.write = ssl_write; + } + signal(SIGPIPE, SIG_IGN); signal(SIGINT, SIG_IGN); cfg.addr = *addr; @@ -227,6 +223,7 @@ void *thread_main(void *arg) { for (uint64_t i = 0; i < thread->connections; i++, c++) { c->thread = thread; + c->ssl = cfg.ctx ? SSL_new(cfg.ctx) : NULL; connect_socket(thread, c); } @@ -268,21 +265,13 @@ static int connect_socket(thread *thread, connection *c) { flags = 1; setsockopt(fd, IPPROTO_TCP, TCP_NODELAY, &flags, sizeof(flags)); - if (aeCreateFileEvent(loop, fd, AE_WRITABLE, socket_writeable, c) != AE_OK) { - goto error; + flags = AE_READABLE | AE_WRITABLE; + if (aeCreateFileEvent(loop, fd, flags, socket_connected, c) == AE_OK) { + c->parser.data = c; + c->fd = fd; + return fd; } - if (aeCreateFileEvent(loop, fd, AE_READABLE, socket_readable, c) != AE_OK) { - goto error; - } - - http_parser_init(&c->parser, HTTP_RESPONSE); - c->parser.data = c; - c->fd = fd; - c->written = 0; - - return fd; - error: thread->errors.connect++; close(fd); @@ -291,6 +280,7 @@ static int connect_socket(thread *thread, connection *c) { static int reconnect_socket(thread *thread, connection *c) { aeDeleteFileEvent(thread->loop, c->fd, AE_WRITABLE | AE_READABLE); + sock.close(c); close(c->fd); return connect_socket(thread, c); } @@ -389,12 +379,40 @@ static int check_timeouts(aeEventLoop *loop, long long id, void *data) { return TIMEOUT_INTERVAL_MS; } +static void socket_connected(aeEventLoop *loop, int fd, void *data, int mask) { + connection *c = data; + + switch (sock.connect(c)) { + case OK: break; + case ERROR: goto error; + case RETRY: return; + } + + http_parser_init(&c->parser, HTTP_RESPONSE); + c->written = 0; + + aeCreateFileEvent(c->thread->loop, c->fd, AE_READABLE, socket_readable, c); + aeCreateFileEvent(c->thread->loop, c->fd, AE_WRITABLE, socket_writeable, c); + + return; + + error: + c->thread->errors.connect++; + reconnect_socket(c->thread, c); + +} + static void socket_writeable(aeEventLoop *loop, int fd, void *data, int mask) { connection *c = data; size_t len = req.size - c->written; - ssize_t n; + size_t n; + + switch (sock.write(c, req.buf + c->written, len, &n)) { + case OK: break; + case ERROR: goto error; + case RETRY: return; + } - if ((n = write(fd, req.buf + c->written, len)) < 0) goto error; if (!c->written) c->start = time_us(); c->written += n; @@ -410,11 +428,17 @@ static void socket_writeable(aeEventLoop *loop, int fd, void *data, int mask) { reconnect_socket(c->thread, c); } + static void socket_readable(aeEventLoop *loop, int fd, void *data, int mask) { connection *c = data; - ssize_t n; + size_t n; + + switch (sock.read(c, &n)) { + case OK: break; + case ERROR: goto error; + case RETRY: return; + } - if ((n = read(fd, c->buf, sizeof(c->buf))) <= 0) goto error; if (http_parser_execute(&c->parser, &parser_settings, c->buf, n) != n) goto error; c->thread->bytes += n; diff --git a/src/wrk.h b/src/wrk.h index 845805c..df0be4d 100644 --- a/src/wrk.h +++ b/src/wrk.h @@ -6,12 +6,14 @@ #include #include +#include +#include + #include "stats.h" #include "ae.h" #include "http_parser.h" -#include "tinymt64.h" -#define VERSION "2.1.0" +#define VERSION "2.2.0" #define RECVBUF 8192 #define SAMPLES 100000000 @@ -49,33 +51,10 @@ typedef struct connection { thread *thread; http_parser parser; int fd; + SSL *ssl; uint64_t start; size_t written; char buf[RECVBUF]; } connection; -struct config; - -static void *thread_main(void *); -static int connect_socket(thread *, connection *); -static int reconnect_socket(thread *, connection *); - -static int calibrate(aeEventLoop *, long long, void *); -static int sample_rate(aeEventLoop *, long long, void *); -static int check_timeouts(aeEventLoop *, long long, void *); - -static void socket_writeable(aeEventLoop *, int, void *, int); -static void socket_readable(aeEventLoop *, int, void *, int); -static int request_complete(http_parser *); - -static uint64_t time_us(); - -static char *extract_url_part(char *, struct http_parser_url *, enum http_parser_url_fields); -static char *format_request(char *, char *, char *, char **); - -static int parse_args(struct config *, char **, char **, int, char **); -static void print_stats_header(); -static void print_stats(char *, stats *, char *(*)(long double)); -static void print_stats_latency(stats *); - #endif /* WRK_H */