From c7a3d708cc15fc753bb6a9aef6876303d367fca1 Mon Sep 17 00:00:00 2001
From: Charlie Jonas <charlie@charliejonas.co.uk>
Date: Mon, 19 Oct 2020 21:58:12 +0100
Subject: [PATCH] Allow pinger.Stop() to be safely called repeatedly. Fix #94

Signed-off-by: Charlie Jonas <charlie@charliejonas.co.uk>
---
 README.md        | 10 ++++++++--
 cmd/ping/ping.go |  4 +++-
 ping.go          | 36 ++++++++++++++++++++++++++----------
 3 files changed, 37 insertions(+), 13 deletions(-)

diff --git a/README.md b/README.md
index 939b994..fa27d1a 100644
--- a/README.md
+++ b/README.md
@@ -30,7 +30,10 @@ c := make(chan os.Signal, 1)
 signal.Notify(c, os.Interrupt)
 go func() {
 	for _ = range c {
-		pinger.Stop()
+		if err := pinger.Stop(); err != nil {
+			// pinger.Stop() will return an error if called more than once
+			fmt.Println(err.Error())
+		}
 	}
 }()
 
@@ -52,7 +55,10 @@ pinger.Run()
 
 It sends ICMP Echo Request packet(s) and waits for an Echo Reply in
 response. If it receives a response, it calls the `OnRecv` callback.
-When it's finished, it calls the `OnFinish` callback.
+When it's finished, it calls the `OnFinish` callback. Note that whilst
+it is safe to call `pinger.Stop()` repeatedly, it will return errors if
+called more than once. However these errors are purely informational and
+can be safely ignored.
 
 For a full ping example, see
 [cmd/ping/ping.go](https://github.com/go-ping/ping/blob/master/cmd/ping/ping.go)
diff --git a/cmd/ping/ping.go b/cmd/ping/ping.go
index a3f028f..f9f7069 100644
--- a/cmd/ping/ping.go
+++ b/cmd/ping/ping.go
@@ -60,7 +60,9 @@ func main() {
 	signal.Notify(c, os.Interrupt)
 	go func() {
 		for range c {
-			pinger.Stop()
+			if err := pinger.Stop(); err != nil {
+				fmt.Println(err.Error())
+			}
 		}
 	}()
 
diff --git a/ping.go b/ping.go
index 19fe1e7..41ca84b 100644
--- a/ping.go
+++ b/ping.go
@@ -69,8 +69,9 @@ const (
 )
 
 var (
-	ipv4Proto = map[string]string{"icmp": "ip4:icmp", "udp": "udp4"}
-	ipv6Proto = map[string]string{"icmp": "ip6:ipv6-icmp", "udp": "udp6"}
+	ipv4Proto   = map[string]string{"icmp": "ip4:icmp", "udp": "udp4"}
+	ipv6Proto   = map[string]string{"icmp": "ip6:ipv6-icmp", "udp": "udp6"}
+	ErrTornDown = errors.New("Stop() called more than once")
 )
 
 // New returns a new Pinger struct pointer.
@@ -84,7 +85,7 @@ func New(addr string) *Pinger {
 		Tracker:  r.Int63n(math.MaxInt64),
 
 		addr:     addr,
-		done:     make(chan bool),
+		done:     make(chan interface{}),
 		id:       r.Intn(math.MaxInt16),
 		ipaddr:   nil,
 		ipv4:     false,
@@ -140,8 +141,9 @@ type Pinger struct {
 	// Source is the source IP address
 	Source string
 
-	// stop chan bool
-	done chan bool
+	// Cleardown channel and mutex
+	done chan interface{}
+	lock sync.Mutex
 
 	ipaddr *net.IPAddr
 	addr   string
@@ -350,7 +352,7 @@ func (p *Pinger) Run() error {
 			wg.Wait()
 			return nil
 		case <-timeout.C:
-			close(p.done)
+			_ = p.Stop()
 			wg.Wait()
 			return nil
 		case <-interval.C:
@@ -370,15 +372,29 @@ func (p *Pinger) Run() error {
 			}
 		}
 		if p.Count > 0 && p.PacketsRecv >= p.Count {
-			close(p.done)
+			_ = p.Stop()
 			wg.Wait()
 			return nil
 		}
 	}
 }
 
-func (p *Pinger) Stop() {
+func (p *Pinger) Stop() error {
+	p.lock.Lock()
+	defer p.lock.Unlock()
+
+	ok := true
+	select {
+	case _, ok = <-p.done:
+	default:
+	}
+
+	if !ok {
+		return ErrTornDown
+	}
+
 	close(p.done)
+	return nil
 }
 
 func (p *Pinger) finish() {
@@ -466,7 +482,7 @@ func (p *Pinger) recvICMP(
 						// Read timeout
 						continue
 					} else {
-						close(p.done)
+						_ = p.Stop()
 						return err
 					}
 				}
@@ -601,7 +617,7 @@ func (p *Pinger) sendICMP(conn *icmp.PacketConn) error {
 func (p *Pinger) listen(netProto string) (*icmp.PacketConn, error) {
 	conn, err := icmp.ListenPacket(netProto, p.Source)
 	if err != nil {
-		close(p.done)
+		_ = p.Stop()
 		return nil, err
 	}
 	return conn, nil