From 72d0191c98aa19925e9172d32fe8d289aaf849be Mon Sep 17 00:00:00 2001
From: zijiren233 <pyh1670605849@gmail.com>
Date: Mon, 9 Oct 2023 23:08:40 +0800
Subject: [PATCH] Feat: Message length verification

---
 room/room.go                 |  3 +++
 room/user.go                 | 10 ++++++++--
 server/handlers/websocket.go |  7 +++++++
 3 files changed, 18 insertions(+), 2 deletions(-)

diff --git a/room/room.go b/room/room.go
index 0130496..acdffc2 100644
--- a/room/room.go
+++ b/room/room.go
@@ -16,6 +16,7 @@ import (
 
 var (
 	ErrRoomIDEmpty        = errors.New("roomid is empty")
+	ErrRoomIDTooLong      = errors.New("roomid is too long")
 	ErrAdminPassWordEmpty = errors.New("admin password is empty")
 )
 
@@ -65,6 +66,8 @@ func WithRootUser(u *User) RoomConf {
 func NewRoom(RoomID string, Password string, rtmps *rtmps.Server, conf ...RoomConf) (*Room, error) {
 	if RoomID == "" {
 		return nil, ErrRoomIDEmpty
+	} else if len(RoomID) > 32 {
+		return nil, ErrRoomIDTooLong
 	}
 	now := time.Now().UnixMilli()
 	r := &Room{
diff --git a/room/user.go b/room/user.go
index c1a9d87..044dcda 100644
--- a/room/user.go
+++ b/room/user.go
@@ -22,8 +22,10 @@ type User struct {
 }
 
 var (
-	ErrUsernameEmpty     = errors.New("user name is empty")
-	ErrUserPasswordEmpty = errors.New("user password is empty")
+	ErrUsernameEmpty       = errors.New("user name is empty")
+	ErrUsernameTooLong     = errors.New("user name is too long")
+	ErrUserPasswordEmpty   = errors.New("user password is empty")
+	ErrUserPasswordTooLong = errors.New("user password is too long")
 )
 
 type UserConf func(*User)
@@ -43,9 +45,13 @@ func WithUserAdmin(admin bool) UserConf {
 func NewUser(id string, password string, room *Room, conf ...UserConf) (*User, error) {
 	if id == "" {
 		return nil, ErrUsernameEmpty
+	} else if len(id) > 32 {
+		return nil, ErrUsernameTooLong
 	}
 	if password == "" {
 		return nil, ErrUserPasswordEmpty
+	} else if len(password) > 32 {
+		return nil, ErrUserPasswordTooLong
 	}
 	hashedPassword, err := bcrypt.GenerateFromPassword(stream.StringToBytes(password), bcrypt.DefaultCost)
 	if err != nil {
diff --git a/server/handlers/websocket.go b/server/handlers/websocket.go
index df97eda..00e06ea 100644
--- a/server/handlers/websocket.go
+++ b/server/handlers/websocket.go
@@ -122,6 +122,13 @@ func handleReaderMessage(c *room.Client) error {
 			}
 			switch msg.Type {
 			case room.ChatMessage:
+				if len(msg.Message) > 4096 {
+					c.Send(&room.ElementMessage{
+						Type:    room.Error,
+						Message: "message too long",
+					})
+					continue
+				}
 				c.Broadcast(&room.ElementMessage{
 					Type:    room.ChatMessage,
 					Sender:  c.Username(),