Use zerolog + intercept SIGINT (#2)

* Use zerolog + intercept SIGINT
* Use RFC3339 for time format and add -json flag to disable pretty printing
* Added new flags to README and fixed issue when command exit was not successful

README.md

@@ -36,6 +36,10 @@ $ ./pcap-broker --help
 Usage of ./pcap-broker:
   -cmd string
         command to execute for pcap data (eg: tcpdump -i eth0 -n --immediate-mode -s 65535 -U -w -)
+  -debug
+        enable debug logging
+  -json
+        enable json logging
   -listen string
         listen address for pcap-over-ip (eg: localhost:4242)
   -n    disable reverse lookup of connecting PCAP-over-IP client IP address

cmd/pcap-broker/main.go

@@ -2,12 +2,13 @@ package main
 
 import (
 	"context"
+	"errors"
 	"flag"
 	"fmt"
-	"log"
 	"net"
 	"os"
 	"os/exec"
+	"os/signal"
 	"time"
 
 	"github.com/google/shlex"
@@ -15,6 +16,9 @@ import (
 	"github.com/google/gopacket"
 	"github.com/google/gopacket/pcap"
 	"github.com/google/gopacket/pcapgo"
+
+	"github.com/rs/zerolog"
+	"github.com/rs/zerolog/log"
 )
 
 type PcapClient struct {
@@ -50,17 +54,34 @@ func lookupHostnameWithTimeout(addr net.Addr, timeout time.Duration) (string, st
 	return names[0], port, nil
 }
 
-func main() {
+var (
+	pcapCommand     = flag.String("cmd", "", "command to execute for pcap data (eg: tcpdump -i eth0 -n --immediate-mode -s 65535 -U -w -)")
+	listenAddress   = flag.String("listen", "", "listen address for pcap-over-ip (eg: localhost:4242)")
+	noReverseLookup = flag.Bool("n", false, "disable reverse lookup of connecting PCAP-over-IP client IP address")
+	debug           = flag.Bool("debug", false, "enable debug logging")
+	json            = flag.Bool("json", false, "enable json logging")
+)
 
-	pcapCommand := flag.String("cmd", "", "command to execute for pcap data (eg: tcpdump -i eth0 -n --immediate-mode -s 65535 -U -w -)")
-	listenAddress := flag.String("listen", "", "listen address for pcap-over-ip (eg: localhost:4242)")
-	noReverseLookup := flag.Bool("n", false, "disable reverse lookup of connecting PCAP-over-IP client IP address")
+func main() {
 	flag.Parse()
 
+	if !*json {
+		log.Logger = log.Output(zerolog.ConsoleWriter{
+			Out:        os.Stderr,
+			TimeFormat: time.RFC3339,
+		})
+	}
+
+	if *debug {
+		zerolog.SetGlobalLevel(zerolog.DebugLevel)
+	} else {
+		zerolog.SetGlobalLevel(zerolog.InfoLevel)
+	}
+
 	if *pcapCommand == "" {
 		*pcapCommand = os.Getenv("PCAP_COMMAND")
 		if *pcapCommand == "" {
-			log.Fatalf("Error: PCAP_COMMAND or -cmd not set, see --help for usage")
+			log.Fatal().Msg("PCAP_COMMAND or -cmd not set, see --help for usage")
 		}
 	}
 
@@ -71,68 +92,83 @@ func main() {
 		}
 	}
 
-	log.Printf("config PCAP_COMMAND = %q", *pcapCommand)
-	log.Printf("config LISTEN_ADDRESS = %q", *listenAddress)
+	log.Debug().Str("pcapCommand", *pcapCommand).Send()
+	log.Debug().Str("listenAddress", *listenAddress).Send()
+
+	ctx, cancelFunc := signal.NotifyContext(context.Background(), os.Interrupt)
 
 	// Create connections to PcapClient map
-	var connMap = map[net.Conn]PcapClient{}
+	connMap := map[net.Conn]PcapClient{}
 
 	// Create a pipe for the command to write to, will be read by pcap.OpenOfflineFile
 	rStdout, wStdout, err := os.Pipe()
 	if err != nil {
-		log.Fatal(err)
+		log.Fatal().Err(err).Msg("failed to create pipe")
 	}
 
-	// Important or these will eventually be garbage collected and the pipe will close
-	defer rStdout.Close()
-	defer wStdout.Close()
-
 	// Acquire pcap data
 	args, err := shlex.Split(*pcapCommand)
 	if err != nil {
-		log.Fatal(err)
+		log.Fatal().Err(err).Msg("failed to parse PCAP_COMMAND")
 	}
-	cmd := exec.Command(args[0], args[1:]...)
-	log.Printf("cmd = %v", cmd.Args)
+	cmd := exec.CommandContext(ctx, args[0], args[1:]...)
+	log.Debug().Strs("args", args).Send()
+
 	cmd.Stdout = wStdout
-	cmd.Stderr = os.Stderr
+	cmd.Stderr = log.Logger
 
 	err = cmd.Start()
 	if err != nil {
-		log.Fatal(err)
+		log.Fatal().Err(err).Msg("failed to start command")
 	}
 
-	log.Printf("PID %v", cmd.Process.Pid)
+	log.Debug().Int("pid", cmd.Process.Pid).Msg("started process")
+
+	// close context on process exit
 	go func() {
 		err := cmd.Wait()
 		if err != nil {
-			log.Fatal("Process exited with error: ", err)
+			log.Fatal().Err(err).Msg("command exited with error")
 		}
-		log.Printf("process exited")
-		os.Exit(0)
+		cancelFunc()
 	}()
 
 	// Read from process stdout pipe
 	handle, err := pcap.OpenOfflineFile(rStdout)
 	if err != nil {
-		log.Fatal(err)
+		log.Fatal().Err(err).Msg("failed to open pcap file")
 	}
 
 	packetSource := gopacket.NewPacketSource(handle, handle.LinkType())
 	packetSource.Lazy = true
 	packetSource.NoCopy = true
-	go processPackets(packetSource, connMap)
 
-	log.Printf("PCAP-over-IP server listening on %v", *listenAddress)
-	l, err := net.Listen("tcp", *listenAddress)
+	go processPackets(ctx, packetSource, connMap)
+
+	log.Info().Msgf("PCAP-over-IP server listening on %v. press CTRL-C to exit", *listenAddress)
+
+	config := net.ListenConfig{}
+	l, err := config.Listen(ctx, "tcp", *listenAddress)
 	if err != nil {
-		log.Fatal(err)
+		log.Fatal().Err(err).Msg("failed to listen")
 	}
 
+	// close listener on context cancel
+	go func() {
+		<-ctx.Done()
+		cancelFunc()
+		err := l.Close()
+		if err != nil {
+			log.Err(err).Msg("failed to close listener")
+		}
+	}()
+
 	for {
 		conn, err := l.Accept()
-		if err != nil {
-			log.Fatal(err)
+		if err != nil && ctx.Err() == nil {
+			log.Fatal().Err(err).Msg("failed to accept connection")
+		} else if errors.Is(ctx.Err(), context.Canceled) {
+			break
 		}
 
 		if *noReverseLookup {
@@ -149,22 +185,56 @@ func main() {
 		writer := pcapgo.NewWriter(conn)
 
 		// Write pcap header
-		writer.WriteFileHeader(65535, handle.LinkType())
+		err = writer.WriteFileHeader(65535, handle.LinkType())
+		if err != nil {
+			log.Err(err).Msg("failed to write pcap header")
+			err := conn.Close()
+			if err != nil {
+				log.Err(err).Msg("failed to close connection")
+			}
+
+			continue
+		}
 
 		// add connection to map
 		connMap[conn] = PcapClient{writer: writer}
 	}
+
+	log.Info().Msg("PCAP-over-IP server exiting")
+
+	err = rStdout.Close()
+	if err != nil {
+		log.Err(err).Msg("failed to close read pipe")
 	}
 
-func processPackets(packetSource *gopacket.PacketSource, connMap map[net.Conn]PcapClient) {
+	err = wStdout.Close()
+	if err != nil {
+		log.Err(err).Msg("failed to close write pipe")
+	}
+}
+
+func processPackets(
+	ctx context.Context,
+	packetSource *gopacket.PacketSource,
+	connMap map[net.Conn]PcapClient,
+) {
 	for packet := range packetSource.Packets() {
+		select {
+		case <-ctx.Done():
+			return
+		default:
+		}
+
 		for conn, stats := range connMap {
 			ci := packet.Metadata().CaptureInfo
 			err := stats.writer.WritePacket(ci, packet.Data())
 			if err != nil {
-				log.Println(err)
+				log.Err(err).Msg("failed to write packet to connection")
 				delete(connMap, conn)
-				conn.Close()
+				err := conn.Close()
+				if err != nil {
+					log.Err(err).Msg("failed to close connection")
+				}
 				continue
 			}
 			stats.totalPackets += 1

go.mod

@@ -8,6 +8,9 @@ require (
 )
 
 require (
+	github.com/mattn/go-colorable v0.1.13 // indirect
+	github.com/mattn/go-isatty v0.0.20 // indirect
+	github.com/rs/zerolog v1.31.0 // indirect
 	golang.org/x/net v0.18.0 // indirect
 	golang.org/x/sys v0.14.0 // indirect
 )

go.sum

@@ -1,7 +1,19 @@
+github.com/coreos/go-systemd/v22 v22.5.0/go.mod h1:Y58oyj3AT4RCenI/lSvhwexgC+NSVTIJ3seZv2GcEnc=
+github.com/godbus/dbus/v5 v5.0.4/go.mod h1:xhWf0FNVPg57R7Z0UbKHbJfkEywrmjJnf7w5xrFpKfA=
 github.com/google/gopacket v1.1.19 h1:ves8RnFZPGiFnTS0uPQStjwru6uO6h+nlr9j6fL7kF8=
 github.com/google/gopacket v1.1.19/go.mod h1:iJ8V8n6KS+z2U1A8pUwu8bW5SyEMkXJB8Yo/Vo+TKTo=
 github.com/google/shlex v0.0.0-20191202100458-e7afc7fbc510 h1:El6M4kTTCOh6aBiKaUGG7oYTSPP8MxqL4YI3kZKwcP4=
 github.com/google/shlex v0.0.0-20191202100458-e7afc7fbc510/go.mod h1:pupxD2MaaD3pAXIBCelhxNneeOaAeabZDe5s4K6zSpQ=
+github.com/mattn/go-colorable v0.1.13 h1:fFA4WZxdEF4tXPZVKMLwD8oUnCTTo08duU7wxecdEvA=
+github.com/mattn/go-colorable v0.1.13/go.mod h1:7S9/ev0klgBDR4GtXTXX8a3vIGJpMovkB8vQcUbaXHg=
+github.com/mattn/go-isatty v0.0.16/go.mod h1:kYGgaQfpe5nmfYZH+SKPsOc2e4SrIfOl2e/yFXSvRLM=
+github.com/mattn/go-isatty v0.0.19/go.mod h1:W+V8PltTTMOvKvAeJH7IuucS94S2C6jfK/D7dTCTo3Y=
+github.com/mattn/go-isatty v0.0.20 h1:xfD0iDuEKnDkl03q4limB+vH+GxLEtL/jb4xVJSWWEY=
+github.com/mattn/go-isatty v0.0.20/go.mod h1:W+V8PltTTMOvKvAeJH7IuucS94S2C6jfK/D7dTCTo3Y=
+github.com/pkg/errors v0.9.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0=
+github.com/rs/xid v1.5.0/go.mod h1:trrq9SKmegXys3aeAKXMUTdJsYXVwGY3RLcfgqegfbg=
+github.com/rs/zerolog v1.31.0 h1:FcTR3NnLWW+NnTwwhFWiJSZr4ECLpqCm6QsEnyvbV4A=
+github.com/rs/zerolog v1.31.0/go.mod h1:/7mN4D5sKwJLZQ2b/znpjC3/GQWY/xaDXUM0kKWRHss=
 golang.org/x/crypto v0.0.0-20190308221718-c2843e01d9a2/go.mod h1:djNgcEr1/C05ACkg1iLfiJU5Ep61QUkGW8qpdssI0+w=
 golang.org/x/crypto v0.0.0-20191011191535-87dc89f01550/go.mod h1:yigFU9vqHzYiE8UmvKecakEJjdnWj3jj499lnFckfCI=
 golang.org/x/lint v0.0.0-20200302205851-738671d3881b/go.mod h1:3xt1FjdF8hUf6vQPIChWIBhFzV8gjjsPE/fR3IyQdNY=
@@ -13,6 +25,9 @@ golang.org/x/net v0.18.0/go.mod h1:/czyP5RqHAH4odGYxBJ1qz0+CE5WZ+2j1YgoEo8F2jQ=
 golang.org/x/sync v0.0.0-20190423024810-112230192c58/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
 golang.org/x/sys v0.0.0-20190215142949-d0b11bdaac8a/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
 golang.org/x/sys v0.0.0-20190412213103-97732733099d/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
+golang.org/x/sys v0.0.0-20220811171246-fbc7d0a398ab/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
+golang.org/x/sys v0.6.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
+golang.org/x/sys v0.12.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 golang.org/x/sys v0.14.0 h1:Vz7Qs629MkJkGyHxUlRHizWJRG2j8fbQKjELVSNhy7Q=
 golang.org/x/sys v0.14.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA=
 golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ=