Merge pull request #24 from buger/refactoring

Refactoring and speed improvements

listener/listener_test.go

@@ -5,34 +5,36 @@ import (
 	"fmt"
 	"net"
 	"testing"
+	//"time"
 )
 
 func getTCPMessage() (msg *TCPMessage) {
-	packet := &TCPPacket{Data: []byte("GET /pub/WWW/ HTTP/1.1\nHost: www.w3.org\r\n\r\n")}
+	packet1 := &TCPPacket{Data: []byte("GET /pub/WWW/ HTTP/1.1\nHost: www.w3.org\r\n\r\n")}
+	packet2 := &TCPPacket{Data: []byte("asd=asdasd&zxc=qwe\r\n\r\n")}
 
-	return &TCPMessage{packets: []*TCPPacket{packet}}
+	return &TCPMessage{packets: []*TCPPacket{packet1, packet2}}
 }
 
-func mockReplayServer() (listener net.Listener) {
-	listener, _ = net.Listen("tcp", "127.0.0.1:0")
+func mockServer() (replay net.Listener) {
+	replay, _ = net.Listen("tcp", "127.0.0.1:0")
 
-	Settings.ReplayAddress = listener.Addr().String()
-
-	fmt.Println(listener.Addr().String())
+	fmt.Println(replay.Addr().String())
 
 	return
 }
 
 func TestSendMessage(t *testing.T) {
-	Settings.Verbose = true
+	Settings.Verbose = false
 
-	listener := mockReplayServer()
+	replay := mockServer()
+
+	Settings.ReplayAddress = replay.Addr().String()
 
 	msg := getTCPMessage()
 
 	sendMessage(msg)
 
-	conn, _ := listener.Accept()
+	conn, _ := replay.Accept()
 	defer conn.Close()
 
 	buf := make([]byte, 1024)
@@ -43,3 +45,37 @@ func TestSendMessage(t *testing.T) {
 		t.Errorf("Original and received requests does not match")
 	}
 }
+
+/*
+func TestPerformance(t *testing.T) {
+	Settings.Verbose = true
+
+	replay := mockReplayServer()
+
+	msg := getTCPMessage()
+
+	for y := 0; y < 10; y++ {
+		go func() {
+			for {
+				conn, _ := replay.Accept()
+
+				go func() {
+					buf := make([]byte, 1024)
+					n, _ := conn.Read(buf)
+					buf = buf[0:n]
+
+					conn.Close()
+				}()
+			}
+		}()
+	}
+
+	for y := 0; y < 10; y++ {
+		for i := 0; i < 500; i++ {
+			go sendMessage(msg)
+		}
+
+		time.Sleep(time.Millisecond * 1000)
+	}
+}
+*/

listener/raw_tcp_listener.go

@@ -14,7 +14,7 @@ import (
 // Since we can't use default TCP libraries RAWTCPLitener implements own TCP layer
 // TCP packets is parsed using tcp_packet.go, and flow control is managed by tcp_message.go
 type RAWTCPListener struct {
-	messages []*TCPMessage // buffer of TCPMessages waiting to be send
+	messages map[uint32]*TCPMessage // buffer of TCPMessages waiting to be send
 
 	c_packets  chan *TCPPacket
 	c_messages chan *TCPMessage // Messages ready to be send to client
@@ -29,9 +29,10 @@ type RAWTCPListener struct {
 func RAWTCPListen(addr string, port int) (listener *RAWTCPListener) {
 	listener = &RAWTCPListener{}
 
-	listener.c_packets = make(chan *TCPPacket)
-	listener.c_messages = make(chan *TCPMessage)
-	listener.c_del_message = make(chan *TCPMessage)
+	listener.c_packets = make(chan *TCPPacket, 100)
+	listener.c_messages = make(chan *TCPMessage, 100)
+	listener.c_del_message = make(chan *TCPMessage, 100)
+	listener.messages = make(map[uint32]*TCPMessage)
 
 	listener.addr = addr
 	listener.port = port
@@ -48,7 +49,7 @@ func (t *RAWTCPListener) listen() {
 		// If message ready for deletion it means that its also complete or expired by timeout
 		case message := <-t.c_del_message:
 			t.c_messages <- message
-			t.deleteMessage(message)
+			delete(t.messages, message.Ack)
 
 		// We need to use channels to process each packet to avoid data races
 		case packet := <-t.c_packets:
@@ -57,32 +58,6 @@ func (t *RAWTCPListener) listen() {
 	}
 }
 
-// Deleting messages that came from t.c_del_message channel
-func (t *RAWTCPListener) deleteMessage(message *TCPMessage) bool {
-	idx := -1
-
-	// Searching for given message in messages buffer
-	for i, m := range t.messages {
-		if m.Ack == message.Ack {
-			idx = i
-			break
-		}
-	}
-
-	if idx == -1 {
-		return false
-	}
-
-	// Delete element from array
-	// Note: that this version for arrays that consist of pointers
-	// https://code.google.com/p/go-wiki/wiki/SliceTricks
-	copy(t.messages[idx:], t.messages[idx+1:])
-	t.messages[len(t.messages)-1] = nil // Ensure that value will be garbage-collected.
-	t.messages = t.messages[:len(t.messages)-1]
-
-	return true
-}
-
 func (t *RAWTCPListener) readRAWSocket() {
 	conn, e := net.ListenPacket("ip4:tcp", t.addr)
 	defer conn.Close()
@@ -99,56 +74,57 @@ func (t *RAWTCPListener) readRAWSocket() {
 
 		if err != nil {
 			Debug("Error:", err)
+			continue
 		}
 
 		if n > 0 {
-			// To avoid full packet parsing every time, we manually parsing values needed for packet filtering
-			// http://en.wikipedia.org/wiki/Transmission_Control_Protocol
-			dest_port := binary.BigEndian.Uint16(buf[2:4])
-
-			// Because RAW_SOCKET can't be bound to port, we have to control it by ourself
-			if int(dest_port) == t.port {
-				// Check TCPPacket code for more description
-				flags := binary.BigEndian.Uint16(buf[12:14]) & 0x1FF
-				f_psh := (flags & TCP_PSH) != 0
-
-				// We need only packets with data inside
-				// TCP PSH flag indicate that packet have data inside
-				if f_psh {
-					// We should create new buffer because go slices is pointers. So buffer data shoud be immutable.
-					new_buf := make([]byte, n)
-					copy(new_buf, buf[:n])
-
-					// To avoid socket locking processing packet in new goroutine
-					go func(buf []byte) {
-						packet := NewTCPPacket(new_buf)
-						t.c_packets <- packet
-					}(new_buf)
-				}
-			}
+			t.parsePacket(buf[:n])
 		}
 	}
 }
 
+func (t *RAWTCPListener) parsePacket(buf []byte) {
+	if t.isIncomingDataPacket(buf) {
+		new_buf := make([]byte, len(buf))
+		copy(new_buf, buf)
+
+		t.c_packets <- ParseTCPPacket(new_buf)
+	}
+}
+
+func (t *RAWTCPListener) isIncomingDataPacket(buf []byte) bool {
+	// To avoid full packet parsing every time, we manually parsing values needed for packet filtering
+	// http://en.wikipedia.org/wiki/Transmission_Control_Protocol
+	dest_port := binary.BigEndian.Uint16(buf[2:4])
+
+	// Because RAW_SOCKET can't be bound to port, we have to control it by ourself
+	if int(dest_port) == t.port {
+		// Check TCPPacket code for more description
+		flags := binary.BigEndian.Uint16(buf[12:14]) & 0x1FF
+
+		// We need only packets with data inside
+		// TCP PSH flag indicate that packet have data inside
+		if (flags & TCP_PSH) != 0 {
+			// We should create new buffer because go slices is pointers. So buffer data shoud be immutable.
+			return true
+		}
+	}
+
+	return false
+}
+
 // Trying to add packet to existing message or creating new message
 //
 // For TCP message unique id is Acknowledgment number (see tcp_packet.go)
 func (t *RAWTCPListener) processTCPPacket(packet *TCPPacket) {
 	var message *TCPMessage
 
-	// Searching for message with same Ack
-	for _, msg := range t.messages {
-		if msg.Ack == packet.Ack {
-			message = msg
-			break
-		}
-	}
+	message, ok := t.messages[packet.Ack]
 
-	if message == nil {
+	if !ok {
 		// We sending c_del_message channel, so message object can communicate with Listener and notify it if message completed
 		message = NewTCPMessage(packet.Ack, t.c_del_message)
-
-		t.messages = append(t.messages, message)
+		t.messages[packet.Ack] = message
 	}
 
 	// Adding packet to message

listener/raw_tcp_listener_test.go

@@ -0,0 +1,84 @@
+package listener
+
+import (
+	"encoding/binary"
+	"math/rand"
+	"net"
+	"strconv"
+	"sync"
+	"testing"
+)
+
+func createHeader(ack uint32, port int) (header []byte, o_ack uint32) {
+	if ack == 0 {
+		ack = rand.Uint32()
+	}
+
+	seq := rand.Uint32()
+
+	header = make([]byte, 256)
+
+	binary.BigEndian.PutUint16(header[2:4], uint16(port))
+	binary.BigEndian.PutUint32(header[4:8], seq)
+	binary.BigEndian.PutUint32(header[8:12], ack)
+	header[12] = 4 << 4
+	header[13] = 8 // Setting PSH flag
+
+	return header, ack
+}
+
+func getPackets(port int) [][]byte {
+	if rand.Int()%2 == 0 {
+		tcp, _ := createHeader(uint32(0), port)
+		tcp = append(tcp, []byte("GET /pub/WWW/ HTTP/1.1\nHost: www.w3.org\r\n\r\n")...)
+
+		return [][]byte{tcp}
+	} else {
+		tcp1, ack := createHeader(uint32(0), port)
+		tcp1 = append(tcp1, []byte("POST /pub/WWW/ HTTP/1.1\nHost: www.w3.org\r\n\r\n")...)
+
+		tcp2, _ := createHeader(ack, port)
+		tcp2 = append(tcp2, []byte("a=1&b=2\r\n\r\n")...)
+
+		return [][]byte{tcp1, tcp2}
+	}
+
+}
+
+func TestRawTCPListener(t *testing.T) {
+	Settings.Verbose = true
+
+	server := mockServer()
+	//server_addr := server.Addr().String()
+	host, port_str, _ := net.SplitHostPort(server.Addr().String())
+	port, _ := strconv.Atoi(port_str)
+
+	// Accept all packets
+	go func() {
+		conn, _ := server.Accept()
+		conn.Close()
+	}()
+
+	listener := RAWTCPListen(host, port)
+
+	var wg sync.WaitGroup
+
+	go func() {
+		for {
+			listener.Receive()
+			wg.Done()
+		}
+	}()
+
+	for i := 0; i < 10000; i++ {
+		wg.Add(1)
+
+		packets := getPackets(port)
+
+		for _, packet := range packets {
+			listener.parsePacket(packet)
+		}
+	}
+
+	wg.Wait()
+}

listener/tcp_message.go

@@ -19,10 +19,7 @@ type TCPMessage struct {
 
 	timer *time.Timer // Used for expire check
 
-	expired bool
-
 	c_packets chan *TCPPacket
-	c_closing chan int
 
 	c_del_message chan *TCPMessage
 }
@@ -32,7 +29,6 @@ func NewTCPMessage(Ack uint32, c_del chan *TCPMessage) (msg *TCPMessage) {
 	msg = &TCPMessage{Ack: Ack}
 
 	msg.c_packets = make(chan *TCPPacket)
-	msg.c_closing = make(chan int)
 	msg.c_del_message = c_del // used for notifying that message completed or expired
 
 	// Every time we receive packet we reset this timer
@@ -46,19 +42,21 @@ func NewTCPMessage(Ack uint32, c_del chan *TCPMessage) (msg *TCPMessage) {
 func (t *TCPMessage) listen() {
 	for {
 		select {
-		case <-t.c_closing:
-			close(t.c_packets)
-			return // Stop loop if message completed/expired
-		case packet := <-t.c_packets:
-			t.AddPacket(packet)
+		case packet, more := <-t.c_packets:
+			if more {
+				t.AddPacket(packet)
+			} else {
+				// Stop loop if channel closed
+				return
+			}
 		}
 	}
 }
 
 // Timeout notifies message to stop listening, close channel and message ready to be sent
 func (t *TCPMessage) Timeout() {
-	t.c_closing <- 1     // Notify to stop listen loop and close channel
-	t.c_del_message <- t // Notify RAWListener that message is ready to be sent to replay server
+	close(t.c_packets)   // Notify to stop listen loop and close channel
+	t.c_del_message <- t // Notify RAWListener that message is ready to be send to replay server
 }
 
 // Bytes sorts packets in right orders and return message content
@@ -83,11 +81,6 @@ func (t *TCPMessage) Bytes() (output []byte) {
 // AddPacket to the message and ensure packet uniqueness
 // TCP allows that packet can be re-send multiple times
 func (t *TCPMessage) AddPacket(packet *TCPPacket) {
-	if t.expired {
-		Debug("Adding packet to expired message")
-		return
-	}
-
 	packetFound := false
 
 	for _, pkt := range t.packets {

listener/tcp_packet.go

@@ -36,12 +36,11 @@ type TCPPacket struct {
 	Data []byte
 }
 
-// NewTCPPacket pointer from data in []byte
-func NewTCPPacket(b []byte) (t *TCPPacket) {
-	t = &TCPPacket{Data: b}
-	t.ParseBasic()
+func ParseTCPPacket(b []byte) (p *TCPPacket) {
+	p = &TCPPacket{Data: b}
+	p.ParseBasic()
 
-	return t
+	return p
 }
 
 // Parse TCP Packet, inspired by: https://github.com/miekg/pcap/blob/master/packet.go