Akkariin/wrk
1
0
Fork 0
mirror of https://github.com/wg/wrk synced 2025-01-08 23:32:54 +08:00
Code Issues Releases Wiki Activity

finalize and integrate SSL support

Browse Source
This commit is contained in:
Will 2013-06-22 15:48:48 +09:00
parent 87a5dae12c
commit 256b4756d3
8 changed files with 300 additions and 75 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
Makefile
View File

@ -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

55
src/main.h Normal file
View File

@ -0,0 +1,55 @@
#ifndef MAIN_H
#define MAIN_H
#include <ctype.h>
#include <errno.h>
#include <fcntl.h>
#include <getopt.h>
#include <math.h>
#include <netdb.h>
#include <netinet/in.h>
#include <netinet/tcp.h>
#include <stdarg.h>
#include <stdbool.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <signal.h>
#include <time.h>
#include <unistd.h>
#include <sys/socket.h>
#include <sys/time.h>
#include <sys/uio.h>
#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 */

32
src/net.c Normal file
View File

@ -0,0 +1,32 @@
// Copyright (C) 2013 - Will Glozer. All rights reserved.
#include <errno.h>
#include <unistd.h>
#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;
}

27
src/net.h Normal file
View File

@ -0,0 +1,27 @@
#ifndef NET_H
#define NET_H
#include <stdint.h>
#include <openssl/ssl.h>
#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 */

95
src/ssl.c Normal file
View File

@ -0,0 +1,95 @@
// Copyright (C) 2013 - Will Glozer. All rights reserved.
#include <pthread.h>
#include <openssl/evp.h>
#include <openssl/err.h>
#include <openssl/ssl.h>
#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;
}

13
src/ssl.h Normal file
View File

@ -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 */

118
src/wrk.c
View File

@ -1,32 +1,7 @@
// Copyright (C) 2012 - Will Glozer. All rights reserved.
#include "wrk.h"
#include <ctype.h>
#include <errno.h>
#include <fcntl.h>
#include <getopt.h>
#include <math.h>
#include <netdb.h>
#include <netinet/in.h>
#include <netinet/tcp.h>
#include <stdarg.h>
#include <stdbool.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <signal.h>
#include <time.h>
#include <unistd.h>
#include <sys/socket.h>
#include <sys/time.h>
#include <sys/uio.h>
#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;

31
src/wrk.h
View File

@ -6,12 +6,14 @@
#include <inttypes.h>
#include <sys/types.h>
#include <openssl/ssl.h>
#include <openssl/err.h>
#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 */