From da6f5dfb85379870be1e922f81f4748e5a83dbc3 Mon Sep 17 00:00:00 2001
From: Arpit Bhayani <arpit.b.bhayani@gmail.com>
Date: Wed, 26 Oct 2022 20:41:43 +0530
Subject: [PATCH] prevented repeated allocations of temp buffer

---
 config/main.go |  1 +
 core/io.go     | 38 ++++++++++++++++++++++++--------------
 2 files changed, 25 insertions(+), 14 deletions(-)

diff --git a/config/main.go b/config/main.go
index 62a387d1d..ce8d56cf5 100644
--- a/config/main.go
+++ b/config/main.go
@@ -13,3 +13,4 @@ var AOFFile string = "./dice-master.aof"
 
 // Network
 var IOBufferLength int = 512
+var IOBufferLengthMAX int = 50 * 1024
diff --git a/core/io.go b/core/io.go
index 6956ba8c4..bf9dc3052 100644
--- a/core/io.go
+++ b/core/io.go
@@ -3,14 +3,17 @@ package core
 import (
 	"bytes"
 	"errors"
+	"fmt"
 	"io"
+	"log"
 
 	"github.com/dicedb/dice/config"
 )
 
 type RESPParser struct {
-	c   io.ReadWriter
-	buf *bytes.Buffer
+	c    io.ReadWriter
+	buf  *bytes.Buffer
+	tbuf []byte
 }
 
 func NewRESPParser(c io.ReadWriter) *RESPParser {
@@ -24,36 +27,42 @@ func NewRESPParserWithBytes(c io.ReadWriter, initBytes []byte) *RESPParser {
 	return &RESPParser{
 		c:   c,
 		buf: buf,
+		// assigning temporary buffer to read 512 bytes in one shot
+		// and reading them in a loop until we have all the data
+		// we want.
+		// note: the size 512 is arbitrarily chosen, and we can put
+		// a decent thought into deciding the optimal value (in case it affects the perf)
+		tbuf: make([]byte, config.IOBufferLength),
 	}
 }
 
 func (rp *RESPParser) DecodeOne() (interface{}, error) {
-	// assigning temporary buffer to read 512 bytes in one shot
-	// and reading them in a loop until we have all the data
-	// we want.
-	// note: the size 512 is arbitrarily chosen, and we can put
-	// a decent thought into deciding the optimal value (in case it affects the perf)
-	var buf []byte = make([]byte, config.IOBufferLength)
-
-	// keep reading the bytes from the buffer until the first /r is found
-	// note: there may be extra bytes read over the socket post /r
+	// keep reading the bytes from the buffer until the first \r is found
+	// note: there may be extra bytes read over the socket post \r
 	// but that is fine.
 	for {
-		n, err := rp.c.Read(buf)
+		n, err := rp.c.Read(rp.tbuf)
 		if n <= 0 {
 			break
 		}
 		if err != nil {
 			if err == io.EOF {
+				if n > 0 {
+					rp.buf.Write(rp.tbuf[:n])
+				}
 				break
 			}
 			return nil, err
 		}
-		rp.buf.Write(buf[:n])
+		rp.buf.Write(rp.tbuf[:n])
 
-		if bytes.Contains(buf, []byte{'\r', '\n'}) {
+		if bytes.Contains(rp.tbuf, []byte{'\r', '\n'}) {
 			break
 		}
+
+		if rp.buf.Len() > config.IOBufferLengthMAX {
+			return nil, fmt.Errorf("input too long. max input can be %d bytes", config.IOBufferLengthMAX)
+		}
 	}
 
 	b, err := rp.buf.ReadByte()
@@ -79,6 +88,7 @@ func (rp *RESPParser) DecodeOne() (interface{}, error) {
 	// not start with any of the above special chars will be a potential
 	// attack.
 	// Details: https://bou.ke/blog/hacking-developers/
+	log.Println("possible cross protocol scripting attack detected. dropping the request.")
 	return nil, errors.New("possible cross protocol scripting attack detected")
 }