Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Add cookie support to WebSockets #2052

Closed
wants to merge 4 commits into from
Closed

Add cookie support to WebSockets #2052

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
3dd5a6b
Add Cookie header support for WebSockets
lcd1232 Oct 4, 2020
c9ab77a
Remove usage of common.RunString
May 20, 2021
831813e
Initialize jar variable closer to where it's used
May 20, 2021
3778915
Fix linter issues
Jun 8, 2021
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
76 changes: 76 additions & 0 deletions js/modules/k6/ws/ws.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -28,6 +28,9 @@ import (
"io/ioutil"
"net"
"net/http"
"net/http/cookiejar"
netURL "net/url"
"reflect"
"strconv"
"strings"
"sync"
Expand Down Expand Up @@ -112,6 +115,16 @@ func (*WS) Connect(ctx context.Context, url string, args ...goja.Value) (*WSHTTP

tags := state.CloneTags()

u, err := netURL.Parse(url)
if err != nil {
return nil, err
}

cookiesMap := make(map[string]http.Cookie)
for _, cookie := range state.CookieJar.Cookies(u) {
cookiesMap[cookie.Name] = *cookie
}

// Parse the optional second argument (params)
if !goja.IsUndefined(paramsV) && !goja.IsNull(paramsV) {
params := paramsV.ToObject(rt)
Expand All @@ -130,6 +143,43 @@ func (*WS) Connect(ctx context.Context, url string, args ...goja.Value) (*WSHTTP
for _, key := range headersObj.Keys() {
header.Set(key, headersObj.Get(key).String())
}
case "cookies":
cookiesV := params.Get(k)
if goja.IsUndefined(cookiesV) || goja.IsNull(cookiesV) {
continue
}
cookies := cookiesV.ToObject(rt)
if cookies == nil {
continue
}
for _, key := range cookies.Keys() {
replace := false
value := ""
cookieV := cookies.Get(key)
if goja.IsUndefined(cookieV) || goja.IsNull(cookieV) {
continue
}
switch cookieV.ExportType() {
case reflect.TypeOf(map[string]interface{}{}):
cookie := cookieV.ToObject(rt)
for _, attr := range cookie.Keys() {
switch strings.ToLower(attr) {
case "replace":
replace = cookie.Get(attr).ToBoolean()
case "value":
value = cookie.Get(attr).String()
}
}
default:
value = cookieV.String()
}
if _, ok := cookiesMap[key]; !ok || replace {
cookiesMap[key] = http.Cookie{
Name: key,
Value: value,
}
}
}
case "tags":
tagsV := params.Get(k)
if goja.IsUndefined(tagsV) || goja.IsNull(tagsV) {
Expand All @@ -147,6 +197,31 @@ func (*WS) Connect(ctx context.Context, url string, args ...goja.Value) (*WSHTTP

}

cookies := make([]*http.Cookie, 0, len(cookiesMap))
for _, cookie := range cookiesMap {
v := cookie
cookies = append(cookies, &v)
}

var jar http.CookieJar
if len(cookies) > 0 {
jar, err = cookiejar.New(nil)
if err != nil {
return nil, err
}

// SetCookies looking for http or https scheme. Otherwise it does nothing. This is mini hack to bypass it.
oldScheme := u.Scheme
switch u.Scheme {
case "ws":
u.Scheme = "http"
case "wss":
u.Scheme = "https"
}
jar.SetCookies(u, cookies)
u.Scheme = oldScheme
}

if state.Options.SystemTags.Has(stats.TagURL) {
tags["url"] = url
}
Expand All @@ -165,6 +240,7 @@ func (*WS) Connect(ctx context.Context, url string, args ...goja.Value) (*WSHTTP
NetDialContext: state.Dialer.DialContext,
Proxy: http.ProxyFromEnvironment,
TLSClientConfig: tlsConfig,
Jar: jar,
}

start := time.Now()
Expand Down
21 changes: 20 additions & 1 deletion js/modules/k6/ws/ws_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -164,7 +164,7 @@ func TestSession(t *testing.T) {
socket.send("test")
})
socket.on("message", function (data){
if (!data=="test") {
if (!(data=="test")) {
throw new Error ("echo'd data doesn't match our message!");
}
socket.close()
Expand Down Expand Up @@ -315,6 +315,25 @@ func TestSession(t *testing.T) {
})
assertSessionMetricsEmitted(t, stats.GetBufferedSamples(samples), "", sr("WSBIN_URL/ws-echo"), 101, "")

t.Run("client_cookie", func(t *testing.T) {
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I tried adding a test for replace, but it seems that cookies in the global cookie jar (http.cookieJar()) are not propagated to ws.connect(). I.e. this script doesn't send the cookie:

import http from 'k6/http';
import ws from 'k6/ws';

export default function () {
  const jar = http.cookieJar();
  jar.set("ws://127.0.0.1:8080/", "Session", "123");
  var res = ws.connect("ws://127.0.0.1:8080/", null, function(socket){
    socket.on("message", function (data){
      console.log(data);
      socket.close();
    });
  });
}

So my question is: should it? Based on this comment probably not as it would be a breaking change. But should we then have an equivalent ws.cookieJar()?

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think it should and I also think that is the better first change.

Unfortunately the web socket library we use takes only jars and as adding both a cookie and using the jar should probably not add the cookie to the jar you will need to:

  1. create a new jar to copy the values from the provide one (or the default one) and add the cookie to that
  2. set any cookies set by the websocket one back in the original jar

Given that ... I am more for not having a way to specify a cookie but to

  1. have a way to specify a jar
  2. use the default one by default ;)

And then whoever wants to specify a Cookie can just add it to the jar and pass the jar or continue setting the Header ... both of which arguably aren't that much harder especially as I would expect the bigger problem will be that you would've made an HTTP request it would set you the cookie in the jar, but then that won't be used by the websocket. And the someone less likely variant of websocket HTTP call setting a cookie that isn';t set on the default jar

Copy link
Contributor Author

@imiric imiric Jun 9, 2021 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@mstoykov I'm a bit confused by the changes you want 😅 Can you summarize them as commit messages? :)

FWIW gorilla/websocket does support setting a Cookie header and not just via a jar. See Dialer.DialContext. There's an issue about the header overriding the jar, but IMO this seems like desired behavior.

So the way I see it to handle this "properly" and achieve feature-parity with how it works for http we should:

  • Add a per-VU WS-only cookie jar (ws.cookieJar()) from which we will copy to the Dialer.Jar for each new connection.

  • Keep this params approach to override cookies per connection, maybe without using Dialer.Jar but simply setting it as a header.

  • Allow passing a standalone jar via params (new ws.CookieJar()) that will be set to Dialer.Jar after merging any cookies set via params and in the global jar. We should document how overriding works to avoid surprises.

I would've preferred if all this could work with http.cookieJar(), since the only reason we need a WS approach is to avoid breaking existing scripts, but hopefully we can unify this when we refactor WS and after the new HTTP API.

WDYT? cc: @na--

EDIT: Needless to say, I doubt this will be tested and ready for v0.33.0 so we might want to postpone it.

Also this expands on the scope of #1226 which merely wanted it passed via params. So we could merge a variation of this change and figure out jar handling later.(?)

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

FWIW gorilla/websocket does support setting a Cookie header and not just via a jar. See Dialer.DialContext. There's an issue about the header overriding the jar, but IMO this seems like desired behavior.

As mentioned in this comment from the "fix" this seems like the correct behaviour ;). And I agree with that.

So the way I see it to handle this "properly" and achieve feature-parity with how it works for http we should:

This is IMO the wrong thing to try to do given #1923 which literally spells out how not ... good and confusing the cookie works for HTTP. Also it was provoked by me trying to work on this exact PR/issue.

IMO this should :

  1. Use the default cookiejar from lib.state by default.
  2. be able to take Params.jar which is http.Cookiejar .. .we can probably move it and keep it's old name as well ... although a bit finicky.
  3. Params.headers.Cookie to overwrite any of the above (so no change I think)
  4. Params.Cookies should not be supported, at least for now. (look at how the merging works for HTTP and you will see it as PITA, and likely is ... completely not useful as I will try to explain below.)

Reasons:

  1. I would expect that if I do an HTTP request and I then make a websocket one they will use the same jars. I would argue this will remove the reason for most of the people requesting this, after all the reports are about authenticating, usually with a session which likely got received through and HTTP request.
  2. Having a way to specify another jar lets them set cookies in a more uniform way even if they don't use HTTP request to get the cookie
  3. The header Cookie and how the order in it works is not very well standardized, so having us merge stuff into it "somehow" just makes things confusing IMO and requires reading way too much documentation from the user to learn how it works (also we will need to write it). IMO if we have a way to get the cookies for a given URL out of the cookie jar (which we do), have a way to modify those values and then seriliaze them to a Cookie value so people can do stuff like
var jar  = ... // cookiejar
var cookies = jar.cookiesForURL(myurl);
cookies[mycoolCookie]="specialValue";
var headers = {Cookie: serializeToCookieValue(cookies)};

will remove the whole need for having Cookie in the params. It might be 2 lines more or something like that, but given how much cleaner it will be and less confusing IMO that should've been what was done instead of the current merging of cookies between params.jar and params.cookies that is happening ...

I would've preferred if all this could work with http.cookieJar(), since the only reason we need a WS approach is to avoid breaking existing scripts, but hopefully we can unify this when we refactor WS and after the new HTTP API.

I see no reason why not do it now ... or just not do these changes at all for now ;) .

This has been the case forever and even now users can always just set the Cookie header. The main benefit of providing a jar IMO is that it makes things less confusing when you need to get a cookie from HTTP request and give it to the websocket one (which is still HTTP first) and , secondly, for when a websocket connection set cookies ... which currently needs to be parsed by the user somehow.

The only useful thing about params.Cookie currently is:

  1. it lets you set a Cookie for every redirect after that ... which probably is a bug now that I write it
  2. It lets you to maybe overwrite a cookie set in the jar given that those two match
  3. It lets you merge easier with jar cookies
  4. It lets you set cookies easier then setting the header.

The problems with those are:

  1. this IMO is probably not what anyone expects ?!?
  2. it only works if you set the replace:true part which makes the syntax horrible and I doubt there are many people who even know about this or use it. (it also means that it will change if the jar suddenly gets the same cookie for some reason)
  3. this is probably the most useful one, but even this one can probably be done easier and IMO more consistently with some js code
  4. great ... unless look at 2. at which point it doesn't.

The whole problem IMO comes from the fact that understanding how the 3 ways interact and interact with redirects makes the whole thing way too confusing to the point that even the (technically small) improvements that it does in the case 2-4 likely get removed by the fact that anyone who needs to set cookies specifically as the usual cookie jar support doesn't work will also hit one of the many problems listed above in some respect.

t.Parallel()
_, err := rt.RunString(sr(`
var params = {
cookies: { "Session": "123" },
};
var res = ws.connect("WSBIN_URL/ws-echo-cookie", params, function(socket){
socket.on("message", function (data){
if (!(data == "Session=123")) {
throw new Error ("echo'd data doesn't match our message!");
}
socket.close();
});
});
`))
assert.NoError(t, err)
assertSessionMetricsEmitted(t, stats.GetBufferedSamples(samples), "", sr("WSBIN_URL/ws-echo-cookie"), 101, "")
})

serverCloseTests := []struct {
name string
endpoint string
Expand Down
22 changes: 17 additions & 5 deletions lib/testutils/httpmultibin/httpmultibin.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -106,7 +106,9 @@ type jsonBody struct {
Compression string `json:"compression"`
}

func getWebsocketHandler(echo bool, closePrematurely bool) http.Handler {
// TODO: Refactor this
//nolint:cyclop
func getWebsocketHandler(echo, closePrematurely, echoCookies bool) http.Handler {
return http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
conn, err := (&websocket.Upgrader{}).Upgrade(w, req, w.Header())
if err != nil {
Expand All @@ -129,6 +131,15 @@ func getWebsocketHandler(echo bool, closePrematurely bool) http.Handler {
return
}
}
if echoCookies {
var cookies []string
for _, cookie := range req.Cookies() {
cookies = append(cookies, cookie.String())
}
if err = conn.WriteMessage(websocket.TextMessage, []byte(strings.Join(cookies, "; "))); err != nil {
return
}
}
// closePrematurely=true mimics an invalid WS server that doesn't
// send a close control frame before closing the connection.
if !closePrematurely {
Expand Down Expand Up @@ -257,10 +268,11 @@ func NewHTTPMultiBin(t testing.TB) *HTTPMultiBin {
// Create a http.ServeMux and set the httpbin handler as the default
mux := http.NewServeMux()
mux.Handle("/brotli", getEncodedHandler(t, "br"))
mux.Handle("/ws-echo", getWebsocketHandler(true, false))
mux.Handle("/ws-echo-invalid", getWebsocketHandler(true, true))
mux.Handle("/ws-close", getWebsocketHandler(false, false))
mux.Handle("/ws-close-invalid", getWebsocketHandler(false, true))
mux.Handle("/ws-echo", getWebsocketHandler(true, false, false))
mux.Handle("/ws-echo-cookie", getWebsocketHandler(false, false, true))
mux.Handle("/ws-echo-invalid", getWebsocketHandler(true, true, false))
mux.Handle("/ws-close", getWebsocketHandler(false, false, false))
mux.Handle("/ws-close-invalid", getWebsocketHandler(false, true, false))
mux.Handle("/zstd", getEncodedHandler(t, "zstd"))
mux.Handle("/zstd-br", getZstdBrHandler(t))
mux.Handle("/", httpbin.New().Handler())
Expand Down