feat: network monitoring from agents (#2266, #1911)

Co-authored-by: Sven van Ginkel <svenvanginkel@icloud.com>
Co-authored-by: xiaomiku01 <xiaomiku01@outlook.com>

internal/hub/systems/network_monitor_sync_test.go

@@ -0,0 +1,144 @@
+//go:build testing
+
+package systems
+
+import (
+	"net/http"
+	"net/http/httptest"
+	"strings"
+	"testing"
+	"time"
+
+	"github.com/blang/semver"
+	"github.com/fxamacker/cbor/v2"
+	"github.com/henrygd/beszel/internal/common"
+	"github.com/henrygd/beszel/internal/entities/monitor"
+	"github.com/henrygd/beszel/internal/hub/ws"
+	"github.com/lxzan/gws"
+	"github.com/pocketbase/pocketbase/core"
+	"github.com/stretchr/testify/require"
+)
+
+type monitorSyncClient struct {
+	gws.BuiltinEventHandler
+	requests chan common.HubRequest[monitor.SyncRequest]
+}
+
+func (c *monitorSyncClient) OnMessage(conn *gws.Conn, message *gws.Message) {
+	defer message.Close()
+	var req common.HubRequest[monitor.SyncRequest]
+	if err := cbor.Unmarshal(message.Bytes(), &req); err != nil {
+		return
+	}
+	c.requests <- req
+	data, _ := cbor.Marshal(monitor.SyncResponse{})
+	response, _ := cbor.Marshal(common.AgentResponse{Id: req.Id, Data: data})
+	_ = conn.WriteMessage(gws.OpcodeBinary, response)
+}
+
+// Avoid the production delayed disconnect notification; these tests explicitly
+// remove each connection from the manager before reconnecting.
+type monitorSyncServer struct{ ws.Handler }
+
+func (*monitorSyncServer) OnClose(*gws.Conn, error) {}
+
+func TestNetworkMonitorSyncSkipsOlderAgents(t *testing.T) {
+	for _, version := range []string{"0.0.0", "0.18.0", "0.19.0"} {
+		t.Run(version, func(t *testing.T) {
+			// No transport: attempting to send any request would fail.
+			sys := &System{agentVersion: semver.MustParse(version)}
+			require.NoError(t, sys.SyncNetworkMonitors(nil))
+			result, err := sys.UpsertNetworkMonitor(monitor.Config{ID: "test"}, true)
+			require.NoError(t, err)
+			require.Nil(t, result)
+			require.NoError(t, sys.DeleteNetworkMonitor("test"))
+		})
+	}
+}
+
+func TestNetworkMonitorReconnectSync(t *testing.T) {
+	for _, change := range []string{"delete", "disable"} {
+		t.Run(change, func(t *testing.T) {
+			sys, app := newTestSystemWithHub(t)
+			record, err := app.FindRecordById("systems", sys.Id)
+			require.NoError(t, err)
+			// Suppress unrelated system-stat requests while exercising reconnects.
+			record.Set("status", paused)
+			require.NoError(t, app.SaveNoValidate(record))
+			collection, err := app.FindCachedCollectionByNameOrId("network_monitors")
+			require.NoError(t, err)
+			probe := core.NewRecord(collection)
+			probe.Load(map[string]any{
+				"system": sys.Id, "target": "localhost", "protocol": "tcp",
+				"port": 80, "interval": 60, "enabled": true,
+			})
+			require.NoError(t, app.SaveNoValidate(probe))
+
+			sm := NewSystemManager(stubHub{app})
+			t.Cleanup(func() {
+				sm.cancel()
+				_ = sm.RemoveSystem(sys.Id)
+				sm.smartFetchMap.StopCleaner()
+				sm.zfsFetchMap.StopCleaner()
+			})
+			version := semver.MustParse("0.20.0")
+			connections := make(chan *ws.WsConn, 1)
+			upgrader := gws.NewUpgrader(&monitorSyncServer{}, nil)
+			server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+				conn, err := upgrader.Upgrade(w, r)
+				if err != nil {
+					t.Error(err)
+					return
+				}
+				wsConn := ws.NewWsConnection(conn, version)
+				conn.Session().Store("wsConn", wsConn)
+				connections <- wsConn
+				conn.ReadLoop()
+			}))
+			t.Cleanup(server.Close)
+			client := &monitorSyncClient{requests: make(chan common.HubRequest[monitor.SyncRequest], 2)}
+			connect := func() monitor.SyncRequest {
+				t.Helper()
+				conn, _, err := gws.NewClient(client, &gws.ClientOption{Addr: "ws" + strings.TrimPrefix(server.URL, "http")})
+				require.NoError(t, err)
+				t.Cleanup(func() { _ = conn.NetConn().Close() })
+				go conn.ReadLoop()
+				select {
+				case wsConn := <-connections:
+					require.NoError(t, sm.AddWebSocketSystem(sys.Id, version, wsConn))
+				case <-time.After(3 * time.Second):
+					t.Fatal("websocket connection was not established")
+				}
+				select {
+				case req := <-client.requests:
+					require.Equal(t, common.SyncNetworkMonitors, req.Action)
+					require.Equal(t, monitor.SyncActionReplace, req.Data.Action)
+					return req.Data
+				case <-time.After(3 * time.Second):
+					t.Fatal("reconnected agent did not receive a monitor replacement")
+					return monitor.SyncRequest{}
+				}
+			}
+
+			initial := connect()
+			require.Len(t, initial.Configs, 1)
+			require.Equal(t, probe.Id, initial.Configs[0].ID)
+			require.NoError(t, sm.RemoveSystem(sys.Id))
+			if change == "delete" {
+				require.NoError(t, app.Delete(probe))
+			} else {
+				probe.Set("enabled", false)
+				require.NoError(t, app.SaveNoValidate(probe))
+			}
+			require.Empty(t, connect().Configs, "reconnect must clear the agent's previous probe")
+		})
+	}
+}
+
+func TestGetMonitorConfigsForSystemQueryError(t *testing.T) {
+	sys, app := newTestSystemWithHub(t)
+	_, err := app.DB().NewQuery("DROP TABLE network_monitors").Execute()
+	require.NoError(t, err)
+	_, err = sys.manager.GetMonitorConfigsForSystem(sys.Id)
+	require.Error(t, err, "a failed query must not be treated as an empty monitor set")
+}