From 0de2a988cc701cb294d917964d05f3ef0b4ff474 Mon Sep 17 00:00:00 2001
From: Steffen Vogel <post@steffenvogel.de>
Date: Thu, 15 Jul 2021 17:54:21 +0200
Subject: [PATCH] kg: add new handler for rendering the topology graph

---
 Dockerfile         |   2 +-
 cmd/kg/handlers.go | 118 +++++++++++++++++++++++++++++++++++++++++++++
 cmd/kg/main.go     |   5 +-
 3 files changed, 121 insertions(+), 4 deletions(-)
 create mode 100644 cmd/kg/handlers.go

diff --git a/Dockerfile b/Dockerfile
index 2b2402f..c688907 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -11,7 +11,7 @@ ARG GOARCH
 ARG ALPINE_VERSION=v3.12
 LABEL maintainer="squat <lserven@gmail.com>"
 RUN echo -e "https://alpine.global.ssl.fastly.net/alpine/$ALPINE_VERSION/main\nhttps://alpine.global.ssl.fastly.net/alpine/$ALPINE_VERSION/community" > /etc/apk/repositories && \
-    apk add --no-cache ipset iptables ip6tables wireguard-tools
+    apk add --no-cache ipset iptables ip6tables wireguard-tools graphviz
 COPY --from=cni bridge host-local loopback portmap /opt/cni/bin/
 COPY bin/linux/$GOARCH/kg /opt/bin/
 ENTRYPOINT ["/opt/bin/kg"]
diff --git a/cmd/kg/handlers.go b/cmd/kg/handlers.go
new file mode 100644
index 0000000..0072429
--- /dev/null
+++ b/cmd/kg/handlers.go
@@ -0,0 +1,118 @@
+package main
+
+import (
+	"bytes"
+	"fmt"
+	"io"
+	"mime"
+	"net/http"
+	"os"
+	"os/exec"
+
+	"github.com/squat/kilo/pkg/mesh"
+)
+
+type GraphHandler struct {
+	mesh        *mesh.Mesh
+	granularity mesh.Granularity
+}
+
+func (h *GraphHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
+	ns, err := h.mesh.Nodes().List()
+	if err != nil {
+		http.Error(w, fmt.Sprintf("failed to list nodes: %v", err), 500)
+		return
+	}
+	ps, err := h.mesh.Peers().List()
+	if err != nil {
+		http.Error(w, fmt.Sprintf("failed to list peers: %v", err), 500)
+		return
+	}
+
+	var hostname string
+	subnet := mesh.DefaultKiloSubnet
+	nodes := make(map[string]*mesh.Node)
+	for _, n := range ns {
+		if n.Ready() {
+			nodes[n.Name] = n
+			hostname = n.Name
+		}
+		if n.WireGuardIP != nil {
+			subnet = n.WireGuardIP
+		}
+	}
+	subnet.IP = subnet.IP.Mask(subnet.Mask)
+	if len(nodes) == 0 {
+		http.Error(w, "did not find any valid Kilo nodes in the cluster", 500)
+		return
+	}
+	peers := make(map[string]*mesh.Peer)
+	for _, p := range ps {
+		if p.Ready() {
+			peers[p.Name] = p
+		}
+	}
+	topo, err := mesh.NewTopology(nodes, peers, h.granularity, hostname, 0, []byte{}, subnet, nodes[hostname].PersistentKeepalive, nil)
+	if err != nil {
+		http.Error(w, fmt.Sprintf("failed to create topology: %v", err), 500)
+		return
+	}
+
+	dot, err := topo.Dot()
+	if err != nil {
+		http.Error(w, fmt.Sprintf("failed to generate graph: %v", err), 500)
+	}
+
+	buf := bytes.NewBufferString(dot)
+
+	format := r.URL.Query().Get("format")
+	if format == "" {
+		format = "png"
+	} else if format == ".dot" || format == ".gv" {
+		// If the raw dot data is requested, return it as string.
+		// This allows client-side rendering rather than server-side.
+		w.Write(buf.Bytes())
+		return
+	}
+
+	command := exec.Command("dot", "-T"+format)
+	command.Stderr = os.Stderr
+
+	stdin, err := command.StdinPipe()
+	if err != nil {
+		http.Error(w, err.Error(), 500)
+		return
+	}
+
+	_, err = io.Copy(stdin, buf)
+	if err != nil {
+		http.Error(w, err.Error(), 500)
+		return
+	}
+
+	err = stdin.Close()
+	if err != nil {
+		http.Error(w, err.Error(), 500)
+		return
+	}
+
+	output, err := command.Output()
+	if err != nil {
+		http.Error(w, fmt.Sprintf("unable to execute dot: %v (is graphviz package installed?)", err), 500)
+		return
+	}
+
+	mimeType := mime.TypeByExtension("." + format)
+	if mimeType == "" {
+		mimeType = "application/octet-stream"
+	}
+
+	w.Write(output)
+}
+
+type HealthHandler struct {
+}
+
+func (h *HealthHandler) ServeHTTP(w http.ResponseWriter, _ *http.Request) {
+	w.WriteHeader(http.StatusOK)
+}
diff --git a/cmd/kg/main.go b/cmd/kg/main.go
index fa29883..94b7e2f 100644
--- a/cmd/kg/main.go
+++ b/cmd/kg/main.go
@@ -196,9 +196,8 @@ func Main() error {
 	{
 		// Run the HTTP server.
 		mux := http.NewServeMux()
-		mux.HandleFunc("/health", func(w http.ResponseWriter, _ *http.Request) {
-			w.WriteHeader(http.StatusOK)
-		})
+		mux.Handle("/health", &HealthHandler{})
+		mux.Handle("/graph", &GraphHandler{m, gr})
 		mux.Handle("/metrics", promhttp.HandlerFor(r, promhttp.HandlerOpts{}))
 		l, err := net.Listen("tcp", *listen)
 		if err != nil {