From 87e9e3db2869fb8e5a6bdcef00ca75caba843231 Mon Sep 17 00:00:00 2001
From: Andrew Dunstall <andydunstall@hotmail.co.uk>
Date: Mon, 22 Apr 2024 07:26:17 +0100
Subject: [PATCH] server: add proxy status

---
 cli/server/command.go         |  6 +++++-
 pkg/conn/conn.go              |  1 +
 pkg/conn/websocket.go         |  4 ++++
 pkg/rpc/stream.go             |  4 ++++
 server/proxy/proxy.go         | 21 +++++++++++++++++++++
 server/proxy/status.go        | 29 +++++++++++++++++++++++++++++
 server/server/admin/server.go |  1 +
 7 files changed, 65 insertions(+), 1 deletion(-)
 create mode 100644 server/proxy/status.go

diff --git a/cli/server/command.go b/cli/server/command.go
index 5b93dc0..f7c0f45 100644
--- a/cli/server/command.go
+++ b/cli/server/command.go
@@ -303,14 +303,18 @@ func run(conf *config.Config, logger *log.Logger) {
 	networkMapStatus := netmap.NewStatus(networkMap)
 	adminServer.AddStatus("/netmap", networkMapStatus)
 
+	p := proxy.NewProxy(networkMap, registry, logger)
 	proxyServer := proxyserver.NewServer(
 		conf.Proxy.BindAddr,
-		proxy.NewProxy(networkMap, registry, logger),
+		p,
 		&conf.Proxy,
 		registry,
 		logger,
 	)
 
+	proxyStatus := proxy.NewStatus(p)
+	adminServer.AddStatus("/proxy", proxyStatus)
+
 	ctx, cancel := signal.NotifyContext(
 		context.Background(), syscall.SIGINT, syscall.SIGTERM,
 	)
diff --git a/pkg/conn/conn.go b/pkg/conn/conn.go
index a816f7a..72ad8ae 100644
--- a/pkg/conn/conn.go
+++ b/pkg/conn/conn.go
@@ -9,5 +9,6 @@ type Conn interface {
 	NextReader() (io.Reader, error)
 	WriteMessage(b []byte) error
 	NextWriter() (io.WriteCloser, error)
+	Addr() string
 	Close() error
 }
diff --git a/pkg/conn/websocket.go b/pkg/conn/websocket.go
index 0f408d2..33ef780 100644
--- a/pkg/conn/websocket.go
+++ b/pkg/conn/websocket.go
@@ -58,6 +58,10 @@ func (t *WebsocketConn) NextWriter() (io.WriteCloser, error) {
 	return t.wsConn.NextWriter(websocket.BinaryMessage)
 }
 
+func (t *WebsocketConn) Addr() string {
+	return t.wsConn.RemoteAddr().String()
+}
+
 func (t *WebsocketConn) Close() error {
 	return t.wsConn.Close()
 }
diff --git a/pkg/rpc/stream.go b/pkg/rpc/stream.go
index 274c4bc..5aa1d9b 100644
--- a/pkg/rpc/stream.go
+++ b/pkg/rpc/stream.go
@@ -73,6 +73,10 @@ func NewStream(conn conn.Conn, handler *Handler, logger *log.Logger) *Stream {
 	return stream
 }
 
+func (s *Stream) Addr() string {
+	return s.conn.Addr()
+}
+
 // RPC sends the given request message to the peer and returns the response or
 // an error.
 //
diff --git a/server/proxy/proxy.go b/server/proxy/proxy.go
index cb0cf08..3772b48 100644
--- a/server/proxy/proxy.go
+++ b/server/proxy/proxy.go
@@ -56,6 +56,14 @@ func (e *localEndpoint) Next() *rpc.Stream {
 	return s
 }
 
+func (e *localEndpoint) UpstreamAddrs() []string {
+	var addrs []string
+	for _, stream := range e.streams {
+		addrs = append(addrs, stream.Addr())
+	}
+	return addrs
+}
+
 // Proxy is responsible for forwarding requests to upstream listeners.
 type Proxy struct {
 	localEndpoints map[string]*localEndpoint
@@ -178,6 +186,19 @@ func (p *Proxy) RemoveUpstream(endpointID string, stream *rpc.Stream) {
 	p.metrics.Listeners.Dec()
 }
 
+// LocalEndpoints returns a mapping from endpoint ID to listener addresses for
+// all local registered endpoints.
+func (p *Proxy) LocalEndpoints() map[string][]string {
+	p.mu.Lock()
+	defer p.mu.Unlock()
+
+	endpoints := make(map[string][]string)
+	for endpointID, endpoint := range p.localEndpoints {
+		endpoints[endpointID] = endpoint.UpstreamAddrs()
+	}
+	return endpoints
+}
+
 func (p *Proxy) request(
 	ctx context.Context,
 	endpointID string,
diff --git a/server/proxy/status.go b/server/proxy/status.go
new file mode 100644
index 0000000..bb220ee
--- /dev/null
+++ b/server/proxy/status.go
@@ -0,0 +1,29 @@
+package proxy
+
+import (
+	"net/http"
+
+	"github.com/andydunstall/pico/server/status"
+	"github.com/gin-gonic/gin"
+)
+
+type Status struct {
+	proxy *Proxy
+}
+
+func NewStatus(proxy *Proxy) *Status {
+	return &Status{
+		proxy: proxy,
+	}
+}
+
+func (s *Status) Register(group *gin.RouterGroup) {
+	group.GET("/endpoints", s.listEndpointsRoute)
+}
+
+func (s *Status) listEndpointsRoute(c *gin.Context) {
+	endpoints := s.proxy.LocalEndpoints()
+	c.JSON(http.StatusOK, endpoints)
+}
+
+var _ status.Handler = &Status{}
diff --git a/server/server/admin/server.go b/server/server/admin/server.go
index 258582c..ee89b89 100644
--- a/server/server/admin/server.go
+++ b/server/server/admin/server.go
@@ -87,6 +87,7 @@ func (s *Server) registerRoutes() {
 }
 
 func (s *Server) healthRoute(c *gin.Context) {
+	c.Status(http.StatusOK)
 }
 
 func (s *Server) panicRoute(c *gin.Context, err any) {