diff --git a/go.mod b/go.mod index 0ffd1e2..6a46cb2 100644 --- a/go.mod +++ b/go.mod @@ -2,4 +2,4 @@ module github.com/cryptexctl/gbio go 1.24.4 -require github.com/gorilla/websocket v1.5.3 // indirect +require github.com/gorilla/websocket v1.5.3 diff --git a/main.go b/main.go index b58485b..f06dfbd 100644 --- a/main.go +++ b/main.go @@ -25,6 +25,53 @@ type KeyPress struct { Key string `json:"key"` } +type RateLimiter struct { + connections map[string]int + messages map[string]time.Time + mutex sync.RWMutex +} + +var limiter = &RateLimiter{ + connections: make(map[string]int), + messages: make(map[string]time.Time), +} + +var trollingGifBase64 string +var trollingHTML string +var trollingHTMLOnce sync.Once + +func (rl *RateLimiter) canConnect(ip string) bool { + rl.mutex.Lock() + defer rl.mutex.Unlock() + + if rl.connections[ip] >= 3 { + return false + } + rl.connections[ip]++ + return true +} + +func (rl *RateLimiter) canSendMessage(ip string) bool { + rl.mutex.Lock() + defer rl.mutex.Unlock() + + lastMessage, exists := rl.messages[ip] + if exists && time.Since(lastMessage) < 30*time.Millisecond { + return false + } + rl.messages[ip] = time.Now() + return true +} + +func (rl *RateLimiter) disconnect(ip string) { + rl.mutex.Lock() + defer rl.mutex.Unlock() + + if rl.connections[ip] > 0 { + rl.connections[ip]-- + } +} + type gzipResponseWriter struct { io.Writer http.ResponseWriter @@ -50,6 +97,21 @@ func gzipHandler(next http.HandlerFunc) http.HandlerFunc { } func wsHandler(w http.ResponseWriter, r *http.Request) { + ip := strings.Split(r.RemoteAddr, ":")[0] + + if !limiter.canConnect(ip) { + conn, err := upgrader.Upgrade(w, r, nil) + if err != nil { + log.Printf("WebSocket upgrade error: %v", err) + return + } + conn.WriteJSON(map[string]string{"action": "show_message", "message": "You make too many requests, yk..."}) + conn.Close() + return + } + + defer limiter.disconnect(ip) + conn, err := upgrader.Upgrade(w, r, nil) if err != nil { log.Printf("WebSocket upgrade error: %v", err) @@ -67,6 +129,11 @@ func wsHandler(w http.ResponseWriter, r *http.Request) { break } + if !limiter.canSendMessage(ip) { + conn.WriteJSON(map[string]string{"action": "show_message", "message": "You make too many requests, yk..."}) + continue + } + now := time.Now() if now.Sub(lastPress) > 5*time.Second { buffer = "" @@ -140,19 +207,15 @@ func wsHandler(w http.ResponseWriter, r *http.Request) { } } -func trollingHandler(w http.ResponseWriter, r *http.Request) { - w.Header().Set("Content-Type", "text/html; charset=utf-8") - w.Header().Set("Server", "0BSD_FOR_EVERYONE") - +func initTrollingHTML() { fileContent, err := os.ReadFile("files/trolling.gif") if err != nil { - log.Printf("Ошибка чтения файла: %v", err) - fileContent = []byte("Файл не найден") + log.Printf("Error reading file: %v", err) + fileContent = []byte("File not found") } - base64Content := base64.StdEncoding.EncodeToString(fileContent) - - html := fmt.Sprintf(` + trollingGifBase64 = base64.StdEncoding.EncodeToString(fileContent) + trollingHTML = fmt.Sprintf(` @@ -177,9 +240,16 @@ func trollingHandler(w http.ResponseWriter, r *http.Request) { мы делаем небольшой троллинг -`, base64Content) +`, trollingGifBase64) +} - fmt.Fprint(w, html) +func trollingHandler(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "text/html; charset=utf-8") + w.Header().Set("Server", "0BSD_FOR_EVERYONE") + w.Header().Set("Cache-Control", "public, max-age=3600") + + trollingHTMLOnce.Do(initTrollingHTML) + fmt.Fprint(w, trollingHTML) } func statusHandler(w http.ResponseWriter, r *http.Request) { @@ -297,9 +367,9 @@ func main() { IdleTimeout: 60 * time.Second, } - log.Printf("Сервер запущен на порту %s", p) + log.Printf("Server started on port %s", p) if err := server.ListenAndServe(); err != nil { - log.Printf("Ошибка на порту %s: %v", p, err) + log.Printf("Error on port %s: %v", p, err) } }(port) }