From e846b3833cba858736966a05836926e42dec1747 Mon Sep 17 00:00:00 2001
From: Yun Zheng Hu <hu@fox-it.com>
Date: Thu, 3 Oct 2024 09:33:53 +0000
Subject: [PATCH] Refactor reverse lookups and connection logging

---
 cmd/pcap-broker/main.go | 41 ++++++++---------------------------------
 1 file changed, 8 insertions(+), 33 deletions(-)

diff --git a/cmd/pcap-broker/main.go b/cmd/pcap-broker/main.go
index d9f55c9..8872044 100644
--- a/cmd/pcap-broker/main.go
+++ b/cmd/pcap-broker/main.go
@@ -4,7 +4,6 @@ import (
 	"context"
 	"errors"
 	"flag"
-	"fmt"
 	"net"
 	"os"
 	"os/exec"
@@ -27,33 +26,6 @@ type PcapClient struct {
 	totalBytes   uint64
 }
 
-func lookupHostnameWithTimeout(addr net.Addr, timeout time.Duration) (string, string, error) {
-	// Extract the IP address and port from the Addr object
-	tcpAddr, ok := addr.(*net.TCPAddr)
-	if !ok {
-		return "", "", fmt.Errorf("unsupported address type: %T", addr)
-	}
-	ip := tcpAddr.IP.String()
-	port := fmt.Sprintf("%d", tcpAddr.Port)
-
-	// Create a new context with the given timeout
-	ctx, cancel := context.WithTimeout(context.Background(), timeout)
-	defer cancel()
-
-	// Create a new Resolver and perform the IP lookup with the given context
-	resolver := net.Resolver{}
-	names, err := resolver.LookupAddr(ctx, ip)
-	if err != nil {
-		return "", "", err
-	}
-	if len(names) == 0 {
-		return "", "", fmt.Errorf("no hostnames found for %s", ip)
-	}
-
-	// Return the first IP address found and the original port
-	return names[0], port, nil
-}
-
 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)")
@@ -174,13 +146,16 @@ func main() {
 		}
 
 		if *noReverseLookup {
-			log.Printf("PCAP-over-IP connection from %v", conn.RemoteAddr())
+			log.Info().Msgf("PCAP-over-IP connection from %v", conn.RemoteAddr())
 		} else {
-			ip, port, err := lookupHostnameWithTimeout(conn.RemoteAddr(), 100*time.Millisecond)
-			if err != nil {
-				log.Printf("PCAP-over-IP connection from %v", conn.RemoteAddr())
+			ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond)
+			defer cancel()
+			ipAddr := conn.RemoteAddr().(*net.TCPAddr).IP.String()
+			names, _ := net.DefaultResolver.LookupAddr(ctx, ipAddr)
+			if len(names) == 0 {
+				log.Info().Msgf("PCAP-over-IP connection from %v", conn.RemoteAddr())
 			} else {
-				log.Printf("PCAP-over-IP connection from %s:%s", ip, port)
+				log.Info().Msgf("PCAP-over-IP connection from %v (%v)", conn.RemoteAddr(), names[0])
 			}
 		}