wrk response headers doesn't handle duplicate headers

[erickt]

wrk doesn't seem to handle duplicated response headers like Set-Cookie correctly. This makes it difficult to test services that require multiple cookies, like a session cookie and a CSRF cookie in order to login to a service.

[wg]

Thanks for the issue report @erickt! Currently wrk overwrites each entry in the headers table with the last header seen. Should be easy to fix, although there will be a performance cost.

[xianjimli]

+1 Is there a plan to support multiple cookies?

[njenia]

I have a hack to fix this (didn't have the time to fix the problem from the root): In script.c, function void script_response(), which processes response headers, in case it's a Set-Cookie header I changed the header name to Set-CookiX when X is a running number (Set-Cooki1, Set-Cooki2...), so now the cookie headers are all distinct in name. This is the modified loop at the beginning of the function:

    for (char *c = headers->buffer; c < headers->cursor; ) {
        if (strcmp(c, "Set-Cookie") == 0) { /* New line */
            cookie_count++; /* New line */
            c[9] = cookie_count + '0'; /* New line */
        } /* New line */
        c = buffer_pushlstring(L, c);
        c = buffer_pushlstring(L, c);
        lua_rawset(L, -3);
    }

You need to recreate the wrk executable by running make. Of course you need some code to handle these headers and make them real cookies again, so I handled it in a lua script in the response hook:

response = function(status, headers, body)
    wrk.headers["Cookie"] = ''
    for key, value in pairs(headers) do
        if string.starts(key, "Set-Cooki") then
            wrk.headers["Cookie"] = wrk.headers["Cookie"] .. string.sub(value, 0, string.find(value, ";") - 1) .. ';'
        end
    end
    cookie = true
end

This concatenates the custom Set-CookiX headers into a standard Cookie header. This solution works great for my use case, but it's a dirty hack, a very dirty hack. I just didn't have time to change the underlying data structures that hold the headers and make them support duplicate names... Hope this helps.

[bayupermadi]

Thanks for the issue report @erickt! Currently wrk overwrites each entry in the headers table with the last header seen. Should be easy to fix, although there will be a performance cost.

hi @wg , is this plan already implemented? It would be nice if we have feature for duplicated header responds.

Thank you

[rsegretain]

Inspired by the workaround of @njenia , I further modified the function void script_response() in script.c so that the cookies are stored in a sub-table named Set-Cookie (side-bonus: there can be more than 9) :

void script_response(lua_State *L, int status, buffer *headers, buffer *body) {
    lua_getglobal(L, "response");
    lua_pushinteger(L, status);
    lua_newtable(L); // parent headers table

    int cookie_idx = 1;

    for (char *c = headers->buffer; c < headers->cursor; ) {
        if (strcmp(c, "Set-Cookie") == 0) {
            if (cookie_idx == 1) { // create the cookies table
                // create table
                buffer_pushlstring(L, c); // -2, the name of the new table "Set-Cookie"
                lua_newtable(L); // -1, the new table
                lua_settable(L, -3); // tell the parent table in -3 to save the new table in -1 under the name "Set-Cookie" in -2
            }
            
            // get the cookie table
            c = buffer_pushlstring(L, c);
            lua_gettable(L, -2);

            // save the cookie
            lua_pushnumber(L, cookie_idx++); // -2, push the cookie index
            c = buffer_pushlstring(L, c); // -1, push the cookie value
            lua_settable(L, -3); // save the cookie

            lua_pop(L, 1); // go back to the headers table
        }
        else {
            c = buffer_pushlstring(L, c);
            c = buffer_pushlstring(L, c);
            lua_rawset(L, -3);
        }
    }

    lua_pushlstring(L, body->buffer, body->cursor - body->buffer);
    lua_call(L, 3, 0);

    buffer_reset(headers);
    buffer_reset(body);
}

Hope it help someone !