From 84547f37eb5d61865496f154d9fa3494b618d16a Mon Sep 17 00:00:00 2001 From: Mihail Stoykov Date: Thu, 21 Oct 2021 17:49:30 +0300 Subject: [PATCH] Add support for cookie jar to k6/ws 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. --- js/modules/k6/http/cookiejar.go | 6 ++-- js/modules/k6/http/request.go | 2 +- js/modules/k6/ws/ws.go | 14 ++++++++ js/modules/k6/ws/ws_test.go | 61 +++++++++++++++++++++++++++++++++ 4 files changed, 79 insertions(+), 4 deletions(-) diff --git a/js/modules/k6/http/cookiejar.go b/js/modules/k6/http/cookiejar.go index a0838d09e1f..9722a11c0df 100644 --- a/js/modules/k6/http/cookiejar.go +++ b/js/modules/k6/http/cookiejar.go @@ -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 } diff --git a/js/modules/k6/http/request.go b/js/modules/k6/http/request.go index 0457a30cf77..5317220fc41 100644 --- a/js/modules/k6/http/request.go +++ b/js/modules/k6/http/request.go @@ -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()) diff --git a/js/modules/k6/ws/ws.go b/js/modules/k6/ws/ws.go index d2bc6a6e6ff..9832b24cae8 100644 --- a/js/modules/k6/ws/ws.go +++ b/js/modules/k6/ws/ws.go @@ -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() diff --git a/js/modules/k6/ws/ws_test.go b/js/modules/k6/ws/ws_test.go index 8ee437f5ede..f6444266e09 100644 --- a/js/modules/k6/ws/ws_test.go +++ b/js/modules/k6/ws/ws_test.go @@ -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, "") +}