Skip to content

Commit

Permalink
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Add support for cookie jar to k6/ws
Browse files Browse the repository at this point in the history
This does require the ws module to get access to the http.CookieJar's
jar so we need to make it exported, but through the magic of `js` tags
we can make it not accessible from the js side.
mstoykov committed Oct 25, 2021
1 parent ee45fa7 commit 84547f3
Showing 4 changed files with 79 additions and 4 deletions.
6 changes: 3 additions & 3 deletions js/modules/k6/http/cookiejar.go
Original file line number Diff line number Diff line change
@@ -36,7 +36,7 @@ import (

// HTTPCookieJar is cookiejar.Jar wrapper to be used in js scripts
type HTTPCookieJar struct {
jar *cookiejar.Jar
Jar *cookiejar.Jar `js:"-"`
ctx *context.Context
}

@@ -55,7 +55,7 @@ func (j HTTPCookieJar) CookiesForURL(url string) map[string][]string {
panic(err)
}

cookies := j.jar.Cookies(u)
cookies := j.Jar.Cookies(u)
objs := make(map[string][]string, len(cookies))
for _, c := range cookies {
objs[c.Name] = append(objs[c.Name], c.Value)
@@ -101,6 +101,6 @@ func (j HTTPCookieJar) Set(url, name, value string, opts goja.Value) (bool, erro
}
}
}
j.jar.SetCookies(u, []*http.Cookie{&c})
j.Jar.SetCookies(u, []*http.Cookie{&c})
return true, nil
}
2 changes: 1 addition & 1 deletion js/modules/k6/http/request.go
Original file line number Diff line number Diff line change
@@ -329,7 +329,7 @@ func (h *HTTP) parseRequest(
}
switch v := jarV.Export().(type) {
case *HTTPCookieJar:
result.ActiveJar = v.jar
result.ActiveJar = v.Jar
}
case "compression":
algosString := strings.TrimSpace(params.Get(k).ToString().String())
14 changes: 14 additions & 0 deletions js/modules/k6/ws/ws.go
Original file line number Diff line number Diff line change
@@ -37,6 +37,7 @@ import (
"github.com/gorilla/websocket"

"go.k6.io/k6/js/common"
httpModule "go.k6.io/k6/js/modules/k6/http"
"go.k6.io/k6/lib"
"go.k6.io/k6/lib/metrics"
"go.k6.io/k6/stats"
@@ -114,6 +115,7 @@ func (*WS) Connect(ctx context.Context, url string, args ...goja.Value) (*WSHTTP
enableCompression := false

tags := state.CloneTags()
jar := state.CookieJar

// Parse the optional second argument (params)
if !goja.IsUndefined(paramsV) && !goja.IsNull(paramsV) {
@@ -144,6 +146,14 @@ func (*WS) Connect(ctx context.Context, url string, args ...goja.Value) (*WSHTTP
for _, key := range tagObj.Keys() {
tags[key] = tagObj.Get(key).String()
}
case "jar":
jarV := params.Get(k)
if goja.IsUndefined(jarV) || goja.IsNull(jarV) {
continue
}
if v, ok := jarV.Export().(*httpModule.HTTPCookieJar); ok {
jar = v.Jar
}
case "compression":
// deflate compression algorithm is supported - as defined in RFC7692
// compression here relies on the implementation in gorilla/websocket package, usage is
@@ -184,6 +194,10 @@ func (*WS) Connect(ctx context.Context, url string, args ...goja.Value) (*WSHTTP
Proxy: http.ProxyFromEnvironment,
TLSClientConfig: tlsConfig,
EnableCompression: enableCompression,
Jar: jar,
}
if jar == nil { // this is needed because of how interfaces work and that wsd.Jar is http.Cookiejar
wsd.Jar = nil
}

start := time.Now()
61 changes: 61 additions & 0 deletions js/modules/k6/ws/ws_test.go
Original file line number Diff line number Diff line change
@@ -26,6 +26,7 @@ import (
"fmt"
"io"
"net/http"
"net/http/cookiejar"
"net/http/httptest"
"strconv"
"testing"
@@ -38,9 +39,11 @@ import (
"gopkg.in/guregu/null.v3"

"go.k6.io/k6/js/common"
httpModule "go.k6.io/k6/js/modules/k6/http"
"go.k6.io/k6/lib"
"go.k6.io/k6/lib/metrics"
"go.k6.io/k6/lib/testutils/httpmultibin"

"go.k6.io/k6/stats"
)

@@ -1202,3 +1205,61 @@ func BenchmarkCompression(b *testing.B) {
}
})
}

func TestCookieJar(t *testing.T) {
t.Parallel()
ts := newTestState(t)
sr := ts.tb.Replacer.Replace

ts.tb.Mux.HandleFunc("/ws-echo-someheader", http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
responseHeaders := w.Header().Clone()
if sh, err := req.Cookie("someheader"); err == nil {
responseHeaders.Add("Echo-Someheader", sh.Value)
}

conn, err := (&websocket.Upgrader{}).Upgrade(w, req, responseHeaders)
if err != nil {
t.Fatalf("/ws-echo-someheader cannot upgrade request: %v", err)
}

err = conn.Close()
if err != nil {
t.Logf("error while closing connection in /ws-echo-someheader: %v", err)
}
}))
err := ts.rt.Set("http", common.Bind(ts.rt, httpModule.New().NewModuleInstancePerVU(), ts.ctxPtr))
require.NoError(t, err)
ts.state.CookieJar, _ = cookiejar.New(nil)

_, err = ts.rt.RunString(sr(`
var res = ws.connect("WSBIN_URL/ws-echo-someheader", function(socket){
socket.close()
})
var someheader = res.headers["Echo-Someheader"];
if (someheader !== undefined) {
throw new Error("someheader is echoed back by test server even though it doesn't exist");
}
http.cookieJar().set("HTTPBIN_URL/ws-echo-someheader", "someheader", "defaultjar")
res = ws.connect("WSBIN_URL/ws-echo-someheader", function(socket){
socket.close()
})
someheader = res.headers["Echo-Someheader"];
if (someheader != "defaultjar") {
throw new Error("someheader has wrong value "+ someheader + " instead of defaultjar");
}
var jar = new http.CookieJar();
jar.set("HTTPBIN_URL/ws-echo-someheader", "someheader", "customjar")
res = ws.connect("WSBIN_URL/ws-echo-someheader", {jar: jar}, function(socket){
socket.close()
})
someheader = res.headers["Echo-Someheader"];
if (someheader != "customjar") {
throw new Error("someheader has wrong value "+ someheader + " instead of customjar");
}
`))
assert.NoError(t, err)

assertSessionMetricsEmitted(t, stats.GetBufferedSamples(ts.samples), "", sr("WSBIN_URL/ws-echo-someheader"), statusProtocolSwitch, "")
}

0 comments on commit 84547f3

Please sign in to comment.