From bc819b42956a243b64ed3c68e9c5accd031df02d Mon Sep 17 00:00:00 2001
From: Ivan Krutov <vania-pooh@aerokube.com>
Date: Fri, 26 Mar 2021 11:12:27 +0300
Subject: [PATCH] Removing only selenoid:options and preserving stuff like
 goog:chromeOptions (fixes #1077)

---
 selenoid.go      | 39 ++++++++++++++++++++++++++++++++-------
 selenoid_test.go | 24 ++++++++++++++++++------
 2 files changed, 50 insertions(+), 13 deletions(-)

diff --git a/selenoid.go b/selenoid.go
index 41222d95..49987313 100644
--- a/selenoid.go
+++ b/selenoid.go
@@ -200,17 +200,11 @@ func create(w http.ResponseWriter, r *http.Request) {
 	u := startedService.Url
 	cancel := startedService.Cancel
 
-	browser.Caps.ExtensionCapabilities = nil
-	browser.W3CCaps.Caps.ExtensionCapabilities = nil
-	for _, c := range browser.W3CCaps.FirstMatch {
-		c.ExtensionCapabilities = nil
-	}
-
 	var resp *http.Response
 	i := 1
 	for ; ; i++ {
 		r.URL.Host, r.URL.Path = u.Host, path.Join(u.Path, r.URL.Path)
-		newBody, _ := json.Marshal(browser)
+		newBody := removeSelenoidOptions(body)
 		req, _ := http.NewRequest(http.MethodPost, r.URL.String(), bytes.NewReader(newBody))
 		ctx, done := context.WithTimeout(r.Context(), newSessionAttemptTimeout)
 		defer done()
@@ -357,6 +351,37 @@ func create(w http.ResponseWriter, r *http.Request) {
 	log.Printf("[%d] [SESSION_CREATED] [%s] [%d] [%.2fs]", requestId, s.ID, i, util.SecondsSince(sessionStartTime))
 }
 
+func removeSelenoidOptions(input []byte) []byte {
+	body := make(map[string]interface{})
+	_ = json.Unmarshal(input, &body)
+	const selenoidOptions = "selenoid:options"
+	if raw, ok := body["desiredCapabilities"]; ok {
+		if dc, ok := raw.(map[string]interface{}); ok {
+			delete(dc, selenoidOptions)
+		}
+	}
+	if raw, ok := body["capabilities"]; ok {
+		if c, ok := raw.(map[string]interface{}); ok {
+			if raw, ok := c["alwaysMatch"]; ok {
+				if am, ok := raw.(map[string]interface{}); ok {
+					delete(am, selenoidOptions)
+				}
+			}
+			if raw, ok := c["firstMatch"]; ok {
+				if fm, ok := raw.([]interface{}); ok {
+					for _, raw := range fm {
+						if c, ok := raw.(map[string]interface{}); ok {
+							delete(c, selenoidOptions)
+						}
+					}
+				}
+			}
+		}
+	}
+	ret, _ := json.Marshal(body)
+	return ret
+}
+
 func preprocessSessionId(sid string) string {
 	if ggrHost != nil {
 		return ggrHost.Sum() + sid
diff --git a/selenoid_test.go b/selenoid_test.go
index 34fae77e..057c2202 100644
--- a/selenoid_test.go
+++ b/selenoid_test.go
@@ -4,7 +4,6 @@ import (
 	"bytes"
 	"context"
 	"fmt"
-	"github.com/aerokube/selenoid/session"
 	"github.com/mafredri/cdp"
 	"github.com/mafredri/cdp/rpcc"
 	"io/ioutil"
@@ -392,10 +391,16 @@ func TestSessionCreatedRedirect(t *testing.T) {
 }
 
 func TestSessionCreatedRemoveExtensionCapabilities(t *testing.T) {
-	extensionCapsMissing := true
+	desiredCapabilitiesPresent := true
+	alwaysMatchPresent := true
+	firstMatchPresent := true
+	chromeOptionsPresent := true
+
 	var browser struct {
+		Caps    map[string]interface{} `json:"desiredCapabilities"`
 		W3CCaps struct {
-			Caps session.Caps `json:"alwaysMatch"`
+			AlwaysMatch map[string]interface{}   `json:"alwaysMatch"`
+			FirstMatch  []map[string]interface{} `json:"firstMatch"`
 		} `json:"capabilities"`
 	}
 
@@ -403,14 +408,21 @@ func TestSessionCreatedRemoveExtensionCapabilities(t *testing.T) {
 	root.Handle("/wd/hub/", http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		err := json.NewDecoder(r.Body).Decode(&browser)
 		AssertThat(t, err, Is{nil})
-		extensionCapsMissing = browser.W3CCaps.Caps.ExtensionCapabilities == nil
+		_, desiredCapabilitiesPresent = browser.Caps["selenoid:options"]
+		_, alwaysMatchPresent = browser.W3CCaps.AlwaysMatch["selenoid:options"]
+		_, chromeOptionsPresent = browser.W3CCaps.AlwaysMatch["goog:chromeOptions"]
+		AssertThat(t, len(browser.W3CCaps.FirstMatch), EqualTo{1})
+		_, firstMatchPresent = browser.W3CCaps.FirstMatch[0]["selenoid:options"]
 	}))
 	manager = &HTTPTest{Handler: root}
 
-	resp, err := httpClient.Post(With(srv.URL).Path("/wd/hub/session"), "", bytes.NewReader([]byte(`{"capabilities":{"alwaysMatch":{"browserName": "firefox", "selenoid:options":{"enableVNC": true}}}}`)))
+	resp, err := httpClient.Post(With(srv.URL).Path("/wd/hub/session"), "", bytes.NewReader([]byte(`{"desiredCapabilities": {"browserName": "chrome", "selenoid:options": {"enableVNC": true}}, "capabilities":{"alwaysMatch":{"browserName": "chrome", "goog:chromeOptions": {"args": ["headless"]}, "selenoid:options":{"enableVNC": true}}, "firstMatch": [{"platform": "linux", "selenoid:options": {"enableVideo": true}}]}}`)))
 	AssertThat(t, err, Is{nil})
 	AssertThat(t, resp, Code{http.StatusOK})
-	AssertThat(t, extensionCapsMissing, Is{true})
+	AssertThat(t, desiredCapabilitiesPresent, Is{false})
+	AssertThat(t, alwaysMatchPresent, Is{false})
+	AssertThat(t, chromeOptionsPresent, Is{true})
+	AssertThat(t, firstMatchPresent, Is{false})
 }
 
 func TestProxySession(t *testing.T) {