From 8a0f80a054e0c8c2ecf246c0d3e8278df0b2d07a Mon Sep 17 00:00:00 2001
From: Ivan Krutov <vania-pooh@aerokube.com>
Date: Fri, 1 Jul 2022 14:24:56 +0300
Subject: [PATCH] Sending correct Host header in drivers mode too

---
 selenoid.go        | 9 +++++----
 service/docker.go  | 2 +-
 service/driver.go  | 2 +-
 service/service.go | 1 +
 session/session.go | 2 +-
 5 files changed, 9 insertions(+), 7 deletions(-)

diff --git a/selenoid.go b/selenoid.go
index 55a88103..f914a128 100644
--- a/selenoid.go
+++ b/selenoid.go
@@ -201,8 +201,8 @@ func create(w http.ResponseWriter, r *http.Request) {
 	u := startedService.Url
 	cancel := startedService.Cancel
 	host := "localhost"
-	if startedService.Container != nil {
-		host = startedService.Container.Origin
+	if startedService.Origin != "" {
+		host = startedService.Origin
 	}
 
 	var resp *http.Response
@@ -300,6 +300,7 @@ func create(w http.ResponseWriter, r *http.Request) {
 		URL:       u,
 		Container: startedService.Container,
 		HostPort:  startedService.HostPort,
+		Origin:    startedService.Origin,
 		Timeout:   sessionTimeout,
 		TimeoutCh: onTimeout(sessionTimeout, func() {
 			request{r}.session(s.ID).Delete(requestId)
@@ -532,8 +533,8 @@ func proxy(w http.ResponseWriter, r *http.Request) {
 				}
 				r.URL.Host, r.URL.Path = sess.URL.Host, path.Clean(sess.URL.Path+r.URL.Path)
 				r.Host = "localhost"
-				if sess.Container != nil {
-					r.Host = sess.Container.Origin
+				if sess.Origin != "" {
+					r.Host = sess.Origin
 				}
 				return
 			}
diff --git a/service/docker.go b/service/docker.go
index c79b673a..0334a62e 100644
--- a/service/docker.go
+++ b/service/docker.go
@@ -231,10 +231,10 @@ func (d *Docker) StartWithCancel() (*StartedService, error) {
 		Container: &session.Container{
 			ID:        browserContainerId,
 			IPAddress: getContainerIP(d.Environment.Network, stat),
-			Origin:    origin,
 			Ports:     publishedPortsInfo,
 		},
 		HostPort: hostPort,
+		Origin:   origin,
 		Cancel: func() {
 			if videoContainerId != "" {
 				stopVideoContainer(ctx, cl, requestId, videoContainerId, d.Environment)
diff --git a/service/driver.go b/service/driver.go
index 587d37cd..25348eb0 100644
--- a/service/driver.go
+++ b/service/driver.go
@@ -81,7 +81,7 @@ func (d *Driver) StartWithCancel() (*StartedService, error) {
 	if d.Caps.VNC {
 		hp.VNC = "127.0.0.1:5900"
 	}
-	return &StartedService{Url: u, HostPort: hp, Cancel: func() { d.stopProcess(cmd) }}, nil
+	return &StartedService{Url: u, HostPort: hp, Origin: fmt.Sprintf("localhost:%s", port), Cancel: func() { d.stopProcess(cmd) }}, nil
 }
 
 func (d *Driver) stopProcess(cmd *exec.Cmd) {
diff --git a/service/service.go b/service/service.go
index c46890e8..bcb17f3d 100644
--- a/service/service.go
+++ b/service/service.go
@@ -45,6 +45,7 @@ type StartedService struct {
 	Url       *url.URL
 	Container *session.Container
 	HostPort  session.HostPort
+	Origin    string
 	Cancel    func()
 }
 
diff --git a/session/session.go b/session/session.go
index 6beaf6ee..2a18bc50 100644
--- a/session/session.go
+++ b/session/session.go
@@ -65,7 +65,6 @@ func (c *Caps) BrowserName() string {
 type Container struct {
 	ID        string            `json:"id"`
 	IPAddress string            `json:"ip"`
-	Origin    string            `json:"origin,omitempty"`
 	Ports     map[string]string `json:"exposedPorts,omitempty"`
 }
 
@@ -76,6 +75,7 @@ type Session struct {
 	URL       *url.URL
 	Container *Container
 	HostPort  HostPort
+	Origin    string
 	Cancel    func()
 	Timeout   time.Duration
 	TimeoutCh chan struct{}