From 14b6df68ce75fec4edf166d9bf248b745aa12fb9 Mon Sep 17 00:00:00 2001
From: Alexey Khit <alexey.khit@gmail.com>
Date: Tue, 6 Sep 2022 18:09:44 +0300
Subject: [PATCH] Adds support nginx with wrong port

---
 cmd/api/api.go |  1 +
 cmd/api/ws.go  | 38 ++++++++++++++++++++++++++++++++------
 2 files changed, 33 insertions(+), 6 deletions(-)

diff --git a/cmd/api/api.go b/cmd/api/api.go
index 36d7aeeb..8df82847 100644
--- a/cmd/api/api.go
+++ b/cmd/api/api.go
@@ -34,6 +34,7 @@ func Init() {
 	log = app.GetLogger("api")
 
 	initStatic(cfg.Mod.StaticDir)
+	initWS()
 
 	HandleFunc("/api/frame.mp4", frameHandler)
 	HandleFunc("/api/frame.raw", frameHandler)
diff --git a/cmd/api/ws.go b/cmd/api/ws.go
index 7fc51343..8c2f80b0 100644
--- a/cmd/api/ws.go
+++ b/cmd/api/ws.go
@@ -4,16 +4,42 @@ import (
 	"github.com/AlexxIT/go2rtc/pkg/streamer"
 	"github.com/gorilla/websocket"
 	"net/http"
+	"net/url"
+	"strings"
 	"sync"
 )
 
-type WSHandler func(ctx *Context, msg *streamer.Message)
-
-var apiWsUp = websocket.Upgrader{
-	ReadBufferSize:  1024,
-	WriteBufferSize: 512000,
+func initWS() {
+	wsUp = &websocket.Upgrader{
+		ReadBufferSize:  1024,
+		WriteBufferSize: 512000,
+	}
+	wsUp.CheckOrigin = func(r *http.Request) bool {
+		origin := r.Header["Origin"]
+		if len(origin) == 0 {
+			return true
+		}
+		o, err := url.Parse(origin[0])
+		if err != nil {
+			return false
+		}
+		if o.Host == r.Host {
+			return true
+		}
+		log.Trace().Msgf("[api.ws] origin: %s, host: %s", o.Host, r.Host)
+		// some users change Nginx external port using Docker port
+		// so origin will be with a port and host without
+		if i := strings.IndexByte(o.Host, ':'); i > 0 {
+			return o.Host[:i] == r.Host
+		}
+		return false
+	}
 }
 
+var wsUp *websocket.Upgrader
+
+type WSHandler func(ctx *Context, msg *streamer.Message)
+
 type Context struct {
 	Conn     *websocket.Conn
 	Request  *http.Request
@@ -24,7 +50,7 @@ type Context struct {
 }
 
 func (ctx *Context) Upgrade(w http.ResponseWriter, r *http.Request) (err error) {
-	ctx.Conn, err = apiWsUp.Upgrade(w, r, nil)
+	ctx.Conn, err = wsUp.Upgrade(w, r, nil)
 	ctx.Request = r
 	return
 }