Failover: Return early by waiting for a single pong

1 files changed, 26 insertions, 0 deletions

diff --git a/internal/failover/ping.go b/internal/failover/ping.go
index 2ffedd4..15d4adb 100644
--- a/internal/failover/ping.go
+++ b/internal/failover/ping.go
@@ -4,6 +4,7 @@ import (
 	"fmt"
 	"net"
 	"os"
+	"time"
 
 	"golang.org/x/net/icmp"
 	"golang.org/x/net/ipv4"
@@ -27,6 +28,30 @@ func NewPinger(size int) (*Pinger, error) {
 	return &Pinger{listener: l, buffer: make([]byte, size-mtuOverhead)}, nil
 }
 
+func (p Pinger) Read(deadline time.Time) error {
+	// First set the deadline to read
+	err := p.listener.SetReadDeadline(deadline)
+	if err != nil {
+		return err
+	}
+
+	r := make([]byte, 1500)
+	n, _, err := p.listener.ReadFrom(r)
+	if err != nil {
+		return err
+	}
+	got, err := icmp.ParseMessage(ipv4.ICMPTypeEchoReply.Protocol(), r[:n])
+	if err != nil {
+		return err
+	}
+	switch got.Type {
+	case ipv4.ICMPTypeEchoReply:
+		return nil
+	default:
+		return errors.Errorf("Not a ping echo reply, got %+v", got)
+	}
+}
+
 func (p Pinger) Send(gateway string, seq int) error {
 	errorMessage := fmt.Sprintf("failed sending ping, seq %d", seq)
 	// Make a new ICMP message
@@ -47,5 +72,6 @@ func (p Pinger) Send(gateway string, seq int) error {
 	if err != nil {
 		return errors.WrapPrefix(err, errorMessage, 0)
 	}
+
 	return nil
 }