feat(experimental): Add --embed-keylog-to-pcapng flag for embeding …

…TLS key logs into pcapng file

experimental: only support unstripped Golang binary and must combined with `-- CMD [ARGS]`

.circleci/config.yml

@@ -92,8 +92,19 @@ jobs:
       - run:
           name: e2e (test nat)
           command: |
+            sudo bash testdata/test_nat.sh ./ptcpdump || \
+            sudo bash testdata/test_nat.sh ./ptcpdump || \
             sudo bash testdata/test_nat.sh ./ptcpdump
 
+      - run:
+          name: e2e (test go tls keylog)
+          command: |
+            sudo apt install -y tshark
+            make -C testdata/gohttpapp build
+            sudo bash testdata/test_gotls_keylog.sh ./ptcpdump || \
+            sudo bash testdata/test_gotls_keylog.sh ./ptcpdump || \
+            sudo bash testdata/test_gotls_keylog.sh ./ptcpdump
+
   docker-e2e:
     machine:
       image: ubuntu-2204:2024.04.4

.github/workflows/test.yml

@@ -246,4 +246,24 @@ jobs:
             uname -a
             cat /etc/issue
 
+            bash /host/testdata/test_nat.sh /host/ptcpdump/ptcpdump || \
+            bash /host/testdata/test_nat.sh /host/ptcpdump/ptcpdump || \
             bash /host/testdata/test_nat.sh /host/ptcpdump/ptcpdump
+
+      - name: install tshark and build demo app
+        run: |
+          sudo apt install -y tshark
+          make -C testdata/gohttpapp build
+
+      - name: Test go tls keylog
+        uses: cilium/little-vm-helper@97c89f004bd0ab4caeacfe92ebc956e13e362e6b # v0.0.19
+        with:
+          provision: 'false'
+          cmd: |
+            set -ex
+            uname -a
+            cat /etc/issue
+
+            sudo bash testdata/test_gotls_keylog.sh /host/ptcpdump/ptcpdump || \
+            sudo bash testdata/test_gotls_keylog.sh /host/ptcpdump/ptcpdump || \
+            sudo bash testdata/test_gotls_keylog.sh /host/ptcpdump/ptcpdump

.gitignore

