diff --git a/proxy.go b/proxy.go index 9df9869..758160a 100644 --- a/proxy.go +++ b/proxy.go @@ -213,8 +213,26 @@ loop: } switch status { case browserStarted: - sess := resp["sessionId"].(string) - resp["sessionId"] = h.sum() + sess + sess, ok := resp["sessionId"].(string) + if !ok { + protocolError := func() { + reply(w, errMsg("protocol error"), http.StatusBadGateway) + log.Printf("[%d] [BAD_RESPONSE] [%s] [%s] [%s] [%s]\n", id, user, remote, fmtBrowser(browser, version), h.net()) + } + value, ok := resp["value"] + if !ok { + protocolError() + return + } + sess, ok := value.(map[string]interface{})["sessionId"].(string) + if !ok { + protocolError() + return + } + resp["value"].(map[string]interface{})["sessionId"] = h.sum() + sess + } else { + resp["sessionId"] = h.sum() + sess + } reply(w, resp, http.StatusOK) log.Printf("[%d] [%.2fs] [SESSION_CREATED] [%s] [%s] [%s] [%s] [%s] [%d]\n", id, float64(time.Now().Sub(start).Seconds()), user, remote, fmtBrowser(browser, version), h.net(), sess, count) return diff --git a/proxy_test.go b/proxy_test.go index beada69..b9512bf 100644 --- a/proxy_test.go +++ b/proxy_test.go @@ -572,6 +572,129 @@ func TestStartSessionBrowserFailWrongMsg(t *testing.T) { AssertThat(t, rsp, AllOf{Code{http.StatusInternalServerError}, Message{"cannot create session browser-1.0 on any hosts after 1 attempt(s)"}}) } +func TestStartSessionFailJSONWireProtocol(t *testing.T) { + mux := http.NewServeMux() + mux.HandleFunc("/wd/hub/session", postOnly(func(w http.ResponseWriter, r *http.Request) { + //w.WriteHeader(http.StatusBadGateway) + w.Write([]byte(`{}`)) + })) + selenium := httptest.NewServer(mux) + defer selenium.Close() + + host, port := hostportnum(selenium.URL) + node := Host{Name: host, Port: port, Count: 1} + + test.Lock() + defer test.Unlock() + + browsers := Browsers{Browsers: []Browser{ + {Name: "browser", DefaultVersion: "1.0", Versions: []Version{ + {Number: "1.0", Regions: []Region{ + {Hosts: Hosts{ + node, + }}, + }}, + }}}} + updateQuota(user, browsers) + + rsp, err := createSession(`{"desiredCapabilities":{"browserName":"browser", "version":"1.0"}}`) + + AssertThat(t, err, Is{nil}) + AssertThat(t, rsp, AllOf{Code{http.StatusBadGateway}, Message{"protocol error"}}) +} + +func TestStartSessionFailJSONWireProtocolNoSessionID(t *testing.T) { + mux := http.NewServeMux() + mux.HandleFunc("/wd/hub/session", postOnly(func(w http.ResponseWriter, r *http.Request) { + w.Write([]byte(`{"value":{}}`)) + })) + selenium := httptest.NewServer(mux) + defer selenium.Close() + + host, port := hostportnum(selenium.URL) + node := Host{Name: host, Port: port, Count: 1} + + test.Lock() + defer test.Unlock() + + browsers := Browsers{Browsers: []Browser{ + {Name: "browser", DefaultVersion: "1.0", Versions: []Version{ + {Number: "1.0", Regions: []Region{ + {Hosts: Hosts{ + node, + }}, + }}, + }}}} + updateQuota(user, browsers) + + rsp, err := createSession(`{"desiredCapabilities":{"browserName":"browser", "version":"1.0"}}`) + + AssertThat(t, err, Is{nil}) + AssertThat(t, rsp, AllOf{Code{http.StatusBadGateway}, Message{"protocol error"}}) +} + +func TestStartSessionFailJSONWireProtocolWrongType(t *testing.T) { + mux := http.NewServeMux() + mux.HandleFunc("/wd/hub/session", postOnly(func(w http.ResponseWriter, r *http.Request) { + w.Write([]byte(`{"value":{"sessionId":123}}`)) + })) + selenium := httptest.NewServer(mux) + defer selenium.Close() + + host, port := hostportnum(selenium.URL) + node := Host{Name: host, Port: port, Count: 1} + + test.Lock() + defer test.Unlock() + + browsers := Browsers{Browsers: []Browser{ + {Name: "browser", DefaultVersion: "1.0", Versions: []Version{ + {Number: "1.0", Regions: []Region{ + {Hosts: Hosts{ + node, + }}, + }}, + }}}} + updateQuota(user, browsers) + + rsp, err := createSession(`{"desiredCapabilities":{"browserName":"browser", "version":"1.0"}}`) + + AssertThat(t, err, Is{nil}) + AssertThat(t, rsp, AllOf{Code{http.StatusBadGateway}, Message{"protocol error"}}) +} + +func TestStartSessionJSONWireProtocol(t *testing.T) { + mux := http.NewServeMux() + mux.HandleFunc("/wd/hub/session", postOnly(func(w http.ResponseWriter, r *http.Request) { + w.Write([]byte(`{"value":{"sessionId":"123"}}`)) + })) + selenium := httptest.NewServer(mux) + defer selenium.Close() + + host, port := hostportnum(selenium.URL) + node := Host{Name: host, Port: port, Count: 1} + + test.Lock() + defer test.Unlock() + + browsers := Browsers{Browsers: []Browser{ + {Name: "{browser}", DefaultVersion: "1.0", Versions: []Version{ + {Number: "1.0", Regions: []Region{ + {Hosts: Hosts{ + node, + }}, + }}, + }}}} + updateQuota(user, browsers) + + rsp, err := createSession(`{"desiredCapabilities":{"browserName":"{browser}", "version":"1.0"}}`) + + AssertThat(t, err, Is{nil}) + var value map[string]interface{} + AssertThat(t, rsp, AllOf{Code{http.StatusOK}, IsJson{&value}}) + AssertThat(t, value["value"].(map[string]interface{})["sessionId"], EqualTo{fmt.Sprintf("%s123", node.sum())}) +} + func TestDeleteSession(t *testing.T) { mux := http.NewServeMux() mux.HandleFunc("/wd/hub/session/", func(w http.ResponseWriter, r *http.Request) { diff --git a/quota/test.xml b/quota/test.xml index 4a69c74..0b7f4fc 100644 --- a/quota/test.xml +++ b/quota/test.xml @@ -1,5 +1,10 @@ - + + + + + +