From 2a4b64033a11a30bbb98d70e9c7e52c4161c60e3 Mon Sep 17 00:00:00 2001
From: Will <will@glozer.net>
Date: Sun, 10 Nov 2013 16:16:53 +0900
Subject: [PATCH] fix Host header with port or IPv6 address

---
 scripts/wrk.lua | 24 ++++++++++++++++++++----
 src/wrk.h       |  2 +-
 2 files changed, 21 insertions(+), 5 deletions(-)

diff --git a/scripts/wrk.lua b/scripts/wrk.lua
index 03cbced..1bfa571 100644
--- a/scripts/wrk.lua
+++ b/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.request()  return req         end
+function wrk.init(args)
+   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
diff --git a/src/wrk.h b/src/wrk.h
index f113570..785c779 100644
--- a/src/wrk.h
+++ b/src/wrk.h
@@ -14,7 +14,7 @@
 #include "script.h"
 #include "http_parser.h"
 
-#define VERSION  "3.0.1"
+#define VERSION  "3.0.3"
 #define RECVBUF  8192
 #define SAMPLES  100000000