Add support for cookie jar to k6/ws

This was all easy until I actually had to get the jar from the
httpCookiejar and unfortunately I don't know of an alternative to the
helper and it definitely will be *way* more code

js/modules/k6/http/cookiejar.go

@@ -40,6 +40,12 @@ type HTTPCookieJar struct {
 	ctx *context.Context
 }
 
+// GetJarFromHTTPCookieJar is a helper function for modules outside of http to get the jar from inside the cookie jar
+// TODO figure out how to do this through only goja
+func GetJarFromHTTPCookieJar(jar *HTTPCookieJar) *cookiejar.Jar {
+	return jar.jar // binks
+}
+
 func newCookieJar(ctxPtr *context.Context) *HTTPCookieJar {
 	jar, err := cookiejar.New(nil)
 	if err != nil {

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 = httpModule.GetJarFromHTTPCookieJar(v)
+				}
 			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, "")
+}