Add support for cookie jar to k6/ws

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
 }

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())

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()

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, "")
+}