1 // Copyright 2012 Google, Inc. All rights reserved.
3 // Use of this source code is governed by a BSD-style license
4 // that can be found in the LICENSE file in the root of the source
7 // This binary provides an example of connecting up bidirectional streams from
8 // the unidirectional streams provided by gopacket/tcpassembly.
14 "github.com/google/gopacket"
15 "github.com/google/gopacket/examples/util"
16 "github.com/google/gopacket/layers"
17 "github.com/google/gopacket/pcap"
18 "github.com/google/gopacket/tcpassembly"
23 var iface = flag.String("i", "eth0", "Interface to get packets from")
24 var snaplen = flag.Int("s", 16<<10, "SnapLen for pcap packet capture")
25 var filter = flag.String("f", "tcp", "BPF filter for pcap")
26 var logAllPackets = flag.Bool("v", false, "Logs every packet in great detail")
28 // key is used to map bidirectional streams to each other.
30 net, transport gopacket.Flow
33 // String prints out the key in a human-readable fashion.
34 func (k key) String() string {
35 return fmt.Sprintf("%v:%v", k.net, k.transport)
38 // timeout is the length of time to wait befor flushing connections and
39 // bidirectional stream pairs.
40 const timeout time.Duration = time.Minute * 5
42 // myStream implements tcpassembly.Stream
43 type myStream struct {
44 bytes int64 // total bytes seen on this stream.
45 bidi *bidi // maps to my bidirectional twin.
46 done bool // if true, we've seen the last packet we're going to for this stream.
49 // bidi stores each unidirectional side of a bidirectional stream.
51 // When a new stream comes in, if we don't have an opposite stream, a bidi is
52 // created with 'a' set to the new stream. If we DO have an opposite stream,
53 // 'b' is set to the new stream.
55 key key // Key of the first stream, mostly for logging.
56 a, b *myStream // the two bidirectional streams.
57 lastPacketSeen time.Time // last time we saw a packet from either stream.
60 // myFactory implements tcpassmebly.StreamFactory
61 type myFactory struct {
62 // bidiMap maps keys to bidirectional stream pairs.
66 // New handles creating a new tcpassembly.Stream.
67 func (f *myFactory) New(netFlow, tcpFlow gopacket.Flow) tcpassembly.Stream {
68 // Create a new stream.
71 // Find the bidi bidirectional struct for this stream, creating a new one if
72 // one doesn't already exist in the map.
73 k := key{netFlow, tcpFlow}
76 bd = &bidi{a: s, key: k}
77 log.Printf("[%v] created first side of bidirectional stream", bd.key)
78 // Register bidirectional with the reverse key, so the matching stream going
79 // the other direction will find it.
80 f.bidiMap[key{netFlow.Reverse(), tcpFlow.Reverse()}] = bd
82 log.Printf("[%v] found second side of bidirectional stream", bd.key)
84 // Clear out the bidi we're using from the map, just in case.
91 // emptyStream is used to finish bidi that only have one stream, in
93 var emptyStream = &myStream{done: true}
95 // collectOldStreams finds any streams that haven't received a packet within
96 // 'timeout', and sets/finishes the 'b' stream inside them. The 'a' stream may
97 // still receive packets after this.
98 func (f *myFactory) collectOldStreams() {
99 cutoff := time.Now().Add(-timeout)
100 for k, bd := range f.bidiMap {
101 if bd.lastPacketSeen.Before(cutoff) {
102 log.Printf("[%v] timing out old stream", bd.key)
103 bd.b = emptyStream // stub out b with an empty stream.
104 delete(f.bidiMap, k) // remove it from our map.
105 bd.maybeFinish() // if b was the last stream we were waiting for, finish up.
110 // Reassembled handles reassembled TCP stream data.
111 func (s *myStream) Reassembled(rs []tcpassembly.Reassembly) {
112 for _, r := range rs {
113 // For now, we'll simply count the bytes on each side of the TCP stream.
114 s.bytes += int64(len(r.Bytes))
116 s.bytes += int64(r.Skip)
118 // Mark that we've received new packet data.
119 // We could just use time.Now, but by using r.Seen we handle the case
120 // where packets are being read from a file and could be very old.
121 if s.bidi.lastPacketSeen.After(r.Seen) {
122 s.bidi.lastPacketSeen = r.Seen
127 // ReassemblyComplete marks this stream as finished.
128 func (s *myStream) ReassemblyComplete() {
133 // maybeFinish will wait until both directions are complete, then print out
135 func (bd *bidi) maybeFinish() {
138 log.Fatalf("[%v] a should always be non-nil, since it's set when bidis are created", bd.key)
140 log.Printf("[%v] still waiting on first stream", bd.key)
142 log.Printf("[%v] no second stream yet", bd.key)
144 log.Printf("[%v] still waiting on second stream", bd.key)
146 log.Printf("[%v] FINISHED, bytes: %d tx, %d rx", bd.key, bd.a.bytes, bd.b.bytes)
152 log.Printf("starting capture on interface %q", *iface)
153 // Set up pcap packet capture
154 handle, err := pcap.OpenLive(*iface, int32(*snaplen), true, pcap.BlockForever)
158 if err := handle.SetBPFFilter(*filter); err != nil {
163 streamFactory := &myFactory{bidiMap: make(map[key]*bidi)}
164 streamPool := tcpassembly.NewStreamPool(streamFactory)
165 assembler := tcpassembly.NewAssembler(streamPool)
167 log.Println("reading in packets")
168 // Read in packets, pass to assembler.
169 packetSource := gopacket.NewPacketSource(handle, handle.LinkType())
170 packets := packetSource.Packets()
171 ticker := time.Tick(timeout / 4)
174 case packet := <-packets:
178 if packet.NetworkLayer() == nil || packet.TransportLayer() == nil || packet.TransportLayer().LayerType() != layers.LayerTypeTCP {
179 log.Println("Unusable packet")
182 tcp := packet.TransportLayer().(*layers.TCP)
183 assembler.AssembleWithTimestamp(packet.NetworkLayer().NetworkFlow(), tcp, packet.Metadata().Timestamp)
186 // Every minute, flush connections that haven't seen activity in the past minute.
187 log.Println("---- FLUSHING ----")
188 assembler.FlushOlderThan(time.Now().Add(-timeout))
189 streamFactory.collectOldStreams()