From 4cf68c94ddbd60dac20e6c10e7357c83b3553ae2 Mon Sep 17 00:00:00 2001
From: pufferffish <github@bandersnatch.anonaddy.com>
Date: Wed, 10 Apr 2024 02:46:29 +0100
Subject: [PATCH] implement metric endpoint

---
 cmd/wireproxy/main.go | 15 +++++++++++++--
 routine.go            | 41 +++++++++++++++++++++++++++++++++++++++++
 wireguard.go          |  1 +
 3 files changed, 55 insertions(+), 2 deletions(-)

diff --git a/cmd/wireproxy/main.go b/cmd/wireproxy/main.go
index 2bbe4a2..9c8a7ca 100644
--- a/cmd/wireproxy/main.go
+++ b/cmd/wireproxy/main.go
@@ -4,6 +4,7 @@ import (
 	"context"
 	"fmt"
 	"log"
+	"net/http"
 	"os"
 	"os/exec"
 	"os/signal"
@@ -78,6 +79,7 @@ func main() {
 	config := parser.String("c", "config", &argparse.Options{Help: "Path of configuration file"})
 	silent := parser.Flag("s", "silent", &argparse.Options{Help: "Silent mode"})
 	daemon := parser.Flag("d", "daemon", &argparse.Options{Help: "Make wireproxy run in background"})
+	info := parser.String("i", "info", &argparse.Options{Help: "Specify the address and port for exposing health status"})
 	printVerison := parser.Flag("v", "version", &argparse.Options{Help: "Print version"})
 	configTest := parser.Flag("n", "configtest", &argparse.Options{Help: "Configtest mode. Only check the configuration file for validity."})
 
@@ -140,13 +142,22 @@ func main() {
 	// no file access is allowed from now on, only networking
 	pledgeOrPanic("stdio inet dns")
 
-	tnet, err := wireproxy.StartWireguard(conf.Device, logLevel)
+	tun, err := wireproxy.StartWireguard(conf.Device, logLevel)
 	if err != nil {
 		log.Fatal(err)
 	}
 
 	for _, spawner := range conf.Routines {
-		go spawner.SpawnRoutine(tnet)
+		go spawner.SpawnRoutine(tun)
+	}
+
+	if *info != "" {
+		go func() {
+			err := http.ListenAndServe(*info, tun)
+			if err != nil {
+				panic(err)
+			}
+		}()
 	}
 
 	<-ctx.Done()
diff --git a/routine.go b/routine.go
index 13fa944..6f7429f 100644
--- a/routine.go
+++ b/routine.go
@@ -1,15 +1,20 @@
 package wireproxy
 
 import (
+	"bytes"
 	"context"
 	"crypto/subtle"
 	"errors"
+	"golang.zx2c4.com/wireguard/device"
 	"io"
 	"log"
 	"math/rand"
 	"net"
+	"net/http"
 	"os"
+	"path"
 	"strconv"
+	"strings"
 
 	"github.com/sourcegraph/conc"
 	"github.com/things-go/go-socks5"
@@ -32,6 +37,7 @@ type CredentialValidator struct {
 // VirtualTun stores a reference to netstack network and DNS configuration
 type VirtualTun struct {
 	Tnet      *netstack.Net
+	Dev       *device.Device
 	SystemDNS bool
 }
 
@@ -330,3 +336,38 @@ 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":
+		w.WriteHeader(http.StatusOK)
+	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)
+	}
+}
diff --git a/wireguard.go b/wireguard.go
index b98bc35..bd0b9c1 100644
--- a/wireguard.go
+++ b/wireguard.go
@@ -82,6 +82,7 @@ func StartWireguard(conf *DeviceConfig, logLevel int) (*VirtualTun, error) {
 
 	return &VirtualTun{
 		Tnet:      tnet,
+		Dev:       dev,
 		SystemDNS: len(setting.dns) == 0,
 	}, nil
 }