Add health status endpoint (#107)

* implement metric endpoint

* implement ICMP ping

* fix linting

* fix IPv6 pings

* Add documentation for --info

routine.go

@@ -1,15 +1,27 @@
 package wireproxy
 
 import (
+	"bytes"
 	"context"
+	srand "crypto/rand"
 	"crypto/subtle"
+	"encoding/binary"
+	"encoding/json"
 	"errors"
+	"golang.org/x/net/icmp"
+	"golang.org/x/net/ipv4"
+	"golang.org/x/net/ipv6"
+	"golang.zx2c4.com/wireguard/device"
 	"io"
 	"log"
 	"math/rand"
 	"net"
+	"net/http"
 	"os"
+	"path"
 	"strconv"
+	"strings"
+	"time"
 
 	"github.com/sourcegraph/conc"
 	"github.com/things-go/go-socks5"
@@ -32,7 +44,11 @@ type CredentialValidator struct {
 // VirtualTun stores a reference to netstack network and DNS configuration
 type VirtualTun struct {
 	Tnet      *netstack.Net
+	Dev       *device.Device
 	SystemDNS bool
+	Conf      *DeviceConfig
+	// PingRecord stores the last time an IP was pinged
+	PingRecord map[string]uint64
 }
 
 // RoutineSpawner spawns a routine (e.g. socks5, tcp static routes) after the configuration is parsed
@@ -148,16 +164,16 @@ func (config *Socks5Config) SpawnRoutine(vt *VirtualTun) {
 
 // SpawnRoutine spawns a http server.
 func (config *HTTPConfig) SpawnRoutine(vt *VirtualTun) {
-	http := &HTTPServer{
+	server := &HTTPServer{
 		config: config,
 		dial:   vt.Tnet.Dial,
 		auth:   CredentialValidator{config.Username, config.Password},
 	}
 	if config.Username != "" || config.Password != "" {
-		http.authRequired = true
+		server.authRequired = true
 	}
 
-	if err := http.ListenAndServe("tcp", config.BindAddress); err != nil {
+	if err := server.ListenAndServe("tcp", config.BindAddress); err != nil {
 		log.Fatal(err)
 	}
 }
@@ -330,3 +346,151 @@ func (conf *TCPServerTunnelConfig) SpawnRoutine(vt *VirtualTun) {
 		go tcpServerForward(vt, raddr, conn)
 	}
 }
+
+func (d VirtualTun) ServeHTTP(w http.ResponseWriter, r *http.Request) {
+	log.Printf("Health metric request: %s\n", r.URL.Path)
+	switch path.Clean(r.URL.Path) {
+	case "/readyz":
+		body, err := json.Marshal(d.PingRecord)
+		if err != nil {
+			errorLogger.Printf("Failed to get device metrics: %s\n", err.Error())
+			w.WriteHeader(http.StatusInternalServerError)
+			return
+		}
+
+		status := http.StatusOK
+		for _, record := range d.PingRecord {
+			lastPong := time.Unix(int64(record), 0)
+			// +2 seconds to account for the time it takes to ping the IP
+			if time.Since(lastPong) > time.Duration(d.Conf.CheckAliveInterval+2)*time.Second {
+				status = http.StatusServiceUnavailable
+				break
+			}
+		}
+
+		w.WriteHeader(status)
+		_, _ = w.Write(body)
+		_, _ = w.Write([]byte("\n"))
+	case "/metrics":
+		get, err := d.Dev.IpcGet()
+		if err != nil {
+			errorLogger.Printf("Failed to get device metrics: %s\n", err.Error())
+			w.WriteHeader(http.StatusInternalServerError)
+			return
+		}
+		var buf bytes.Buffer
+		for _, peer := range strings.Split(get, "\n") {
+			pair := strings.SplitN(peer, "=", 2)
+			if len(pair) != 2 {
+				buf.WriteString(peer)
+				continue
+			}
+			if pair[0] == "private_key" || pair[0] == "preshared_key" {
+				pair[1] = "REDACTED"
+			}
+			buf.WriteString(pair[0])
+			buf.WriteString("=")
+			buf.WriteString(pair[1])
+			buf.WriteString("\n")
+		}
+
+		w.WriteHeader(http.StatusOK)
+		_, _ = w.Write(buf.Bytes())
+	default:
+		w.WriteHeader(http.StatusNotFound)
+	}
+}
+
+func (d VirtualTun) pingIPs() {
+	for _, addr := range d.Conf.CheckAlive {
+		socket, err := d.Tnet.Dial("ping", addr.String())
+		if err != nil {
+			errorLogger.Printf("Failed to ping %s: %s\n", addr, err.Error())
+			continue
+		}
+
+		data := make([]byte, 16)
+		_, _ = srand.Read(data)
+
+		requestPing := icmp.Echo{
+			Seq:  rand.Intn(1 << 16),
+			Data: data,
+		}
+
+		var icmpBytes []byte
+		if addr.Is4() {
+			icmpBytes, _ = (&icmp.Message{Type: ipv4.ICMPTypeEcho, Code: 0, Body: &requestPing}).Marshal(nil)
+		} else if addr.Is6() {
+			icmpBytes, _ = (&icmp.Message{Type: ipv6.ICMPTypeEchoRequest, Code: 0, Body: &requestPing}).Marshal(nil)
+		} else {
+			errorLogger.Printf("Failed to ping %s: invalid address: %s\n", addr, addr.String())
+			continue
+		}
+
+		_ = socket.SetReadDeadline(time.Now().Add(time.Duration(d.Conf.CheckAliveInterval) * time.Second))
+		_, err = socket.Write(icmpBytes)
+		if err != nil {
+			errorLogger.Printf("Failed to ping %s: %s\n", addr, err.Error())
+			continue
+		}
+
+		addr := addr
+		go func() {
+			n, err := socket.Read(icmpBytes[:])
+			if err != nil {
+				errorLogger.Printf("Failed to read ping response from %s: %s\n", addr, err.Error())
+				return
+			}
+
+			replyPacket, err := icmp.ParseMessage(1, icmpBytes[:n])
+			if err != nil {
+				errorLogger.Printf("Failed to parse ping response from %s: %s\n", addr, err.Error())
+				return
+			}
+
+			if addr.Is4() {
+				replyPing, ok := replyPacket.Body.(*icmp.Echo)
+				if !ok {
+					errorLogger.Printf("Failed to parse ping response from %s: invalid reply type: %s\n", addr, replyPacket.Type)
+					return
+				}
+				if !bytes.Equal(replyPing.Data, requestPing.Data) || replyPing.Seq != requestPing.Seq {
+					errorLogger.Printf("Failed to parse ping response from %s: invalid ping reply: %v\n", addr, replyPing)
+					return
+				}
+			}
+
+			if addr.Is6() {
+				replyPing, ok := replyPacket.Body.(*icmp.RawBody)
+				if !ok {
+					errorLogger.Printf("Failed to parse ping response from %s: invalid reply type: %s\n", addr, replyPacket.Type)
+					return
+				}
+
+				seq := binary.BigEndian.Uint16(replyPing.Data[2:4])
+				pongBody := replyPing.Data[4:]
+				if !bytes.Equal(pongBody, requestPing.Data) || int(seq) != requestPing.Seq {
+					errorLogger.Printf("Failed to parse ping response from %s: invalid ping reply: %v\n", addr, replyPing)
+					return
+				}
+			}
+
+			d.PingRecord[addr.String()] = uint64(time.Now().Unix())
+
+			defer socket.Close()
+		}()
+	}
+}
+
+func (d VirtualTun) StartPingIPs() {
+	for _, addr := range d.Conf.CheckAlive {
+		d.PingRecord[addr.String()] = 0
+	}
+
+	go func() {
+		for {
+			d.pingIPs()
+			time.Sleep(time.Duration(d.Conf.CheckAliveInterval) * time.Second)
+		}
+	}()
+}