user: improve pcapng writer, flush every 2s.

Signed-off-by: cfc4n <[email protected]>

user/module/probe_gotls.go

@@ -37,7 +37,6 @@ var (
 
 // GoTLSProbe represents a probe for Go SSL
 type GoTLSProbe struct {
-	Module
 	MTCProbe
 	bpfManager        *manager.Manager
 	bpfManagerOptions manager.Options
@@ -213,20 +212,6 @@ func (g *GoTLSProbe) DecodeFun(m *ebpf.Map) (event.IEventStruct, bool) {
 }
 
 func (g *GoTLSProbe) Close() error {
-
-	if g.eBPFProgramType == TlsCaptureModelTypePcap {
-		g.logger.Printf("%s\tsaving pcapng file %s\n", g.Name(), g.pcapngFilename)
-		i, err := g.savePcapng()
-		if err != nil {
-			g.logger.Printf("%s\tsave pcanNP failed, error:%v. \n", g.Name(), err)
-		}
-		if i == 0 {
-			g.logger.Printf("nothing captured, please check your network interface, see \"ecapture tls -h\" for more information.")
-		} else {
-			g.logger.Printf("%s\t save %d packets into pcapng file.\n", g.Name(), i)
-		}
-	}
-
 	g.logger.Printf("%s\tclose. \n", g.Name())
 	if err := g.bpfManager.Stop(manager.CleanAll); err != nil {
 		return fmt.Errorf("couldn't stop manager %v .", err)

user/module/probe_gotls_pcap.go

@@ -59,6 +59,11 @@ func (g *GoTLSProbe) setupManagersPcap() error {
 		return err
 	}
 
+	// Serve pcapng writer to flush pcapng file
+	go func() {
+		g.ServePcap()
+	}()
+
 	var (
 		sec string
 		fn  string

user/module/probe_openssl.go

@@ -56,7 +56,6 @@ const (
 )
 
 type MOpenSSLProbe struct {
-	Module
 	MTCProbe
 	bpfManager        *manager.Manager
 	bpfManagerOptions manager.Options
@@ -108,6 +107,8 @@ func (m *MOpenSSLProbe) Init(ctx context.Context, logger *log.Logger, conf confi
 		if err != nil {
 			return err
 		}
+		m.tcPacketsChan = make(chan *TcPacket, 2048)
+		m.tcPackets = make([]*TcPacket, 0, 256)
 		m.pcapngFilename = fileInfo
 	case config.TlsCaptureModelText:
 		fallthrough
@@ -129,7 +130,6 @@ func (m *MOpenSSLProbe) Init(ctx context.Context, logger *log.Logger, conf confi
 	m.startTime = uint64(startTime)
 	m.bootTime = uint64(bootTime)
 
-	m.tcPackets = make([]*TcPacket, 0, 1024)
 	m.tcPacketLocker = &sync.Mutex{}
 	m.masterKeyBuffer = bytes.NewBuffer([]byte{})
 
@@ -231,18 +231,6 @@ func (m *MOpenSSLProbe) start() error {
 }
 
 func (m *MOpenSSLProbe) Close() error {
-	if m.eBPFProgramType == TlsCaptureModelTypePcap {
-		m.logger.Printf("%s\tsaving pcapng file %s\n", m.Name(), m.pcapngFilename)
-		i, err := m.savePcapng()
-		if err != nil {
-			m.logger.Printf("%s\tsave pcanNP failed, error:%v. \n", m.Name(), err)
-		}
-		if i == 0 {
-			m.logger.Printf("nothing captured, please check your network interface, see \"ecapture tls -h\" for more information.")
-		} else {
-			m.logger.Printf("%s\t save %d packets into pcapng file.\n", m.Name(), i)
-		}
-	}
 
 	m.logger.Printf("%s\tclose. \n", m.Name())
 	if err := m.bpfManager.Stop(manager.CleanAll); err != nil {

user/module/probe_openssl_pcap.go

@@ -86,6 +86,11 @@ func (m *MOpenSSLProbe) setupManagersPcap() error {
 		return err
 	}
 
+	// Serve pcapng writer to flush pcapng file
+	go func() {
+		m.ServePcap()
+	}()
+
 	m.bpfManager = &manager.Manager{
 		Probes: []*manager.Probe{
 			// customize deleteed TC filter

user/module/probe_tc.go → user/module/probe_pcap.go

@@ -5,6 +5,7 @@ import (
 	"ecapture/pkg/util/ethernet"
 	"ecapture/user/event"
 	"encoding/binary"
+	"errors"
 	"fmt"
 	"github.com/google/gopacket"
 	"github.com/google/gopacket/layers"
@@ -17,6 +18,8 @@ import (
 	"time"
 )
 
+var eOverflow = errors.New("pcapNG channel overflow")
+
 // packets of TC probe
 type TcPacket struct {
 	info gopacket.CaptureInfo
@@ -68,8 +71,7 @@ func (NetCaptureData) GetSizeBytes() uint32 {
 }
 
 type MTCProbe struct {
-	//logger          *log.Logger
-	//mName           string
+	Module
 	pcapngFilename  string
 	ifIdex          int
 	ifName          string
@@ -79,6 +81,7 @@ type MTCProbe struct {
 	tcPackets       []*TcPacket
 	masterKeyBuffer *bytes.Buffer
 	tcPacketLocker  *sync.Mutex
+	tcPacketsChan   chan *TcPacket
 }
 
 func (t *MTCProbe) dumpTcSkb(tcEvent *event.TcSkbEvent) error {
@@ -146,7 +149,10 @@ func (t *MTCProbe) savePcapng() (i int, err error) {
 		return
 	}
 	t.tcPacketLocker.Lock()
-	defer t.tcPacketLocker.Unlock()
+	defer func() {
+		t.tcPackets = t.tcPackets[:0]
+		t.tcPacketLocker.Unlock()
+	}()
 	for _, packet := range t.tcPackets {
 		err = t.pcapWriter.WritePacket(packet.info, packet.data)
 		i++
@@ -221,7 +227,7 @@ func (t *MTCProbe) createPcapng(netIfs []net.Interface) error {
 func (t *MTCProbe) writePacket(dataLen uint32, timeStamp time.Time, packetBytes []byte) error {
 
 	// TODO add packetMeta info (e.g: process. pid, commom, etc.)
-	
+
 	info := gopacket.CaptureInfo{
 		Timestamp:     timeStamp,
 		CaptureLength: int(dataLen),
@@ -235,11 +241,63 @@ func (t *MTCProbe) writePacket(dataLen uint32, timeStamp time.Time, packetBytes
 
 	packet := &TcPacket{info: info, data: packetBytes}
 
-	t.tcPackets = append(t.tcPackets, packet)
-	return nil
+	select {
+	case t.tcPacketsChan <- packet:
+		return nil
+	default:
+		return eOverflow
+	}
 }
 
 func (t *MTCProbe) savePcapngSslKeyLog(sslKeyLog []byte) (err error) {
 	_, e := t.masterKeyBuffer.Write(sslKeyLog)
 	return e
 }
+
+// ServePcap is used to serve pcapng file
+func (t *MTCProbe) ServePcap() {
+	var ti = time.NewTicker(2 * time.Second)
+	t.logger.Printf("%s\tsaving pcapng file %s\n", t.Name(), t.pcapngFilename)
+	var allCount int
+	defer func() {
+		if allCount == 0 {
+			t.logger.Printf("nothing captured, please check your network interface, see \"ecapture tls -h\" for more information.")
+		} else {
+			t.logger.Printf("%s\t save %d packets into pcapng file.\n", t.Name(), allCount)
+		}
+		ti.Stop()
+	}()
+
+	for {
+		select {
+		case <-ti.C:
+			// append tcPackets to tcPackets Array from tcPacketsChan
+			var i int
+			for packet := range t.tcPacketsChan {
+				t.tcPackets = append(t.tcPackets, packet)
+				i++
+			}
+			if i == 0 {
+				continue
+			}
+
+			i, e := t.savePcapng()
+			if e != nil {
+				t.logger.Printf("save pcapng err:%s\n", e.Error())
+			} else {
+				t.logger.Printf("save pcapng success, count:%d\n", i)
+				allCount += i
+			}
+		case <-t.ctx.Done():
+			// TODO: save all data to file
+			i, e := t.savePcapng()
+			if e != nil {
+				t.logger.Printf("save pcapng err:%s\n", e.Error())
+			} else {
+				t.logger.Printf("save pcapng success, count:%d\n", i)
+				allCount += i
+			}
+			return
+		}
+	}
+}