From 4c84fbe55ca50cff1352cea8842354719b2b1b68 Mon Sep 17 00:00:00 2001
From: Tim Foerster <tim.foerster@hetzner.com>
Date: Wed, 20 Feb 2019 23:20:15 +0100
Subject: [PATCH] mtr: Add ptr lookups

---
 cli/cli.go     | 16 ++++++----------
 pkg/hop/hop.go | 31 +++++++++++++++++++++++++------
 pkg/mtr/mtr.go | 16 +++++++++++++---
 3 files changed, 44 insertions(+), 19 deletions(-)

diff --git a/cli/cli.go b/cli/cli.go
index 9267641..03a6282 100644
--- a/cli/cli.go
+++ b/cli/cli.go
@@ -2,8 +2,6 @@ package cli
 
 import (
 	"fmt"
-	"log"
-	"net"
 	"sync"
 	"time"
 
@@ -24,6 +22,7 @@ var (
 	MAX_HOPS         = 64
 	MAX_UNKNOWN_HOPS = 10
 	RING_BUFFER_SIZE = 50
+	PTR_LOOKUP       = false
 	jsonFmt          = false
 	srcAddr          = "0.0.0.0"
 	versionFlag      bool
@@ -38,15 +37,11 @@ var RootCmd = &cobra.Command{
 			fmt.Printf("MTR Version: %s, build date: %s\n", version, date)
 			return nil
 		}
-		ip := args[0]
-		if net.ParseIP(args[0]) == nil {
-			addrs, err := net.LookupHost(args[0])
-			if err != nil || len(addrs) == 0 {
-				log.Fatalf("invalid host or ip provided: %s", err)
-			}
-			ip = addrs[0]
+		m, ch, err := mtr.NewMTR(args[0], srcAddr, TIMEOUT, INTERVAL, HOP_SLEEP,
+			MAX_HOPS, MAX_UNKNOWN_HOPS, RING_BUFFER_SIZE, PTR_LOOKUP)
+		if err != nil {
+			return err
 		}
-		m, ch := mtr.NewMTR(ip, srcAddr, TIMEOUT, INTERVAL, HOP_SLEEP, MAX_HOPS, MAX_UNKNOWN_HOPS, RING_BUFFER_SIZE)
 		if jsonFmt {
 			go func(ch chan struct{}) {
 				for {
@@ -93,6 +88,7 @@ func init() {
 	RootCmd.Flags().IntVar(&MAX_UNKNOWN_HOPS, "max-unknown-hops", MAX_UNKNOWN_HOPS, "Maximal hops that do not reply before stopping to look")
 	RootCmd.Flags().IntVar(&RING_BUFFER_SIZE, "buffer-size", RING_BUFFER_SIZE, "Cached packet buffer size")
 	RootCmd.Flags().BoolVar(&jsonFmt, "json", jsonFmt, "Print json results")
+	RootCmd.Flags().BoolVarP(&PTR_LOOKUP, "ptr", "n", PTR_LOOKUP, "Reverse lookup on host")
 	RootCmd.Flags().BoolVar(&versionFlag, "version", false, "Print version")
 	RootCmd.Flags().StringVar(&srcAddr, "address", srcAddr, "The address to be bound the outgoing socket")
 }
diff --git a/pkg/hop/hop.go b/pkg/hop/hop.go
index 2729fdc..ab04ca9 100644
--- a/pkg/hop/hop.go
+++ b/pkg/hop/hop.go
@@ -26,6 +26,7 @@ type HopStatistic struct {
 	Packets        *ring.Ring
 	RingBufferSize int
 	pingSeq        int
+	dnsCache       map[string]string
 }
 
 type packet struct {
@@ -124,7 +125,10 @@ func (h *HopStatistic) packets() []*packet {
 	return v
 }
 
-func (h *HopStatistic) Render() {
+func (h *HopStatistic) Render(ptrLookup bool) {
+	if h.dnsCache == nil {
+		h.dnsCache = map[string]string{}
+	}
 	packets := make([]byte, h.RingBufferSize)
 	i := h.RingBufferSize - 1
 	h.Packets.Do(func(f interface{}) {
@@ -137,14 +141,10 @@ func (h *HopStatistic) Render() {
 		}
 		i--
 	})
-	addr := "???"
-	if h.Target != "" {
-		addr = h.Target
-	}
 	l := fmt.Sprintf("%d", h.RingBufferSize)
 	gm.Printf("%3d:|-- %-20s  %5.1f%%  %4d  %6.1f  %6.1f  %6.1f  %6.1f  %"+l+"s\n",
 		h.TTL,
-		addr,
+		fmt.Sprintf("%.20s", h.lookupAddr(ptrLookup)),
 		h.Loss(),
 		h.Sent,
 		h.Last.Elapsed.Seconds()*1000,
@@ -154,3 +154,22 @@ func (h *HopStatistic) Render() {
 		packets,
 	)
 }
+
+func (h *HopStatistic) lookupAddr(ptrLookup bool) string {
+	addr := "???"
+	if h.Target != "" {
+		addr = h.Target
+		if ptrLookup {
+			if key, ok := h.dnsCache[h.Target]; ok {
+				addr = key
+			} else {
+				names, err := net.LookupAddr(h.Target)
+				if err == nil && len(names) > 0 {
+					addr = names[0]
+				}
+			}
+		}
+		h.dnsCache[h.Target] = addr
+	}
+	return addr
+}
diff --git a/pkg/mtr/mtr.go b/pkg/mtr/mtr.go
index bf1f29e..4d3d4c1 100644
--- a/pkg/mtr/mtr.go
+++ b/pkg/mtr/mtr.go
@@ -24,9 +24,18 @@ type MTR struct {
 	ringBufferSize int
 	maxHops        int
 	maxUnknownHops int
+	ptrLookup      bool
 }
 
-func NewMTR(addr, srcAddr string, timeout time.Duration, interval time.Duration, hopsleep time.Duration, maxHops, maxUnknownHops, ringBufferSize int) (*MTR, chan struct{}) {
+func NewMTR(addr, srcAddr string, timeout time.Duration, interval time.Duration,
+	hopsleep time.Duration, maxHops, maxUnknownHops, ringBufferSize int, ptr bool) (*MTR, chan struct{}, error) {
+	if net.ParseIP(addr) == nil {
+		addrs, err := net.LookupHost(addr)
+		if err != nil || len(addrs) == 0 {
+			return nil, nil, fmt.Errorf("invalid host or ip provided: %s", err)
+		}
+		addr = addrs[0]
+	}
 	return &MTR{
 		SrcAddress:     srcAddr,
 		interval:       interval,
@@ -38,7 +47,8 @@ func NewMTR(addr, srcAddr string, timeout time.Duration, interval time.Duration,
 		maxHops:        maxHops,
 		ringBufferSize: ringBufferSize,
 		maxUnknownHops: maxUnknownHops,
-	}, make(chan struct{})
+		ptrLookup:      ptr,
+	}, make(chan struct{}), nil
 }
 
 func (m *MTR) registerStatistic(ttl int, r icmp.ICMPReturn) *hop.HopStatistic {
@@ -69,7 +79,7 @@ func (m *MTR) Render(offset int) {
 	for i := 1; i <= len(m.Statistic); i++ {
 		gm.MoveCursor(1, offset+i)
 		m.mutex.RLock()
-		m.Statistic[i].Render()
+		m.Statistic[i].Render(m.ptrLookup)
 		m.mutex.RUnlock()
 	}
 }