@@ -18,3 +18,5 @@
 /*.pcapng
 /*.pcap
 /dist/
+keylog.txt
+/testdata/gohttpapp/gohttpapp*

cmd/capture.go

@@ -18,14 +18,9 @@ import (
 )
 
 func capture(ctx context.Context, stop context.CancelFunc, opts Options) error {
-	gcr, err := getGoKeyLogEventConsumer(opts)
-	if err != nil {
-		logFatal(err)
-	}
-
 	btfSpec, btfPath, err := btf.LoadBTFSpec(opts.btfPath)
 	if err != nil {
-		logFatal(err)
+		return err
 	}
 	if btfPath != btf.DefaultPath {
 		log.Warnf("use BTF specs from %s", btfPath)
@@ -38,18 +33,6 @@ func capture(ctx context.Context, stop context.CancelFunc, opts Options) error {
 		pcache.WithContainerCache(cc)
 	}
 
-	var subProcessFinished <-chan struct{}
-	var subProcessLoaderPid int
-	if len(opts.subProgArgs) > 0 {
-		log.Info("start sub process loader")
-		subProcessLoaderPid, subProcessFinished, err = utils.StartSubProcessLoader(ctx, os.Args[0], opts.subProgArgs)
-		if err != nil {
-			return err
-		}
-		opts.pids = []uint{uint(subProcessLoaderPid)}
-		opts.followForks = true
-	}
-
 	writers, fcloser, err := getWriters(opts, pcache)
 	if err != nil {
 		return err
@@ -62,6 +45,23 @@ func capture(ctx context.Context, stop context.CancelFunc, opts Options) error {
 			fcloser()
 		}
 	}()
+	gcr, err := getGoKeyLogEventConsumer(&opts, writers)
+	if err != nil {
+		return err
+	}
+
+	var subProcessFinished <-chan struct{}
+	var subProcessLoaderPid int
+	if len(opts.subProgArgs) > 0 {
+		log.Info("start sub process loader")
+		subProcessLoaderPid, subProcessFinished, err = utils.StartSubProcessLoader(ctx, os.Args[0], opts.subProgArgs)
+		if err != nil {
+			return err
+		}
+		opts.pids = []uint{uint(subProcessLoaderPid)}
+		opts.followForks = true
+	}
+
 	pcache.Start(ctx)
 
 	log.Info("start get current connections")
@@ -104,10 +104,7 @@ func capture(ctx context.Context, stop context.CancelFunc, opts Options) error {
 	go gcr.Start(ctx, goTlsKeyLogEventsCh)
 
 	var stopByInternal bool
-	packetConsumer := consumer.NewPacketEventConsumer(writers)
-	if opts.delayBeforeHandlePacketEvents > 0 {
-		time.Sleep(opts.delayBeforeHandlePacketEvents)
-	}
+	packetConsumer := consumer.NewPacketEventConsumer(writers).WithDelay(opts.delayBeforeHandlePacketEvents)
 	if subProcessLoaderPid > 0 {
 		go func() {
 			log.Infof("notify loader %d to start sub process", subProcessLoaderPid)
@@ -116,6 +113,7 @@ func capture(ctx context.Context, stop context.CancelFunc, opts Options) error {
 			log.Info("sub process exited")
 			time.Sleep(time.Second * 3)
 			stopByInternal = true
+			time.Sleep(opts.delayBeforeHandlePacketEvents)
 			stop()
 		}()
 	}

cmd/gotls.go

@@ -3,6 +3,7 @@ package cmd
 import (
 	"fmt"
 	"os/exec"
+	"time"
 
 	"debug/buildinfo"
 	"github.com/cilium/ebpf/link"
@@ -14,18 +15,29 @@ import (
 
 const goTLSSymbolWriteKeyLog = "crypto/tls.(*Config).writeKeyLog"
 
-func getGoKeyLogEventConsumer(opts Options) (*consumer.GoKeyLogEventConsumer, error) {
-	var writers []writer.KeyLogWriter
+func getGoKeyLogEventConsumer(opts *Options, packetWriters []writer.PacketWriter) (*consumer.GoKeyLogEventConsumer, error) {
+	var keylogWs []writer.KeyLogWriter
 
+	if opts.embedTLSKeyLogToPcapng {
+		for _, pw := range packetWriters {
+			if pngw, ok := pw.(*writer.PcapNGWriter); ok {
+				keylogWs = append(keylogWs, writer.NewKeyLogPcapNGWriter(pngw))
+				if opts.delayBeforeHandlePacketEvents == 0 {
+					opts.delayBeforeHandlePacketEvents = time.Second * 3
+				}
+				break
+			}
+		}
+	}
 	if opts.writeTLSKeyLogPath != "" {
 		w, err := writer.NewKeyLogFileWriter(opts.writeTLSKeyLogPath)
 		if err != nil {
 			return nil, err
 		}
-		writers = append(writers, w)
+		keylogWs = append(keylogWs, w)
 	}
 
-	cr := consumer.NewGoKeyLogEventConsumer(2, writers...)
+	cr := consumer.NewGoKeyLogEventConsumer(10, keylogWs...)
 	return cr, nil
 }
 

cmd/options.go

@@ -60,7 +60,8 @@ type Options struct {
 	criRuntimeEndpoint string
 	btfPath            string
 
-	writeTLSKeyLogPath string
+	writeTLSKeyLogPath     string
+	embedTLSKeyLogToPcapng bool
 
 	subProgArgs []string
 

cmd/root.go

@@ -120,6 +120,8 @@ func init() {
 
 	rootCmd.Flags().StringVar(&opts.writeTLSKeyLogPath, "write-keylog-file", "",
 		"Write TLS Key Log file to this path (experimental: only support unstripped Go binary and must combined with `-- CMD [ARGS]`)")
+	rootCmd.Flags().BoolVar(&opts.embedTLSKeyLogToPcapng, "embed-keylog-to-pcapng", false,
+		"Write TLS Key Log file to this path (experimental: only support unstripped Go binary and must combined with `-- CMD [ARGS]`)")
 
 	silenceKlog()
 }

internal/consumer/net.go

@@ -7,12 +7,15 @@ import (
 	"github.com/mozillazg/ptcpdump/internal/event"
 	"github.com/mozillazg/ptcpdump/internal/log"
 	"github.com/mozillazg/ptcpdump/internal/writer"
+	"time"
 )
 
 type PacketEventConsumer struct {
 	writers        []writer.PacketWriter
 	devices        map[int]dev.Device
 	processedCount int
+
+	delay time.Duration
 }
 
 func NewPacketEventConsumer(writers []writer.PacketWriter) *PacketEventConsumer {
@@ -23,7 +26,15 @@ func NewPacketEventConsumer(writers []writer.PacketWriter) *PacketEventConsumer
 	}
 }
 
+func (c *PacketEventConsumer) WithDelay(delay time.Duration) *PacketEventConsumer {
+	c.delay = delay
+	return c
+}
+
 func (c *PacketEventConsumer) Start(ctx context.Context, ch <-chan bpf.BpfPacketEventWithPayloadT, maxPacketCount uint) {
+	if c.delay > 0 {
+		time.Sleep(c.delay)
+	}
 	for {
 		select {
 		case <-ctx.Done():

internal/writer/keylog.go

@@ -16,6 +16,10 @@ type KeyLogFileWriter struct {
 	f     *os.File
 }
 
+type KeyLogPcapNGWriter struct {
+	w *PcapNGWriter
+}
+
 func NewKeyLogFileWriter(fpath string) (*KeyLogFileWriter, error) {
 	f, err := os.OpenFile(fpath, os.O_CREATE|os.O_WRONLY|os.O_APPEND, 0644)
 	if err != nil {
@@ -27,6 +31,10 @@ func NewKeyLogFileWriter(fpath string) (*KeyLogFileWriter, error) {
 	}, nil
 }
 
+func NewKeyLogPcapNGWriter(w *PcapNGWriter) *KeyLogPcapNGWriter {
+	return &KeyLogPcapNGWriter{w: w}
+}
+
 func (k *KeyLogFileWriter) Write(line string) error {
 	_, err := k.f.WriteString(line)
 	return err
@@ -39,3 +47,16 @@ func (k *KeyLogFileWriter) Flush() error {
 func (k *KeyLogFileWriter) Close() error {
 	return k.f.Close()
 }
+
+func (k *KeyLogPcapNGWriter) Write(line string) error {
+	err := k.w.WriteTLSKeyLog(line)
+	return err
+}
+
+func (k *KeyLogPcapNGWriter) Flush() error {
+	return nil
+}
+
+func (k *KeyLogPcapNGWriter) Close() error {
+	return nil
+}

internal/writer/pcapng.go

@@ -1,7 +1,9 @@
 package writer
 
 import (
+	"bytes"
 	"fmt"
+	"sync"
 
 	"github.com/gopacket/gopacket"
 	"github.com/gopacket/gopacket/pcapgo"
@@ -15,10 +17,12 @@ type PcapNGWriter struct {
 	pcache *metadata.ProcessCache
 
 	noBuffer bool
+	lock     sync.Mutex
+	keylogs  bytes.Buffer
 }
 
 func NewPcapNGWriter(pw *pcapgo.NgWriter, pcache *metadata.ProcessCache) *PcapNGWriter {
-	return &PcapNGWriter{pw: pw, pcache: pcache}
+	return &PcapNGWriter{pw: pw, pcache: pcache, lock: sync.Mutex{}}
 }
 
 func (w *PcapNGWriter) Write(e *event.Packet) error {
@@ -56,6 +60,10 @@ func (w *PcapNGWriter) Write(e *event.Packet) error {
 		)
 	}
 
+	if err := w.writeTLSKeyLog(); err != nil {
+		return err
+	}
+
 	if err := w.pw.WritePacketWithOptions(info, e.Data, opts); err != nil {
 		return fmt.Errorf("writing packet: %w", err)
 	}
@@ -66,6 +74,33 @@ func (w *PcapNGWriter) Write(e *event.Packet) error {
 	return nil
 }
 
+func (w *PcapNGWriter) WriteTLSKeyLog(line string) error {
+	w.lock.Lock()
+	defer w.lock.Unlock()
+
+	w.keylogs.WriteString(line)
+
+	return nil
+}
+
+func (w *PcapNGWriter) writeTLSKeyLog() error {
+	w.lock.Lock()
+	defer w.lock.Unlock()
+
+	lines := w.keylogs.Bytes()
+	if len(lines) == 0 {
+		return nil
+	}
+
+	if err := w.pw.WriteDecryptionSecretsBlock(pcapgo.DSB_SECRETS_TYPE_TLS, lines); err != nil {
+		return fmt.Errorf("writing tls key log: %w", err)
+	}
+
+	w.keylogs.Reset()
+
+	return nil
+}
+
 func (w *PcapNGWriter) Flush() error {
 	return w.pw.Flush()
 }

testdata/gohttpapp/Makefile

@@ -0,0 +1,6 @@
+
+
+
+.PHONY: build
+build:
+	go build main.go

testdata/gohttpapp/main.go

@@ -0,0 +1,30 @@
+package main
+
+import (
+	"fmt"
+	"io"
+	"net/http"
+	"net/http/httptest"
+)
+
+func newHttpServer() *httptest.Server {
+	s := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.Write([]byte("response from TLS Server"))
+	}))
+	return s
+}
+
+func main() {
+	s := newHttpServer()
+	defer s.Close()
+
+	client := s.Client()
+	resp, err := client.Get(s.URL + "/foo/bar")
+	if err != nil {
+		panic(err)
+	}
+	defer resp.Body.Close()
+
+	body, _ := io.ReadAll(resp.Body)
+	fmt.Println(string(body))
+}

testdata/test_gotls_keylog.sh

@@ -0,0 +1,29 @@
+#!/usr/bin/env bash
+
+set -xe
+
+SCRIPT_DIR="$( cd "$( dirname "${BASH_SOURCE[0]}" )" > /dev/null && pwd )"
+CMD="$1"
+APP="${SCRIPT_DIR}/gohttpapp/gohttpapp"
+FILE_PREFIX="/tmp/ptcpdump"
+KEYLOG_PATH="${FILE_PREFIX}_keylog.txt"
+PCAP_FILE="${FILE_PREFIX}_keylog_01.pcap"
+PCAPNG_FILE="${FILE_PREFIX}_keylog_01.pcapng"
+
+function test_keylog_to_file() {
+    ${CMD} -i any --write-keylog-file ${KEYLOG_PATH} -w ${PCAP_FILE} -- ${APP}
+    cat ${KEYLOG_PATH}
+    tshark -r ${PCAP_FILE} -o tls.keylog_file:${KEYLOG_PATH} | grep "GET /foo/bar HTTP/1.1"
+}
+
+function test_keylog_to_pcapng() {
+    ${CMD} -i any --embed-keylog-to-pcapng -w ${PCAPNG_FILE} -- ${APP}
+    tshark -r ${PCAPNG_FILE} | grep "GET /foo/bar HTTP/1.1"
+}
+
+function main() {
+  test_keylog_to_file
+  test_keylog_to_pcapng
+}
+
+main