Modified communication format between server and client.

1 files changed, 62 insertions, 14 deletions

diff --git a/server/main.go b/server/main.go
index ec848ad..7c67ead 100644
--- a/server/main.go
+++ b/server/main.go
@@ -2,7 +2,6 @@ package main
 
 import (
 	"encoding/json"
-	"fmt"
 	"log"
 	"net"
 )
@@ -36,13 +35,49 @@ func parseMessage(data []byte) Message {
 	return m
 }
 
+func toBytes(num uint32) []byte {
+	buf := make([]byte, 4)
+	var mask uint32 = 0x00FF
+	for i := 0; i < 4; i += 1 {
+		buf[4-i-1] = byte(num & mask)
+		num >>= 8
+	}
+	return buf
+}
+
+func fromBytes(data []byte) uint32 {
+	var num uint32
+	for i := 0; i < 4; i += 1 {
+		num += uint32(data[i] << ((4 - i - 1) * 8))
+	}
+	return num
+}
+
+func sendMessage(conn net.Conn, message []byte) error {
+	size := uint32(len(message))
+
+	_, err := conn.Write(toBytes(size))
+	if err != nil {
+		return err
+	}
+
+	_, err = conn.Write(message)
+	if err != nil {
+		return err
+	}
+
+	return nil
+}
+
 func handleConnection(conn net.Conn) {
 	defer conn.Close()
 
+	// 1 мегабайт
+	var maxBodySize uint32 = 1048576
 	maxAuthAttempts := 5
 	authAttempts := 0
 	authenticated := false
-	buf := make([]byte, 4096)
+	buf := make([]byte, 256)
 
 	for {
 		if authAttempts >= maxAuthAttempts {
@@ -50,16 +85,30 @@ func handleConnection(conn net.Conn) {
 				Status: FAILURE,
 				Data:   "Превышено количество неавторизованных запросов",
 			})
-			_, _ = conn.Write(resp)
+
+			_ = sendMessage(conn, resp)
 			break
 		}
 
-		readLen, err := conn.Read(buf)
+		_, err := conn.Read(buf)
 		if err != nil {
 			break
 		}
 
-		message := parseMessage(buf[:readLen])
+		size := fromBytes(buf[:4])
+		if size > maxBodySize {
+			resp, _ := json.Marshal(Response{
+				Status: FAILURE,
+				Data:   "Превышен максимальный размер запроса",
+			})
+
+			_ = sendMessage(conn, resp)
+			break
+		}
+
+		rawMessage := buf[4 : size+4]
+		message := parseMessage(rawMessage)
+		//fmt.Println(message.Data)
 		if message.Type == AUTHENTICATE {
 			// TODO(andrew): Добавить процесс аутентификации
 			if true {
@@ -72,28 +121,26 @@ func handleConnection(conn net.Conn) {
 				Status: FAILURE,
 				Data:   "Для работы с сервером необходима аутентификация",
 			})
-			_, err = conn.Write(resp)
+
+			err = sendMessage(conn, resp)
 			if err != nil {
-				conn.Close()
 				break
-			} else {
-				authAttempts += 1
-				continue
 			}
+
+			authAttempts += 1
+			continue
 		}
 
 		if message.Data == "exit" {
 			break
 		}
 
-		fmt.Println(message.Data)
 		resp, _ := json.Marshal(Response{
 			Status: SUCCESS,
 			Data:   "some data",
 		})
-		_, err = conn.Write(resp)
+		err = sendMessage(conn, resp)
 		if err != nil {
-			err = conn.Close()
 			break
 		}
 	}
@@ -101,6 +148,7 @@ func handleConnection(conn net.Conn) {
 
 func main() {
 	sock, err := net.Listen("tcp", "localhost:8080")
+	defer sock.Close()
 	handleError(err)
 
 	for {
@@ -108,7 +156,7 @@ func main() {
 		if err == nil {
 			go handleConnection(conn)
 		} else {
-			break
+			handleError(err)
 		}
 	}
 }