Merge pull request #119 from LiquidTheDangerous/buffer-size/main

perf: reuse allocated buffer. add config parameter

packet/http.go

@@ -2,6 +2,7 @@ package packet
 
 import (
 	"bufio"
+	"io"
 	"net"
 	"net/http"
 	"strings"
@@ -56,12 +57,16 @@ func ParseUrl(raw []byte) {
 
 }
 
-func NewHttpPacket(raw []byte) (*HttpPacket, error) {
+func NewHttpPacketFromReader(rdr io.Reader) (*HttpPacket, error) {
-	pkt := &HttpPacket{raw: raw}
+	sb := strings.Builder{}
-
+	tee := io.TeeReader(rdr, &sb)
-	pkt.parse()
+	p := &HttpPacket{}
-
+	err := parse(p, bufio.NewReader(tee))
-	return pkt, nil
+	if err != nil {
+		return nil, err
+	}
+	p.raw = []byte(sb.String())
+	return p, nil
 }
 
 func (p *HttpPacket) Raw() []byte {
@@ -123,8 +128,7 @@ func (p *HttpPacket) Tidy() {
 	p.raw = []byte(result)
 }
 
-func (p *HttpPacket) parse() error {
+func parse(p *HttpPacket, reader *bufio.Reader) error {
-	reader := bufio.NewReader(strings.NewReader(string(p.raw)))
 	request, err := http.ReadRequest(reader)
 	if err != nil {
 		return err
@@ -152,6 +156,5 @@ func (p *HttpPacket) parse() error {
 	}
 
 	request.Body.Close()
-
 	return nil
 }

proxy/http.go

@@ -39,7 +39,7 @@ func (pxy *Proxy) handleHttp(lConn *net.TCPConn, pkt *packet.HttpPacket, ip stri
 
 	log.Debug("[HTTP] New connection to the server ", pkt.Domain(), " ", rConn.LocalAddr())
 
-	go Serve(rConn, lConn, "[HTTP]", lConn.RemoteAddr().String(), pkt.Domain(), pxy.timeout)
+	go Serve(rConn, lConn, "[HTTP]", lConn.RemoteAddr().String(), pkt.Domain(), pxy.timeout, pxy.bufferSize)
 
 	_, err = rConn.Write(pkt.Raw())
 	if err != nil {
@@ -49,5 +49,5 @@ func (pxy *Proxy) handleHttp(lConn *net.TCPConn, pkt *packet.HttpPacket, ip stri
 
 	log.Debug("[HTTP] Sent a request to ", pkt.Domain())
 
-	Serve(lConn, rConn, "[HTTP]", lConn.RemoteAddr().String(), pkt.Domain(), pxy.timeout)
+	Serve(lConn, rConn, "[HTTP]", lConn.RemoteAddr().String(), pkt.Domain(), pxy.timeout, pxy.bufferSize)
 }

proxy/https.go

@@ -45,7 +45,8 @@ func (pxy *Proxy) handleHttps(lConn *net.TCPConn, exploit bool, initPkt *packet.
 	log.Debug("[HTTPS] Sent 200 Connection Estabalished to ", lConn.RemoteAddr())
 
 	// Read client hello
-	clientHello, err := ReadBytes(lConn)
+	tmpBuffer := make([]byte, pxy.bufferSize)
+	clientHello, err := ReadBytes(lConn, tmpBuffer)
 	if err != nil {
 		log.Debug("[HTTPS] Error reading client hello from the client", err)
 		return
@@ -60,7 +61,7 @@ func (pxy *Proxy) handleHttps(lConn *net.TCPConn, exploit bool, initPkt *packet.
 	// lConn.SetLinger(3)
 	// rConn.SetLinger(3)
 
-	go Serve(rConn, lConn, "[HTTPS]", rConn.RemoteAddr().String(), initPkt.Domain(), pxy.timeout)
+	go Serve(rConn, lConn, "[HTTPS]", rConn.RemoteAddr().String(), initPkt.Domain(), pxy.timeout, pxy.bufferSize)
 
 	if exploit {
 		log.Debug("[HTTPS] Writing chunked client hello to ", initPkt.Domain())
@@ -77,7 +78,7 @@ func (pxy *Proxy) handleHttps(lConn *net.TCPConn, exploit bool, initPkt *packet.
 		}
 	}
 
-	Serve(lConn, rConn, "[HTTPS]", lConn.RemoteAddr().String(), initPkt.Domain(), pxy.timeout)
+	Serve(lConn, rConn, "[HTTPS]", lConn.RemoteAddr().String(), initPkt.Domain(), pxy.timeout, pxy.bufferSize)
 }
 
 func splitInChunks(bytes []byte, size int) [][]byte {

proxy/io.go

@@ -9,8 +9,6 @@ import (
 	log "github.com/sirupsen/logrus"
 )
 
-const BUF_SIZE = 1024
-
 func WriteChunks(conn *net.TCPConn, c [][]byte) (n int, err error) {
 	total := 0
 	for i := 0; i < len(c); i++ {
@@ -25,37 +23,27 @@ func WriteChunks(conn *net.TCPConn, c [][]byte) (n int, err error) {
 	return total, nil
 }
 
-func ReadBytes(conn *net.TCPConn) ([]byte, error) {
+func ReadBytes(conn *net.TCPConn, dest []byte) ([]byte, error) {
-	ret := make([]byte, 0)
+	n, err := readBytesInternal(conn, dest)
-	buf := make([]byte, BUF_SIZE)
+	return dest[:n], err
+}
 
-	for {
+func readBytesInternal(conn *net.TCPConn, dest []byte) (int, error) {
-		n, err := conn.Read(buf)
+	totalRead, err := conn.Read(dest)
 	if err != nil {
 		switch err.(type) {
 		case *net.OpError:
-				return nil, errors.New("timed out")
+			return totalRead, errors.New("timed out")
 		default:
-				return nil, err
+			return totalRead, err
 		}
 	}
-		ret = append(ret, buf[:n]...)
+	return totalRead, nil
-
-		if n < BUF_SIZE {
-			break
-		}
-	}
-
-	if len(ret) == 0 {
-		return nil, io.EOF
-	}
-
-	return ret, nil
 }
 
-func Serve(from *net.TCPConn, to *net.TCPConn, proto string, fd string, td string, timeout int) {
+func Serve(from *net.TCPConn, to *net.TCPConn, proto string, fd string, td string, timeout int, bufferSize int) {
 	proto += " "
-
+	buf := make([]byte, bufferSize)
 	for {
 		if timeout > 0 {
 			from.SetReadDeadline(
@@ -63,7 +51,7 @@ func Serve(from *net.TCPConn, to *net.TCPConn, proto string, fd string, td strin
 			)
 		}
 
-		buf, err := ReadBytes(from)
+		bytesRead, err := ReadBytes(from, buf)
 		if err != nil {
 			if err == io.EOF {
 				log.Debug(proto, "Finished ", fd)
@@ -73,7 +61,7 @@ func Serve(from *net.TCPConn, to *net.TCPConn, proto string, fd string, td strin
 			return
 		}
 
-		if _, err := to.Write(buf); err != nil {
+		if _, err := to.Write(bytesRead); err != nil {
 			log.Debug(proto, "Error Writing to ", td)
 			return
 		}

proxy/proxy.go

@@ -13,6 +13,8 @@ import (
 	"github.com/xvzc/SpoofDPI/util"
 )
 
+const BUFFER_SIZE = 1024
+
 type Proxy struct {
 	addr           string
 	port           int
@@ -20,6 +22,7 @@ type Proxy struct {
 	resolver       *dns.DnsResolver
 	windowSize     int
 	allowedPattern []*regexp.Regexp
+	bufferSize     int
 }
 
 func New(config *util.Config) *Proxy {
@@ -30,6 +33,7 @@ func New(config *util.Config) *Proxy {
 		windowSize:     *config.WindowSize,
 		allowedPattern: config.AllowedPattern,
 		resolver:       dns.NewResolver(config),
+		bufferSize:     BUFFER_SIZE,
 	}
 }
 
@@ -57,16 +61,15 @@ func (pxy *Proxy) Start() {
 		}
 
 		go func() {
-			b, err := ReadBytes(conn.(*net.TCPConn))
+			pkt, err := packet.NewHttpPacketFromReader(conn)
 			if err != nil {
 				return
 			}
 
-			log.Debug("[PROXY] Request from ", conn.RemoteAddr(), "\n\n", string(b))
+			log.Debug("[PROXY] Request from ", conn.RemoteAddr(), "\n\n", string(pkt.Raw()))
 
-			pkt, err := packet.NewHttpPacket(b)
 			if err != nil {
-				log.Debug("[PROXY] Error while parsing request: ", string(b))
+				log.Debug("[PROXY] Error while parsing request: ", string(pkt.Raw()))
 				conn.Close()
 				return
 			}

util/config.go

@@ -58,7 +58,6 @@ when not given, the client hello packet will be sent in two parts:
 fragmentation for the first data packet and the rest
 `)
 	config.Version = flag.Bool("v", false, "print spoof-dpi's version; this may contain some other relevant information")
-
 	var allowedPattern StringArray
 	flag.Var(
 		&allowedPattern,