From df7cba081eecab6fddf834e542b926218fb3d818 Mon Sep 17 00:00:00 2001
From: Andrei Mihu <andrei@heroiclabs.com>
Date: Thu, 10 Nov 2022 16:23:15 +0000
Subject: [PATCH] Allow the socket acceptor to read session tokens from request
 headers.

---
 CHANGELOG.md         |  3 +++
 server/matchmaker.go |  2 +-
 server/socket_ws.go  | 15 ++++++++++++++-
 3 files changed, 18 insertions(+), 2 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index ee04c11b3..ef7096033 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -4,6 +4,9 @@ All notable changes to this project are documented below.
 The format is based on [keep a changelog](http://keepachangelog.com) and this project uses [semantic versioning](http://semver.org).
 
 ## [Unreleased]
+### Added
+- Allow the socket acceptor to read session tokens from request headers.
+
 ### Changed
 - Stricter validation of limit in runtime storage list operations.
 - Allow subdomain variance in Facebook Limited Login token issuer field.
diff --git a/server/matchmaker.go b/server/matchmaker.go
index f86d6cd04..cef50ed22 100644
--- a/server/matchmaker.go
+++ b/server/matchmaker.go
@@ -550,7 +550,7 @@ func (m *LocalMatchmaker) Process() {
 					// We've removed something, update the known size of the currently considered combo.
 					l = len(foundCombo) + index.Count
 
-					if (len(foundCombo)+index.Count)%index.CountMultiple != 0 {
+					if l%index.CountMultiple != 0 {
 						// Removal was insufficient, the combo is still not valid for the required multiple.
 						continue
 					}
diff --git a/server/socket_ws.go b/server/socket_ws.go
index 2a2396e1b..40ed50eab 100644
--- a/server/socket_ws.go
+++ b/server/socket_ws.go
@@ -18,6 +18,7 @@ import (
 	"net"
 	"net/http"
 	"strconv"
+	"strings"
 
 	"github.com/gofrs/uuid"
 	"github.com/gorilla/websocket"
@@ -55,7 +56,19 @@ func NewSocketWsAcceptor(logger *zap.Logger, config Config, sessionRegistry Sess
 		}
 
 		// Check authentication.
-		token := r.URL.Query().Get("token")
+		var token string
+		if auth := r.Header["Authorization"]; len(auth) >= 1 {
+			// Attempt header based authentication.
+			const prefix = "Bearer "
+			if !strings.HasPrefix(auth[0], prefix) {
+				http.Error(w, "Missing or invalid token", 401)
+				return
+			}
+			token = auth[0][len(prefix):]
+		} else {
+			// Attempt query parameter based authentication.
+			token = r.URL.Query().Get("token")
+		}
 		if token == "" {
 			http.Error(w, "Missing or invalid token", 401)
 			return
-- 
GitLab