From 97ad45e505c696ab2bd34ba748c33efdbcabd540 Mon Sep 17 00:00:00 2001
From: Marco <mgross@sw-gross.de>
Date: Sun, 12 May 2024 15:42:40 +0200
Subject: [PATCH] Rate limit websocket connection

---
 api/{websocket/connection.go => handler/websocket.go} | 7 ++++++-
 main.go                                               | 5 ++---
 2 files changed, 8 insertions(+), 4 deletions(-)
 rename api/{websocket/connection.go => handler/websocket.go} (96%)

diff --git a/api/websocket/connection.go b/api/handler/websocket.go
similarity index 96%
rename from api/websocket/connection.go
rename to api/handler/websocket.go
index 1f18929..5ca8d46 100644
--- a/api/websocket/connection.go
+++ b/api/handler/websocket.go
@@ -1,4 +1,4 @@
-package websocket
+package handler
 
 import (
 	"context"
@@ -21,6 +21,8 @@ var upgrader = gorillaws.Upgrader{
 }
 
 func RegisterWebSocketConnection(c *gin.Context) {
+	limiter.Take()
+
 	log.Println(c.Request)
 	conn, err := upgrader.Upgrade(c.Writer, c.Request, nil)
 	if err != nil {
@@ -31,6 +33,8 @@ func RegisterWebSocketConnection(c *gin.Context) {
 }
 
 func waitForAndHandlePlayerID(ctx context.Context, conn *gorillaws.Conn) {
+	limiter.Take()
+
 	msgType, msg, err := conn.ReadMessage()
 	if err != nil {
 		errorMessage := fmt.Sprintf("Reading from websocket connection did not work: %s", err)
@@ -72,5 +76,6 @@ func waitForAndHandlePlayerID(ctx context.Context, conn *gorillaws.Conn) {
 }
 
 func ConnectWsForGame(c *gin.Context) {
+	limiter.Take()
 
 }
diff --git a/main.go b/main.go
index ab54f4a..87ebd93 100644
--- a/main.go
+++ b/main.go
@@ -4,7 +4,6 @@ import (
 	"flag"
 	"log"
 	"mchess_server/api/handler"
-	"mchess_server/api/websocket"
 
 	"github.com/gin-gonic/gin"
 )
@@ -26,10 +25,10 @@ func main() {
 	router := gin.Default()
 	router.GET("/api/hostPrivate", handler.HostGameHandler)
 	router.POST("/api/joinPrivate", handler.JoinPrivateGame)
-	router.GET("/api/ws", websocket.RegisterWebSocketConnection)
+	router.GET("/api/ws", handler.RegisterWebSocketConnection)
 
 	router.GET("/api/getLobbyForPassphrase/:phrase", handler.GetLobbyForPassphraseHandler)
-	router.GET("/api/registerWsForGame/:id", websocket.ConnectWsForGame)
+	router.GET("/api/registerWsForGame/:id", handler.ConnectWsForGame)
 	router.POST("/api/joinGame/:id", handler.JoinGameHandler)
 
 	if debugMode {