diff --git a/listener/listener_test.go b/listener/listener_test.go index 775399c..34a8772 100644 --- a/listener/listener_test.go +++ b/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) + } +} +*/ diff --git a/listener/raw_tcp_listener.go b/listener/raw_tcp_listener.go index 911c83c..e9f2c14 100644 --- a/listener/raw_tcp_listener.go +++ b/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 diff --git a/listener/raw_tcp_listener_test.go b/listener/raw_tcp_listener_test.go new file mode 100644 index 0000000..ca82bbd --- /dev/null +++ b/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() +} diff --git a/listener/tcp_message.go b/listener/tcp_message.go index 0172149..dca2224 100644 --- a/listener/tcp_message.go +++ b/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 { diff --git a/listener/tcp_packet.go b/listener/tcp_packet.go index 0020279..2282782 100644 --- a/listener/tcp_packet.go +++ b/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