#730 bug(websockets): not properly checked request originv2.3.1-beta.3

#730 bug(websockets): not properly checked request origin

14 files changed, 213 insertions, 46 deletions

diff --git a/.github/workflows/linux.yml b/.github/workflows/linux.yml
index 3b3f7e9f..69269557 100644
--- a/.github/workflows/linux.yml
+++ b/.github/workflows/linux.yml
@@ -92,6 +92,7 @@ jobs:
           go test -v -race -cover -tags=debug -coverpkg=./... -coverprofile=./coverage-ci/rpc.txt -covermode=atomic ./tests/plugins/rpc
           go test -v -race -cover -tags=debug -coverpkg=./... -coverprofile=./coverage-ci/kv_plugin.txt -covermode=atomic ./tests/plugins/kv
           go test -v -race -cover -tags=debug -coverpkg=./... -coverprofile=./coverage-ci/websockets.txt -covermode=atomic ./tests/plugins/websockets
+          go test -v -race -cover -tags=debug -coverpkg=./... -coverprofile=./coverage-ci/ws_origin.txt -covermode=atomic ./plugins/websockets
           docker-compose -f ./tests/docker-compose.yaml down
           cat ./coverage-ci/*.txt > ./coverage-ci/summary.txt
 
diff --git a/.github/workflows/windows.yml b/.github/workflows/windows.yml
index 92d76d2c..227c725b 100644
--- a/.github/workflows/windows.yml
+++ b/.github/workflows/windows.yml
@@ -67,7 +67,6 @@ jobs:
       - name: Run golang tests on Windows
         run: |
           docker-compose -f ./tests/docker-compose.yaml up -d
-          mkdir ./coverage-ci
           go test -v -race ./pkg/transport/pipe
           go test -v -race ./pkg/transport/socket
           go test -v -race ./pkg/pool
@@ -92,5 +91,5 @@ jobs:
           go test -v -race ./tests/plugins/rpc
           go test -v -race ./tests/plugins/kv
           go test -v -race ./tests/plugins/websockets
+          go test -v -race ./plugins/websockets
           docker-compose -f ./tests/docker-compose.yaml down
-          cat ./coverage-ci/*.txt > ./coverage-ci/summary.txt
diff --git a/CHANGELOG.md b/CHANGELOG.md
index cbb9936b..6f7cf2cc 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -8,6 +8,7 @@ v2.3.1 (_.06.2021)
 - 🐛 Fix: Bugs with `boltdb` storage: [Boom](https://github.com/spiral/roadrunner/issues/717), [Boom](https://github.com/spiral/roadrunner/issues/718), [Boom](https://github.com/spiral/roadrunner/issues/719)
 - 🐛 Fix: Bug with incorrect redis initialization and usage [Bug](https://github.com/spiral/roadrunner/issues/720)
 - 🐛 Fix: Bug, Goridge duplicate error messages [Bug](https://github.com/spiral/goridge/issues/128)
+- 🐛 Fix: Bug, incorrect request `origin` check [Bug](https://github.com/spiral/roadrunner/issues/727)
 
 ## 📦 Packages:
 
diff --git a/Makefile b/Makefile
index 15f9e394..fbd5ab76 100755
--- a/Makefile
+++ b/Makefile
@@ -32,6 +32,7 @@ test_coverage:
 	go test -v -race -cover -tags=debug -coverpkg=./... -coverprofile=./coverage/rpc.out -covermode=atomic ./tests/plugins/rpc
 	go test -v -race -cover -tags=debug -coverpkg=./... -coverprofile=./coverage/kv_plugin.out -covermode=atomic ./tests/plugins/kv
 	go test -v -race -cover -tags=debug -coverpkg=./... -coverprofile=./coverage/ws_plugin.out -covermode=atomic ./tests/plugins/websockets
+	go test -v -race -cover -tags=debug -coverpkg=./... -coverprofile=./coverage/ws_origin.out -covermode=atomic ./plugins/websockets
 	cat ./coverage/*.out > ./coverage/summary.out
 	docker-compose -f tests/docker-compose.yaml down
 
@@ -61,6 +62,7 @@ test: ## Run application tests
 	go test -v -race -tags=debug ./tests/plugins/rpc
 	go test -v -race -tags=debug ./tests/plugins/kv
 	go test -v -race -tags=debug ./tests/plugins/websockets
+	go test -v -race -tags=debug ./plugins/websockets
 	docker-compose -f tests/docker-compose.yaml down
 
 testGo1.17beta1: ## Run application tests
@@ -89,4 +91,5 @@ testGo1.17beta1: ## Run application tests
 	go1.17beta1 test -v -race -tags=debug ./tests/plugins/rpc
 	go1.17beta1 test -v -race -tags=debug ./tests/plugins/kv
 	go1.17beta1 test -v -race -tags=debug ./tests/plugins/websockets
+	go1.17beta1 test -v -race -tags=debug ./plugins/websockets
 	docker-compose -f tests/docker-compose.yaml down
diff --git a/plugins/websockets/config.go b/plugins/websockets/config.go
index 93d9ac3b..deb4406c 100644
--- a/plugins/websockets/config.go
+++ b/plugins/websockets/config.go
@@ -1,6 +1,7 @@
 package websockets
 
 import (
+	"strings"
 	"time"
 
 	"github.com/spiral/roadrunner/v2/pkg/pool"
@@ -57,9 +58,15 @@ type Config struct {
 	PubSubs    []string `mapstructure:"pubsubs"`
 	Middleware []string `mapstructure:"middleware"`
 
-	Redis *RedisConfig `mapstructure:"redis"`
+	AllowedOrigin string `mapstructure:"allowed_origin"`
+
+	// wildcard origin
+	allowedWOrigins []wildcard
+	allowedOrigins  []string
+	allowedAll      bool
 
-	Pool *pool.Config `mapstructure:"pool"`
+	Redis *RedisConfig `mapstructure:"redis"`
+	Pool  *pool.Config `mapstructure:"pool"`
 }
 
 // InitDefault initialize default values for the ws config
@@ -67,6 +74,7 @@ func (c *Config) InitDefault() {
 	if c.Path == "" {
 		c.Path = "/ws"
 	}
+
 	if len(c.PubSubs) == 0 {
 		// memory used by default
 		c.PubSubs = append(c.PubSubs, "memory")
@@ -86,10 +94,9 @@ func (c *Config) InitDefault() {
 		if c.Pool.DestroyTimeout == 0 {
 			c.Pool.DestroyTimeout = time.Minute
 		}
-		if c.Pool.Supervisor == nil {
-			return
+		if c.Pool.Supervisor != nil {
+			c.Pool.Supervisor.InitDefaults()
 		}
-		c.Pool.Supervisor.InitDefaults()
 	}
 
 	if c.Redis != nil {
@@ -98,4 +105,22 @@ func (c *Config) InitDefault() {
 			c.Redis.Addrs = append(c.Redis.Addrs, "localhost:6379")
 		}
 	}
+
+	if c.AllowedOrigin == "" {
+		c.AllowedOrigin = "*"
+	}
+
+	// Normalize
+	origin := strings.ToLower(c.AllowedOrigin)
+	if origin == "*" {
+		// If "*" is present in the list, turn the whole list into a match all
+		c.allowedAll = true
+		return
+	} else if i := strings.IndexByte(origin, '*'); i >= 0 {
+		// Split the origin in two: start and end string without the *
+		w := wildcard{origin[0:i], origin[i+1:]}
+		c.allowedWOrigins = append(c.allowedWOrigins, w)
+	} else {
+		c.allowedOrigins = append(c.allowedOrigins, origin)
+	}
 }
diff --git a/plugins/websockets/connection/connection.go b/plugins/websockets/connection/connection.go
index 2b847173..04c29d83 100644
--- a/plugins/websockets/connection/connection.go
+++ b/plugins/websockets/connection/connection.go
@@ -22,7 +22,7 @@ func NewConnection(wsConn *websocket.Conn, log logger.Logger) *Connection {
 	}
 }
 
-func (c *Connection) Write(mt int, data []byte) error {
+func (c *Connection) Write(data []byte) error {
 	c.Lock()
 	defer c.Unlock()
 
@@ -34,7 +34,7 @@ func (c *Connection) Write(mt int, data []byte) error {
 		}
 	}()
 
-	err := c.conn.WriteMessage(mt, data)
+	err := c.conn.WriteMessage(websocket.TextMessage, data)
 	if err != nil {
 		return errors.E(op, err)
 	}
diff --git a/plugins/websockets/executor/executor.go b/plugins/websockets/executor/executor.go
index e3d47166..5f904d26 100644
--- a/plugins/websockets/executor/executor.go
+++ b/plugins/websockets/executor/executor.go
@@ -5,7 +5,6 @@ import (
 	"net/http"
 	"sync"
 
-	"github.com/fasthttp/websocket"
 	json "github.com/json-iterator/go"
 	"github.com/spiral/errors"
 	websocketsv1 "github.com/spiral/roadrunner/v2/pkg/proto/websockets/v1beta"
@@ -100,7 +99,7 @@ func (e *Executor) StartCommandLoop() error { //nolint:gocognit
 					return errors.E(op, fmt.Errorf("%v,%v", err, errJ))
 				}
 
-				errW := e.conn.Write(websocket.BinaryMessage, packet)
+				errW := e.conn.Write(packet)
 				if errW != nil {
 					e.log.Error("error writing payload to the connection", "payload", packet, "error", errW)
 					return errors.E(op, fmt.Errorf("%v,%v", err, errW))
@@ -120,7 +119,7 @@ func (e *Executor) StartCommandLoop() error { //nolint:gocognit
 				return errors.E(op, err)
 			}
 
-			err = e.conn.Write(websocket.BinaryMessage, packet)
+			err = e.conn.Write(packet)
 			if err != nil {
 				e.log.Error("error writing payload to the connection", "payload", packet, "error", err)
 				return errors.E(op, err)
@@ -150,7 +149,7 @@ func (e *Executor) StartCommandLoop() error { //nolint:gocognit
 				return errors.E(op, err)
 			}
 
-			err = e.conn.Write(websocket.BinaryMessage, packet)
+			err = e.conn.Write(packet)
 			if err != nil {
 				e.log.Error("error writing payload to the connection", "payload", packet, "error", err)
 				return errors.E(op, err)
diff --git a/plugins/websockets/origin.go b/plugins/websockets/origin.go
new file mode 100644
index 00000000..c6d9c9b8
--- /dev/null
+++ b/plugins/websockets/origin.go
@@ -0,0 +1,28 @@
+package websockets
+
+import (
+	"strings"
+)
+
+func isOriginAllowed(origin string, cfg *Config) bool {
+	if cfg.allowedAll {
+		return true
+	}
+
+	origin = strings.ToLower(origin)
+	// simple case
+	origin = strings.ToLower(origin)
+	for _, o := range cfg.allowedOrigins {
+		if o == origin {
+			return true
+		}
+	}
+	// check wildcards
+	for _, w := range cfg.allowedWOrigins {
+		if w.match(origin) {
+			return true
+		}
+	}
+
+	return false
+}
diff --git a/plugins/websockets/origin_test.go b/plugins/websockets/origin_test.go
new file mode 100644
index 00000000..e877fad3
--- /dev/null
+++ b/plugins/websockets/origin_test.go
@@ -0,0 +1,67 @@
+package websockets
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestConfig_Origin(t *testing.T) {
+	cfg := &Config{
+		AllowedOrigin: "*",
+	}
+
+	cfg.InitDefault()
+
+	assert.True(t, isOriginAllowed("http://some.some.some.sssome", cfg))
+	assert.True(t, isOriginAllowed("http://", cfg))
+	assert.True(t, isOriginAllowed("http://google.com", cfg))
+	assert.True(t, isOriginAllowed("ws://*", cfg))
+	assert.True(t, isOriginAllowed("*", cfg))
+	assert.True(t, isOriginAllowed("you are bad programmer", cfg)) // True :(
+	assert.True(t, isOriginAllowed("****", cfg))
+	assert.True(t, isOriginAllowed("asde!@#!!@#!%", cfg))
+	assert.True(t, isOriginAllowed("http://*.domain.com", cfg))
+}
+
+func TestConfig_OriginWildCard(t *testing.T) {
+	cfg := &Config{
+		AllowedOrigin: "https://*my.site.com",
+	}
+
+	cfg.InitDefault()
+
+	assert.True(t, isOriginAllowed("https://my.site.com", cfg))
+	assert.False(t, isOriginAllowed("http://", cfg))
+	assert.False(t, isOriginAllowed("http://google.com", cfg))
+	assert.False(t, isOriginAllowed("ws://*", cfg))
+	assert.False(t, isOriginAllowed("*", cfg))
+	assert.False(t, isOriginAllowed("you are bad programmer", cfg)) // True :(
+	assert.False(t, isOriginAllowed("****", cfg))
+	assert.False(t, isOriginAllowed("asde!@#!!@#!%", cfg))
+	assert.False(t, isOriginAllowed("http://*.domain.com", cfg))
+
+	assert.False(t, isOriginAllowed("https://*site.com", cfg))
+	assert.True(t, isOriginAllowed("https://some.my.site.com", cfg))
+}
+
+func TestConfig_OriginWildCard2(t *testing.T) {
+	cfg := &Config{
+		AllowedOrigin: "https://my.*.com",
+	}
+
+	cfg.InitDefault()
+
+	assert.True(t, isOriginAllowed("https://my.site.com", cfg))
+	assert.False(t, isOriginAllowed("http://", cfg))
+	assert.False(t, isOriginAllowed("http://google.com", cfg))
+	assert.False(t, isOriginAllowed("ws://*", cfg))
+	assert.False(t, isOriginAllowed("*", cfg))
+	assert.False(t, isOriginAllowed("you are bad programmer", cfg)) // True :(
+	assert.False(t, isOriginAllowed("****", cfg))
+	assert.False(t, isOriginAllowed("asde!@#!!@#!%", cfg))
+	assert.False(t, isOriginAllowed("http://*.domain.com", cfg))
+
+	assert.False(t, isOriginAllowed("https://*site.com", cfg))
+	assert.True(t, isOriginAllowed("https://my.bad.com", cfg))
+}
diff --git a/plugins/websockets/plugin.go b/plugins/websockets/plugin.go
index 6dfe6ca3..8b708187 100644
--- a/plugins/websockets/plugin.go
+++ b/plugins/websockets/plugin.go
@@ -82,6 +82,9 @@ func (p *Plugin) Init(cfg config.Configurer, log logger.Logger, server server.Se
 		HandshakeTimeout: time.Second * 60,
 		ReadBufferSize:   1024,
 		WriteBufferSize:  1024,
+		CheckOrigin: func(r *http.Request) bool {
+			return isOriginAllowed(r.Header.Get("Origin"), p.cfg)
+		},
 	}
 	p.serveExit = make(chan struct{})
 	p.server = server
diff --git a/plugins/websockets/pool/workers_pool.go b/plugins/websockets/pool/workers_pool.go
index 1a7c6f8a..a196d1f0 100644
--- a/plugins/websockets/pool/workers_pool.go
+++ b/plugins/websockets/pool/workers_pool.go
@@ -3,11 +3,12 @@ package pool
 import (
 	"sync"
 
-	"github.com/fasthttp/websocket"
+	json "github.com/json-iterator/go"
 	websocketsv1 "github.com/spiral/roadrunner/v2/pkg/proto/websockets/v1beta"
 	"github.com/spiral/roadrunner/v2/pkg/pubsub"
 	"github.com/spiral/roadrunner/v2/plugins/logger"
 	"github.com/spiral/roadrunner/v2/plugins/websockets/connection"
+	"github.com/spiral/roadrunner/v2/utils"
 )
 
 type WorkersPool struct {
@@ -67,6 +68,12 @@ func (wp *WorkersPool) get() map[string]struct{} {
 	return wp.resPool.Get().(map[string]struct{})
 }
 
+// Response from the server
+type Response struct {
+	Topic   string `json:"topic"`
+	Payload string `json:"payload"`
+}
+
 func (wp *WorkersPool) do() { //nolint:gocognit
 	go func() {
 		for {
@@ -89,44 +96,52 @@ func (wp *WorkersPool) do() { //nolint:gocognit
 					continue
 				}
 
-				res := wp.get()
-
+				// send a message to every topic
 				for i := 0; i < len(msg.GetTopics()); i++ {
+					// get free map
+					res := wp.get()
+
 					// get connections for the particular topic
 					br.Connections(msg.GetTopics()[i], res)
-				}
 
-				if len(res) == 0 {
-					for i := 0; i < len(msg.GetTopics()); i++ {
+					if len(res) == 0 {
 						wp.log.Info("no such topic", "topic", msg.GetTopics()[i])
+						wp.put(res)
+						continue
 					}
-					wp.put(res)
-					continue
-				}
 
-				for i := range res {
-					c, ok := wp.connections.Load(i)
-					if !ok {
-						for i := 0; i < len(msg.GetTopics()); i++ {
+					// res is a map with a connectionsID
+					for topic := range res {
+						c, ok := wp.connections.Load(topic)
+						if !ok {
 							wp.log.Warn("the user disconnected connection before the message being written to it", "broker", msg.GetBroker(), "topics", msg.GetTopics()[i])
+							wp.put(res)
+							continue
 						}
-						continue
-					}
 
-					conn := c.(*connection.Connection)
+						response := &Response{
+							Topic:   msg.GetTopics()[i],
+							Payload: utils.AsString(msg.GetPayload()),
+						}
 
-					// put data into the bytes buffer
-					err := conn.Write(websocket.BinaryMessage, msg.GetPayload())
-					if err != nil {
-						for i := 0; i < len(msg.GetTopics()); i++ {
-							wp.log.Error("error sending payload over the connection", "error", err, "broker", msg.GetBroker(), "topics", msg.GetTopics()[i])
+						d, err := json.Marshal(response)
+						if err != nil {
+							wp.log.Error("error marshaling response", "error", err)
+							wp.put(res)
+							break
+						}
+
+						// put data into the bytes buffer
+						err = c.(*connection.Connection).Write(d)
+						if err != nil {
+							for i := 0; i < len(msg.GetTopics()); i++ {
+								wp.log.Error("error sending payload over the connection", "error", err, "broker", msg.GetBroker(), "topics", msg.GetTopics()[i])
+							}
+							wp.put(res)
+							continue
 						}
-						continue
 					}
 				}
-
-				// put map with results back
-				wp.put(res)
 			case <-wp.exit:
 				wp.log.Info("get exit signal, exiting from the workers pool")
 				return
diff --git a/plugins/websockets/wildcard.go b/plugins/websockets/wildcard.go
new file mode 100644
index 00000000..2f1c6601
--- /dev/null
+++ b/plugins/websockets/wildcard.go
@@ -0,0 +1,12 @@
+package websockets
+
+import "strings"
+
+type wildcard struct {
+	prefix string
+	suffix string
+}
+
+func (w wildcard) match(s string) bool {
+	return len(s) >= len(w.prefix)+len(w.suffix) && strings.HasPrefix(s, w.prefix) && strings.HasSuffix(s, w.suffix)
+}
diff --git a/tests/plugins/websockets/websocket_plugin_test.go b/tests/plugins/websockets/websocket_plugin_test.go
index 07ee5f12..8321297d 100644
--- a/tests/plugins/websockets/websocket_plugin_test.go
+++ b/tests/plugins/websockets/websocket_plugin_test.go
@@ -356,7 +356,7 @@ func RPCWsMemoryPubAsync(t *testing.T) {
 	_, msg, err = c.ReadMessage()
 	retMsg = utils.AsString(msg)
 	assert.NoError(t, err)
-	assert.Equal(t, "hello, PHP", retMsg)
+	assert.Equal(t, "{\"topic\":\"foo\",\"payload\":\"hello, PHP\"}", retMsg)
 
 	// //// LEAVE foo, foo2 /////////
 	d, err = json.Marshal(messageWS("leave", "memory", []byte("hello websockets"), "foo"))
@@ -386,7 +386,7 @@ func RPCWsMemoryPubAsync(t *testing.T) {
 	_, msg, err = c.ReadMessage()
 	retMsg = utils.AsString(msg)
 	assert.NoError(t, err)
-	assert.Equal(t, "hello, PHP2", retMsg)
+	assert.Equal(t, "{\"topic\":\"foo2\",\"payload\":\"hello, PHP2\"}", retMsg)
 
 	err = c.WriteControl(websocket.CloseMessage, nil, time.Time{})
 	assert.NoError(t, err)
@@ -430,7 +430,7 @@ func RPCWsMemory(t *testing.T) {
 	_, msg, err = c.ReadMessage()
 	retMsg = utils.AsString(msg)
 	assert.NoError(t, err)
-	assert.Equal(t, "hello, PHP", retMsg)
+	assert.Equal(t, "{\"topic\":\"foo\",\"payload\":\"hello, PHP\"}", retMsg)
 
 	// //// LEAVE foo, foo2 /////////
 	d, err = json.Marshal(messageWS("leave", "memory", []byte("hello websockets"), "foo"))
@@ -460,7 +460,7 @@ func RPCWsMemory(t *testing.T) {
 	_, msg, err = c.ReadMessage()
 	retMsg = utils.AsString(msg)
 	assert.NoError(t, err)
-	assert.Equal(t, "hello, PHP2", retMsg)
+	assert.Equal(t, "{\"topic\":\"foo2\",\"payload\":\"hello, PHP2\"}", retMsg)
 
 	err = c.WriteControl(websocket.CloseMessage, nil, time.Time{})
 	assert.NoError(t, err)
@@ -502,7 +502,7 @@ func RPCWsRedis(t *testing.T) {
 	_, msg, err = c.ReadMessage()
 	retMsg = utils.AsString(msg)
 	assert.NoError(t, err)
-	assert.Equal(t, "hello, PHP", retMsg)
+	assert.Equal(t, "{\"topic\":\"foo\",\"payload\":\"hello, PHP\"}", retMsg)
 
 	// //// LEAVE foo, foo2 /////////
 	d, err = json.Marshal(messageWS("leave", "redis", []byte("hello websockets"), "foo"))
@@ -532,7 +532,7 @@ func RPCWsRedis(t *testing.T) {
 	_, msg, err = c.ReadMessage()
 	retMsg = utils.AsString(msg)
 	assert.NoError(t, err)
-	assert.Equal(t, "hello, PHP2", retMsg)
+	assert.Equal(t, "{\"topic\":\"foo2\",\"payload\":\"hello, PHP2\"}", retMsg)
 
 	err = c.WriteControl(websocket.CloseMessage, nil, time.Time{})
 	assert.NoError(t, err)
@@ -873,7 +873,7 @@ func RPCWsMemoryAllow(t *testing.T) {
 	_, msg, err = c.ReadMessage()
 	retMsg = utils.AsString(msg)
 	assert.NoError(t, err)
-	assert.Equal(t, "hello, PHP", retMsg)
+	assert.Equal(t, "{\"topic\":\"foo\",\"payload\":\"hello, PHP\"}", retMsg)
 
 	// //// LEAVE foo, foo2 /////////
 	d, err = json.Marshal(messageWS("leave", "memory", []byte("hello websockets"), "foo"))
@@ -903,7 +903,7 @@ func RPCWsMemoryAllow(t *testing.T) {
 	_, msg, err = c.ReadMessage()
 	retMsg = utils.AsString(msg)
 	assert.NoError(t, err)
-	assert.Equal(t, "hello, PHP2", retMsg)
+	assert.Equal(t, "{\"topic\":\"foo2\",\"payload\":\"hello, PHP2\"}", retMsg)
 
 	err = c.WriteControl(websocket.CloseMessage, nil, time.Time{})
 	assert.NoError(t, err)
diff --git a/tests/worker-origin.php b/tests/worker-origin.php
new file mode 100644
index 00000000..6ce4de59
--- /dev/null
+++ b/tests/worker-origin.php
@@ -0,0 +1,14 @@
+<?php
+
+use Spiral\RoadRunner\Worker;
+use Spiral\RoadRunner\Http\HttpWorker;
+
+require __DIR__ . '/vendor/autoload.php';
+
+$http = new HttpWorker(Worker::create());
+
+while ($req = $http->waitRequest()) {
+    $http->respond(200, 'Response', [
+        'Access-Control-Allow-Origin' => ['*']
+    ]);
+}