From e4801c3989ec86ca9ff5541f5cb1dd3a1b1d073c Mon Sep 17 00:00:00 2001
From: xjasonlyu <xjasonlyu@gmail.com>
Date: Mon, 4 Apr 2022 22:05:47 +0800
Subject: [PATCH] Fix(restapi): debug missing mountpoints

---
 restapi/connections.go |  4 ++++
 restapi/netstats.go    |  4 ++++
 restapi/server.go      | 22 ++++++++++++++++------
 3 files changed, 24 insertions(+), 6 deletions(-)

diff --git a/restapi/connections.go b/restapi/connections.go
index da4157e..8f1e7aa 100644
--- a/restapi/connections.go
+++ b/restapi/connections.go
@@ -16,6 +16,10 @@ import (
 
 const defaultInterval = 1000
 
+func init() {
+	registerMountPoint("/connections", connectionRouter())
+}
+
 func connectionRouter() http.Handler {
 	r := chi.NewRouter()
 	r.Get("/", getConnections)
diff --git a/restapi/netstats.go b/restapi/netstats.go
index d0dc8c2..d22e30f 100644
--- a/restapi/netstats.go
+++ b/restapi/netstats.go
@@ -18,6 +18,10 @@ func SetStatsFunc(s func() tcpip.Stats) {
 	_statsFunc = s
 }
 
+func init() {
+	registerMountPoint("/netstats", http.HandlerFunc(getNetStats))
+}
+
 func getNetStats(w http.ResponseWriter, r *http.Request) {
 	if _statsFunc == nil {
 		render.Status(r, http.StatusInternalServerError)
diff --git a/restapi/server.go b/restapi/server.go
index 57d26e3..8f18fdb 100644
--- a/restapi/server.go
+++ b/restapi/server.go
@@ -20,10 +20,18 @@ import (
 	"github.com/gorilla/websocket"
 )
 
-var _upgrader = websocket.Upgrader{
-	CheckOrigin: func(r *http.Request) bool {
-		return true
-	},
+var (
+	_upgrader = websocket.Upgrader{
+		CheckOrigin: func(r *http.Request) bool {
+			return true
+		},
+	}
+
+	_mountPoints = make(map[string]http.Handler)
+)
+
+func registerMountPoint(pattern string, handler http.Handler) {
+	_mountPoints[pattern] = handler
 }
 
 func Start(addr, token string) error {
@@ -43,8 +51,10 @@ func Start(addr, token string) error {
 		r.Get("/logs", getLogs)
 		r.Get("/traffic", traffic)
 		r.Get("/version", version)
-		r.Get("/netstats", getNetStats)
-		r.Mount("/connections", connectionRouter())
+		// attach HTTP handlers
+		for pattern, handler := range _mountPoints {
+			r.Mount(pattern, handler)
+		}
 	})
 
 	listener, err := net.Listen("tcp", addr)