From ee45fa74e36c693a965ba6fa652856eddf3f6cdb Mon Sep 17 00:00:00 2001 From: Mihail Stoykov Date: Mon, 25 Oct 2021 17:11:44 +0300 Subject: [PATCH 1/2] Add newTestState to ws tests --- js/modules/k6/ws/ws_test.go | 48 ++++++++++++++++++++++++------------- 1 file changed, 31 insertions(+), 17 deletions(-) diff --git a/js/modules/k6/ws/ws_test.go b/js/modules/k6/ws/ws_test.go index b6b04b31bff..8ee437f5ede 100644 --- a/js/modules/k6/ws/ws_test.go +++ b/js/modules/k6/ws/ws_test.go @@ -91,7 +91,15 @@ func assertMetricEmittedCount(t *testing.T, metricName string, sampleContainers assert.Equal(t, count, actualCount, "url %s emitted %s %d times, expected was %d times", url, metricName, actualCount, count) } -func newRuntime(t testing.TB) (*httpmultibin.HTTPMultiBin, chan stats.SampleContainer, *goja.Runtime) { +type testState struct { + ctxPtr *context.Context + rt *goja.Runtime + tb *httpmultibin.HTTPMultiBin + state *lib.State + samples chan stats.SampleContainer +} + +func newTestState(t testing.TB) testState { tb := httpmultibin.NewHTTPMultiBin(t) root, err := lib.NewGroup("", nil) @@ -125,7 +133,13 @@ func newRuntime(t testing.TB) (*httpmultibin.HTTPMultiBin, chan stats.SampleCont err = rt.Set("ws", common.Bind(rt, New(), ctx)) assert.NoError(t, err) - return tb, samples, rt + return testState{ + ctxPtr: ctx, + rt: rt, + tb: tb, + state: state, + samples: samples, + } } func TestSession(t *testing.T) { @@ -982,9 +996,9 @@ func TestCompression(t *testing.T) { t.Parallel() const text string = `Lorem ipsum dolor sit amet, consectetur adipiscing elit. Maecenas sed pharetra sapien. Nunc laoreet molestie ante ac gravida. Etiam interdum dui viverra posuere egestas. Pellentesque at dolor tristique, mattis turpis eget, commodo purus. Nunc orci aliquam.` - tb, samples, rt := newRuntime(t) - sr := tb.Replacer.Replace - tb.Mux.HandleFunc("/ws-compression", http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) { + ts := newTestState(t) + sr := ts.tb.Replacer.Replace + ts.tb.Mux.HandleFunc("/ws-compression", http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) { upgrader := websocket.Upgrader{ EnableCompression: true, ReadBufferSize: 1024, @@ -1010,7 +1024,7 @@ func TestCompression(t *testing.T) { } })) - _, err := rt.RunString(sr(` + _, err := ts.rt.RunString(sr(` // if client supports compression, it has to send the header // 'Sec-Websocket-Extensions:permessage-deflate; server_no_context_takeover; client_no_context_takeover' to server. // if compression is negotiated successfully, server will reply with header @@ -1035,7 +1049,7 @@ func TestCompression(t *testing.T) { `)) assert.NoError(t, err) - assertSessionMetricsEmitted(t, stats.GetBufferedSamples(samples), "", sr("WSBIN_URL/ws-compression"), statusProtocolSwitch, "") + assertSessionMetricsEmitted(t, stats.GetBufferedSamples(ts.samples), "", sr("WSBIN_URL/ws-compression"), statusProtocolSwitch, "") }) t.Run("params", func(t *testing.T) { @@ -1070,9 +1084,9 @@ func TestCompression(t *testing.T) { testCase := testCase t.Run(testCase.compression, func(t *testing.T) { t.Parallel() - tb, _, rt := newRuntime(t) - sr := tb.Replacer.Replace - tb.Mux.HandleFunc("/ws-compression-param", http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) { + ts := newTestState(t) + sr := ts.tb.Replacer.Replace + ts.tb.Mux.HandleFunc("/ws-compression-param", http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) { upgrader := websocket.Upgrader{ EnableCompression: true, ReadBufferSize: 1024, @@ -1092,7 +1106,7 @@ func TestCompression(t *testing.T) { } })) - _, err := rt.RunString(sr(` + _, err := ts.rt.RunString(sr(` var res = ws.connect("WSBIN_URL/ws-compression-param", {"compression":"` + testCase.compression + `"}, function(socket){ socket.close() }); @@ -1122,9 +1136,9 @@ func clearSamples(tb *httpmultibin.HTTPMultiBin, samples chan stats.SampleContai func BenchmarkCompression(b *testing.B) { const textMessage = 1 - tb, samples, rt := newRuntime(b) - sr := tb.Replacer.Replace - go clearSamples(tb, samples) + ts := newTestState(b) + sr := ts.tb.Replacer.Replace + go clearSamples(ts.tb, ts.samples) testCodes := []string{ sr(` @@ -1143,7 +1157,7 @@ func BenchmarkCompression(b *testing.B) { `), } - tb.Mux.HandleFunc("/ws-compression", http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) { + ts.tb.Mux.HandleFunc("/ws-compression", http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) { kbData := bytes.Repeat([]byte("0123456789"), 100) // upgrade connection, send the first (long) message, disconnect @@ -1175,14 +1189,14 @@ func BenchmarkCompression(b *testing.B) { b.ResetTimer() b.Run("compression-enabled", func(b *testing.B) { for i := 0; i < b.N; i++ { - if _, err := rt.RunString(testCodes[0]); err != nil { + if _, err := ts.rt.RunString(testCodes[0]); err != nil { b.Error(err) } } }) b.Run("compression-disabled", func(b *testing.B) { for i := 0; i < b.N; i++ { - if _, err := rt.RunString(testCodes[1]); err != nil { + if _, err := ts.rt.RunString(testCodes[1]); err != nil { b.Error(err) } } From dc36cb240c2086e1c73d6c762af5931b0a3f540f Mon Sep 17 00:00:00 2001 From: Mihail Stoykov Date: Thu, 21 Oct 2021 17:49:30 +0300 Subject: [PATCH 2/2] 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 | 7 ++-- 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, 80 insertions(+), 4 deletions(-) diff --git a/js/modules/k6/http/cookiejar.go b/js/modules/k6/http/cookiejar.go index a0838d09e1f..1ccebf38450 100644 --- a/js/modules/k6/http/cookiejar.go +++ b/js/modules/k6/http/cookiejar.go @@ -36,7 +36,8 @@ import ( // HTTPCookieJar is cookiejar.Jar wrapper to be used in js scripts type HTTPCookieJar struct { - jar *cookiejar.Jar + // js is to make it not be accessible from inside goja/js, the json is because it's used if we return it from setup + Jar *cookiejar.Jar `js:"-" json:"-"` ctx *context.Context } @@ -55,7 +56,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 +102,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, "") +}