diff --git a/raw_socket_listener/listener.go b/raw_socket_listener/listener.go index 50eeebe..34f55c2 100644 --- a/raw_socket_listener/listener.go +++ b/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 diff --git a/raw_socket_listener/tcp_message.go b/raw_socket_listener/tcp_message.go index df8caa9..c614cc2 100644 --- a/raw_socket_listener/tcp_message.go +++ b/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 } diff --git a/raw_socket_listener/tcp_message_test.go b/raw_socket_listener/tcp_message_test.go index f404541..da1c431 100644 --- a/raw_socket_listener/tcp_message_test.go +++ b/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 diff --git a/raw_socket_listener/tcp_packet.go b/raw_socket_listener/tcp_packet.go index b56ef42..f443a16 100644 --- a/raw_socket_listener/tcp_packet.go +++ b/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)),