fix Host header with port or IPv6 address

Makefile

@@ -45,7 +45,7 @@ $(ODIR): $(LDIR)/libluajit.a
 
 $(ODIR)/bytecode.o: scripts/wrk.lua
 	@echo LUAJIT $<
-	@$(SHELL) -c 'cd $(LDIR) && ./luajit -b $(PWD)/$< $(PWD)/$@'
+	@$(SHELL) -c 'cd $(LDIR) && ./luajit -b $(CURDIR)/$< $(CURDIR)/$@'
 
 $(ODIR)/%.o : %.c
 	@echo CC $<

scripts/wrk.lua

@@ -9,14 +9,16 @@ local wrk = {
 }
 
 function wrk.format(method, path, headers, body)
-   local host    = wrk.host
    local method  = method  or wrk.method
    local path    = path    or wrk.path
    local headers = headers or wrk.headers
    local body    = body    or wrk.body
    local s       = {}
 
-   headers["Host"]           = port and (host .. ":" .. port) or host
+   if not headers["Host"] then
+      headers["Host"] = wrk.headers["Host"]
+   end
+
    headers["Content-Length"] = body and string.len(body)
 
    s[1] = string.format("%s %s HTTP/1.1", method, path)
@@ -30,8 +32,22 @@ function wrk.format(method, path, headers, body)
    return table.concat(s, "\r\n")
 end
 
-function wrk.init(args) req = wrk.format() end
+function wrk.init(args)
-function wrk.request()  return req         end
+   if not wrk.headers["Host"] then
+      local host = wrk.host
+      local port = wrk.port
+
+      host = host:find(":") and ("[" .. host .. "]")  or host
+      host = port           and (host .. ":" .. port) or host
+
+      wrk.headers["Host"] = host
+   end
+   req = wrk.format()
+end
+
+function wrk.request()
+   return req
+end
 
 init     = wrk.init
 request  = wrk.request

src/net.c

@@ -2,6 +2,7 @@
 
 #include <errno.h>
 #include <unistd.h>
+#include <sys/ioctl.h>
 
 #include "net.h"
 
@@ -30,3 +31,9 @@ status sock_write(connection *c, char *buf, size_t len, size_t *n) {
     *n = (size_t) r;
     return OK;
 }
+
+size_t sock_readable(connection *c) {
+    int n, rc;
+    rc = ioctl(c->fd, FIONREAD, &n);
+    return rc == -1 ? 0 : n;
+}

src/net.h

@@ -13,15 +13,17 @@ typedef enum {
 } status;
 
 struct sock {
-    status (*connect)(connection *);
+    status ( *connect)(connection *);
-    status (  *close)(connection *);
+    status (   *close)(connection *);
-    status (   *read)(connection *, size_t *);
+    status (    *read)(connection *, size_t *);
-    status (  *write)(connection *, char *, size_t, size_t *);
+    status (   *write)(connection *, char *, size_t, size_t *);
+    size_t (*readable)(connection *);
 };
 
 status sock_connect(connection *);
 status sock_close(connection *);
 status sock_read(connection *, size_t *);
 status sock_write(connection *, char *, size_t, size_t *);
+size_t sock_readable(connection *);
 
 #endif /* NET_H */

src/ssl.c

@@ -93,3 +93,7 @@ status ssl_write(connection *c, char *buf, size_t len, size_t *n) {
     *n = (size_t) r;
     return OK;
 }
+
+size_t ssl_readable(connection *c) {
+    return SSL_pending(c->ssl);
+}

src/ssl.h

@@ -9,5 +9,6 @@ status ssl_connect(connection *);
 status ssl_close(connection *);
 status ssl_read(connection *, size_t *);
 status ssl_write(connection *, char *, size_t, size_t *);
+size_t ssl_readable(connection *);
 
 #endif /* SSL_H */

src/wrk.c

@@ -22,10 +22,11 @@ static struct {
 } statistics;
 
 static struct sock sock = {
-    .connect = sock_connect,
+    .connect  = sock_connect,
-    .close   = sock_close,
+    .close    = sock_close,
-    .read    = sock_read,
+    .read     = sock_read,
-    .write   = sock_write
+    .write    = sock_write,
+    .readable = sock_readable
 };
 
 static struct http_parser_settings parser_settings = {
@@ -114,10 +115,11 @@ int main(int argc, char **argv) {
             ERR_print_errors_fp(stderr);
             exit(1);
         }
-        sock.connect = ssl_connect;
+        sock.connect  = ssl_connect;
-        sock.close   = ssl_close;
+        sock.close    = ssl_close;
-        sock.read    = ssl_read;
+        sock.read     = ssl_read;
-        sock.write   = ssl_write;
+        sock.write    = ssl_write;
+        sock.readable = ssl_readable;
     }
 
     signal(SIGPIPE, SIG_IGN);
@@ -134,6 +136,7 @@ int main(int argc, char **argv) {
 
     for (uint64_t i = 0; i < cfg.threads; i++) {
         thread *t = &threads[i];
+        t->loop        = aeCreateEventLoop(10 + cfg.connections * 3);
         t->connections = connections;
         t->stop_at     = stop_at;
 
@@ -151,9 +154,9 @@ int main(int argc, char **argv) {
             }
         }
 
-        if (pthread_create(&t->thread, NULL, &thread_main, t)) {
+        if (!t->loop || pthread_create(&t->thread, NULL, &thread_main, t)) {
             char *msg = strerror(errno);
-            fprintf(stderr, "unable to create thread %"PRIu64" %s\n", i, msg);
+            fprintf(stderr, "unable to create thread %"PRIu64": %s\n", i, msg);
             exit(2);
         }
     }
@@ -225,10 +228,9 @@ int main(int argc, char **argv) {
 
 void *thread_main(void *arg) {
     thread *thread = arg;
+    aeEventLoop *loop = thread->loop;
 
-    aeEventLoop *loop = aeCreateEventLoop(10 + cfg.connections * 3);
+    thread->cs = zmalloc(thread->connections * sizeof(connection));
-    thread->cs   = zmalloc(thread->connections * sizeof(connection));
-    thread->loop = loop;
     tinymt64_init(&thread->rand, time_us());
     thread->latency = stats_alloc(100000);
 
@@ -496,14 +498,16 @@ static void socket_readable(aeEventLoop *loop, int fd, void *data, int mask) {
     connection *c = data;
     size_t n;
 
-    switch (sock.read(c, &n)) {
+    do {
-        case OK:    break;
+        switch (sock.read(c, &n)) {
-        case ERROR: goto error;
+            case OK:    break;
-        case RETRY: return;
+            case ERROR: goto error;
-    }
+            case RETRY: return;
+        }
 
-    if (http_parser_execute(&c->parser, &parser_settings, c->buf, n) != n) goto error;
+        if (http_parser_execute(&c->parser, &parser_settings, c->buf, n) != n) goto error;
-    c->thread->bytes += n;
+        c->thread->bytes += n;
+    } while (n == RECVBUF && sock.readable(c) > 0);
 
     return;
 

src/wrk.h

@@ -14,7 +14,7 @@
 #include "script.h"
 #include "http_parser.h"
 
-#define VERSION  "3.0.0"
+#define VERSION  "3.0.3"
 #define RECVBUF  8192
 #define SAMPLES  100000000