From 178c76f410b69f7daaa8442ea051c105856cac74 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E9=A5=BA=E5=AD=90w?= <ef@eternalflame.info>
Date: Fri, 15 Mar 2019 01:16:24 +0800
Subject: [PATCH] Fix websocket allowed origin (#150)

---
 api/stream/stream.go      | 11 ++++++-----
 api/stream/stream_test.go |  8 ++++++++
 2 files changed, 14 insertions(+), 5 deletions(-)

diff --git a/api/stream/stream.go b/api/stream/stream.go
index dd225fe..06aa733 100644
--- a/api/stream/stream.go
+++ b/api/stream/stream.go
@@ -4,6 +4,7 @@ import (
 	"net/http"
 	"net/url"
 	"regexp"
+	"strings"
 	"sync"
 	"time"
 
@@ -155,22 +156,22 @@ func (a *API) Close() {
 }
 
 func isAllowedOrigin(r *http.Request, allowedOrigins []*regexp.Regexp) bool {
-	origin := r.Header["Origin"]
-	if len(origin) == 0 {
+	origin := r.Header.Get("origin")
+	if origin == "" {
 		return true
 	}
 
-	u, err := url.Parse(origin[0])
+	u, err := url.Parse(origin)
 	if err != nil {
 		return false
 	}
 
-	if u.Hostname() == r.Host {
+	if strings.ToLower(u.Host) == strings.ToLower(r.Host) {
 		return true
 	}
 
 	for _, allowedOrigin := range allowedOrigins {
-		if allowedOrigin.Match([]byte(u.Hostname())) {
+		if allowedOrigin.Match([]byte(strings.ToLower(u.Hostname()))) {
 			return true
 		}
 	}
diff --git a/api/stream/stream_test.go b/api/stream/stream_test.go
index acc159e..758c4f2 100644
--- a/api/stream/stream_test.go
+++ b/api/stream/stream_test.go
@@ -408,6 +408,14 @@ func Test_sameOrigin_returnsTrue(t *testing.T) {
 	assert.True(t, actual)
 }
 
+func Test_sameOrigin_returnsTrue_withCustomPort(t *testing.T) {
+	mode.Set(mode.Prod)
+	req := httptest.NewRequest("GET", "http://example.com:8080/stream", nil)
+	req.Header.Set("Origin", "http://example.com:8080")
+	actual := isAllowedOrigin(req, nil)
+	assert.True(t, actual)
+}
+
 func Test_isAllowedOrigin_withoutAllowedOrigins_failsWhenNotSameOrigin(t *testing.T) {
 	mode.Set(mode.Prod)
 	req := httptest.NewRequest("GET", "http://example.com/stream", nil)