Fix tests

raw_socket_listener/listener.go

@@ -36,7 +36,7 @@ type Listener struct {
 	mu sync.Mutex
 	// buffer of TCPMessages waiting to be send
 	// ID -> TCPMessage
-	messages map[string]*TCPMessage
+	messages map[[10]byte]*TCPMessage
 
 	// Expect: 100-continue request is send in 2 tcp messages
 	// We store ACK aliases to merge this packets together
@@ -48,7 +48,7 @@ type Listener struct {
 	respAliases map[uint32]*request
 
 	// Ack -> ID
-	respWithoutReq map[uint32]string
+	respWithoutReq map[uint32][10]byte
 
 	// Messages ready to be send to client
 	packetsChan chan []byte
@@ -67,7 +67,7 @@ type Listener struct {
 }
 
 type request struct {
-	id    string
+	id    [10]byte
 	start time.Time
 	ack   uint32
 }
@@ -87,11 +87,11 @@ func NewListener(addr string, port string, engine int, expire time.Duration) (l
 	l.quit = make(chan bool)
 	l.readyCh = make(chan bool, 1)
 
-	l.messages = make(map[string]*TCPMessage)
+	l.messages = make(map[[10]byte]*TCPMessage)
 	l.ackAliases = make(map[uint32]uint32)
 	l.seqWithData = make(map[uint32]uint32)
 	l.respAliases = make(map[uint32]*request)
-	l.respWithoutReq = make(map[uint32]string)
+	l.respWithoutReq = make(map[uint32][10]byte)
 
 	l.addr = addr
 	_port, _ := strconv.Atoi(port)
@@ -131,33 +131,29 @@ func (t *Listener) listen() {
 			}
 			return
 		case data := <- t.packetsChan:
-			packet := ParseTCPPacket(net.IP(data[:4]).String(), data[4:])
+			packet := ParseTCPPacket(data[:4], data[4:])
 			t.processTCPPacket(packet)
 		case <- gcTicker:
 			now := time.Now()
-			// log.Println("GC")
 
-			t.mu.Lock()
 			// Dispatch requests before responses
 			for _, message := range t.messages {
 				if now.Sub(message.End) >= t.messageExpire {
 					t.dispatchMessage(message)
 				}
 			}
-
-			t.mu.Unlock()
 		}
 	}
 }
 
 func (t *Listener) dispatchMessage(message *TCPMessage) {
 	// If already dispatched
-	if _, ok := t.messages[message.ID]; !ok {
+	if _, ok := t.messages[message.ID()]; !ok {
 		return
 	}
 
 	delete(t.ackAliases, message.Ack)
-	delete(t.messages, message.ID)
+	delete(t.messages, message.ID())
 	delete(t.respAliases, message.ResponseAck)
 
 	// log.Println("Dispatching, message", message.Seq, message.Ack, string(message.Bytes()))
@@ -391,7 +387,7 @@ func (t *Listener) processTCPPacket(packet *TCPPacket) {
 		// In case if non-first data chunks comes first
 		for _id, m := range t.messages {
 			// log.Println("Message ack:", m.Ack, m.packets[0].Addr, packet.Addr)
-			if m.Ack == packet.Ack && m.packets[0].Addr == packet.Addr {
+			if m.Ack == packet.Ack && bytes.Equal(m.packets[0].Addr, packet.Addr) {
 				delete(t.messages, _id)
 
 				for _, pkt := range m.packets {
@@ -417,13 +413,11 @@ func (t *Listener) processTCPPacket(packet *TCPPacket) {
 		responseRequest, _ = t.respAliases[packet.Ack]
 	}
 
-	mID := packet.Addr + strconv.Itoa(int(packet.DestPort)) + strconv.Itoa(int(packet.Ack))
-
-	message, ok := t.messages[mID]
+	message, ok := t.messages[packet.ID]
 
 	if !ok {
-		message = NewTCPMessage(mID, packet.Seq, packet.Ack, isIncoming)
-		t.messages[mID] = message
+		message = NewTCPMessage(packet.Seq, packet.Ack, isIncoming)
+		t.messages[packet.ID] = message
 
 		if !isIncoming {
 			if responseRequest != nil {
@@ -431,7 +425,7 @@ func (t *Listener) processTCPPacket(packet *TCPPacket) {
 				message.RequestAck = responseRequest.ack
 				message.RequestID = responseRequest.id
 			} else {
-				t.respWithoutReq[packet.Ack] = mID
+				t.respWithoutReq[packet.Ack] = packet.ID
 			}
 		}
 	}
@@ -475,7 +469,7 @@ func (t *Listener) processTCPPacket(packet *TCPPacket) {
 		}
 
 		message.UpdateResponseAck()
-		t.respAliases[message.ResponseAck] = &request{message.ID, message.Start, message.Ack}
+		t.respAliases[message.ResponseAck] = &request{message.ID(), message.Start, message.Ack}
 	}
 
 	// If message contains only single packet immediately dispatch it

raw_socket_listener/tcp_message.go

@@ -4,6 +4,7 @@ import (
 	"bytes"
 	"crypto/sha1"
 	"encoding/hex"
+	"encoding/binary"
 	"github.com/buger/gor/proto"
 	"log"
 	"strconv"
@@ -17,13 +18,12 @@ import (
 // Message can be compiled from unique packets with same message_id which sorted by sequence
 // Message is received if we didn't receive any packets for 2000ms
 type TCPMessage struct {
-	ID           string // Message ID
 	Seq          uint32
 	Ack          uint32
 	ResponseAck  uint32
 	RequestStart time.Time
 	RequestAck   uint32
-	RequestID    string
+	RequestID    [10]byte
 	Start        time.Time
 	End          time.Time
 	IsIncoming   bool
@@ -34,8 +34,8 @@ type TCPMessage struct {
 }
 
 // NewTCPMessage pointer created from a Acknowledgment number and a channel of messages readuy to be deleted
-func NewTCPMessage(ID string, Seq, Ack uint32, IsIncoming bool) (msg *TCPMessage) {
-	msg = &TCPMessage{ID: ID, Seq: Seq, Ack: Ack, IsIncoming: IsIncoming}
+func NewTCPMessage(Seq, Ack uint32, IsIncoming bool) (msg *TCPMessage) {
+	msg = &TCPMessage{Seq: Seq, Ack: Ack, IsIncoming: IsIncoming}
 	msg.Start = time.Now()
 
 	return
@@ -196,7 +196,18 @@ func (t *TCPMessage) UpdateResponseAck() uint32 {
 	return t.ResponseAck
 }
 
+func (t *TCPMessage) ID() [10]byte {
+	return t.packets[0].ID
+}
+
 // ResponseID generate message ID for request response
-func (t *TCPMessage) ResponseID() string {
-	return t.packets[0].Addr + strconv.Itoa(int(t.packets[0].SrcPort)) + strconv.Itoa(int(t.ResponseAck))
+func (t *TCPMessage) ResponseID() [10]byte {
+	var id [10]byte
+	p := t.packets[0]
+
+	copy(id[:4], p.Addr)
+	copy(id[4:], p.Data[2:4]) // Dest port
+	binary.BigEndian.PutUint32(id[6:10], t.ResponseAck)
+
+	return id
 }

raw_socket_listener/tcp_message_test.go

@@ -3,13 +3,12 @@ package rawSocket
 import (
 	"bytes"
 	_ "log"
-	"strconv"
 	"testing"
 )
 
 func buildPacket(isIncoming bool, Ack, Seq uint32, Data []byte) (packet *TCPPacket) {
 	packet = &TCPPacket{
-		Addr: "",
+		Addr: []byte(""),
 		Ack:  Ack,
 		Seq:  Seq,
 		Data: Data,
@@ -26,14 +25,12 @@ func buildPacket(isIncoming bool, Ack, Seq uint32, Data []byte) (packet *TCPPack
 }
 
 func buildMessage(p *TCPPacket) *TCPMessage {
-	id := p.Addr + strconv.Itoa(int(p.DestPort)) + strconv.Itoa(int(p.Ack))
-
 	isIncoming := false
 	if p.SrcPort == 1 {
 		isIncoming = true
 	}
 
-	m := NewTCPMessage(id, p.Seq, p.Ack, isIncoming)
+	m := NewTCPMessage(p.Seq, p.Ack, isIncoming)
 	m.AddPacket(p)
 
 	return m

raw_socket_listener/tcp_packet.go

@@ -32,38 +32,34 @@ type TCPPacket struct {
 	Checksum   uint16
 	Urgent     uint16
 
+	Raw []byte
 	Data []byte
-
-	Addr string
+	Addr []byte
+	ID [10]byte
 }
 
 // ParseTCPPacket takes address and tcp payload and returns parsed TCPPacket
-func ParseTCPPacket(addr string, b []byte) (p *TCPPacket) {
-	p = &TCPPacket{Data: b}
+func ParseTCPPacket(addr []byte, data []byte) (p *TCPPacket) {
+	p = &TCPPacket{Raw: data}
 	p.ParseBasic()
 	p.Addr = addr
 
-	return p
-}
+	copy(p.ID[:4], addr)
+	copy(p.ID[4:], p.Raw[2:4]) // Dest port
+	copy(p.ID[6:], p.Raw[8:12]) // Ack
 
-// Parse TCP Packet, inspired by: https://github.com/miekg/pcap/blob/master/packet.go
-func (t *TCPPacket) Parse() {
-	t.ParseBasic()
-	t.Flags = binary.BigEndian.Uint16(t.Data[12:14]) & 0x1FF
-	t.Window = binary.BigEndian.Uint16(t.Data[14:16])
-	t.Checksum = binary.BigEndian.Uint16(t.Data[16:18])
-	t.Urgent = binary.BigEndian.Uint16(t.Data[18:20])
+	return
 }
 
 // ParseBasic set of fields
 func (t *TCPPacket) ParseBasic() {
-	t.DestPort = binary.BigEndian.Uint16(t.Data[2:4])
-	t.SrcPort = binary.BigEndian.Uint16(t.Data[0:2])
-	t.Seq = binary.BigEndian.Uint32(t.Data[4:8])
-	t.Ack = binary.BigEndian.Uint32(t.Data[8:12])
-	t.DataOffset = (t.Data[12] & 0xF0) >> 4
+	t.DestPort = binary.BigEndian.Uint16(t.Raw[2:4])
+	t.SrcPort = binary.BigEndian.Uint16(t.Raw[0:2])
+	t.Seq = binary.BigEndian.Uint32(t.Raw[4:8])
+	t.Ack = binary.BigEndian.Uint32(t.Raw[8:12])
+	t.DataOffset = (t.Raw[12] & 0xF0) >> 4
 
-	t.Data = t.Data[t.DataOffset*4:]
+	t.Data = t.Raw[t.DataOffset*4:]
 }
 
 func (t *TCPPacket) Dump() []byte {
@@ -89,7 +85,7 @@ func (t *TCPPacket) String() string {
 	}
 
 	return strings.Join([]string{
-		"Addr: " + t.Addr,
+		"Addr: " + string(t.Addr),
 		"Source port: " + strconv.Itoa(int(t.SrcPort)),
 		"Dest port:" + strconv.Itoa(int(t.DestPort)),
 		"Sequence:" + strconv.Itoa(int(t.Seq)),