beszel-ipv6/internal/hub/systems/network_monitor_ssh_test.go

//go:build testing

package systems

import (
	"context"
	"crypto/ed25519"
	"crypto/rand"
	"net"
	"sync/atomic"
	"testing"
	"time"

	"github.com/fxamacker/cbor/v2"
	"github.com/henrygd/beszel/internal/common"
	"github.com/henrygd/beszel/internal/entities/monitor"
	esystem "github.com/henrygd/beszel/internal/entities/system"
	"github.com/henrygd/beszel/internal/hub/expirymap"
	"github.com/pocketbase/pocketbase/core"
	"github.com/stretchr/testify/require"
	"golang.org/x/crypto/ssh"
)

func TestSSHNetworkMonitorReconnectSync(t *testing.T) {
	sys, app := newTestSystemWithHub(t)
	sys.manager.zfsFetchMap = expirymap.New[zfsFetchState](time.Hour)
	t.Cleanup(sys.manager.zfsFetchMap.StopCleaner)
	sys.ctx = context.Background()
	sys.Status = up
	_, key, err := ed25519.GenerateKey(rand.Reader)
	require.NoError(t, err)
	signer, err := ssh.NewSignerFromKey(key)
	require.NoError(t, err)
	config := &ssh.ServerConfig{NoClientAuth: true, ServerVersion: "SSH-2.0-beszel_0.20.0"}
	config.AddHostKey(signer)
	listener, err := net.Listen("tcp", "127.0.0.1:0")
	require.NoError(t, err)
	t.Cleanup(func() { _ = listener.Close() })
	sys.Host, sys.Port, err = net.SplitHostPort(listener.Addr().String())
	require.NoError(t, err)
	sys.manager.sshConfig = &ssh.ClientConfig{User: "test", HostKeyCallback: ssh.InsecureIgnoreHostKey(), Timeout: time.Second}
	t.Cleanup(sys.closeSSHConnection)
	requests := make(chan monitor.SyncRequest, 10)
	var failSync atomic.Bool
	go func() {
		for {
			conn, err := listener.Accept()
			if err != nil {
				return
			}
			go func() {
				server, channels, reqs, err := ssh.NewServerConn(conn, config)
				if err != nil {
					_ = conn.Close()
					return
				}
				defer server.Close()
				go ssh.DiscardRequests(reqs)
				for channel := range channels {
					ch, reqs, err := channel.Accept()
					if err != nil {
						return
					}
					go func() {
						defer ch.Close()
						for req := range reqs {
							if req.Type != "shell" {
								_ = req.Reply(false, nil)
								continue
							}
							_ = req.Reply(true, nil)
							var request common.HubRequest[cbor.RawMessage]
							if cbor.NewDecoder(ch).Decode(&request) != nil {
								return
							}
							response := common.AgentResponse{}
							switch request.Action {
							case common.GetData:
								response.SystemData = &esystem.CombinedData{}
							case common.SyncNetworkMonitors:
								var syncReq monitor.SyncRequest
								if cbor.Unmarshal(request.Data, &syncReq) != nil {
									return
								}
								requests <- syncReq
								if failSync.Load() {
									response.Error = "test sync failure"
								} else {
									response.Data, _ = cbor.Marshal(monitor.SyncResponse{})
								}
							}
							_ = cbor.NewEncoder(ch).Encode(response)
							_, _ = ch.SendRequest("exit-status", false, ssh.Marshal(struct{ Status uint32 }{0}))
							return
						}
					}()
				}
			}()
		}
	}()
	collection, err := app.FindCachedCollectionByNameOrId("network_monitors")
	require.NoError(t, err)
	configs := []monitor.Config{
		{Target: "localhost", Protocol: "tcp", Port: 80, Interval: 60},
		{Target: "localhost", Protocol: "dns", Interval: 60},
		{Target: "localhost", Protocol: "dns", Server: "127.0.0.1:5353", Interval: 60},
	}
	var probes []*core.Record
	for i := range configs {
		cfg := &configs[i]
		probe := core.NewRecord(collection)
		probe.Load(map[string]any{
			"system": sys.Id, "target": cfg.Target, "protocol": cfg.Protocol,
			"port": cfg.Port, "server": cfg.Server, "interval": cfg.Interval, "enabled": true,
		})
		require.NoError(t, app.SaveNoValidate(probe))
		cfg.ID = probe.Id
		probes = append(probes, probe)
	}
	fetch := func() {
		t.Helper()
		_, err := sys.fetchDataFromAgent(common.DataRequestOptions{})
		require.NoError(t, err, "monitor sync failure must not fail stats fetching")
	}
	receive := func() monitor.SyncRequest {
		t.Helper()
		select {
		case req := <-requests:
			require.Equal(t, monitor.SyncActionReplace, req.Action)
			return req
		case <-time.After(time.Second):
			t.Fatal("missing full monitor sync")
			return monitor.SyncRequest{}
		}
	}
	fetch()
	require.ElementsMatch(t, configs, receive().Configs)
	require.False(t, sys.monitorsNeedSync.Load())
	fetch()
	require.Empty(t, requests, "steady-state fetch must not resync")

	// Simulate loss of the agent process/connection and its in-memory monitors.
	require.NoError(t, sys.client.Load().Close())
	fetch()
	require.ElementsMatch(t, configs, receive().Configs)
	require.False(t, sys.monitorsNeedSync.Load())

	// Failed replacements are retried on the next successful stats fetch.
	require.NoError(t, sys.client.Load().Close())
	failSync.Store(true)
	fetch()
	require.ElementsMatch(t, configs, receive().Configs)
	require.True(t, sys.monitorsNeedSync.Load())
	failSync.Store(false)
	fetch()
	require.ElementsMatch(t, configs, receive().Configs)
	require.False(t, sys.monitorsNeedSync.Load())

	for _, probe := range probes {
		probe.Set("enabled", false)
		require.NoError(t, app.SaveNoValidate(probe))
	}
	require.NoError(t, sys.client.Load().Close())
	fetch()
	require.Empty(t, receive().Configs, "empty replacement must clear stale monitors")
}

func TestPendingNetworkMonitorSyncQueryFailure(t *testing.T) {
	sys, app := newTestSystemWithHub(t)
	_, err := app.DB().NewQuery("DROP TABLE network_monitors").Execute()
	require.NoError(t, err)
	sys.monitorsNeedSync.Store(true)
	sys.syncPendingNetworkMonitors()
	require.True(t, sys.monitorsNeedSync.Load())
}