diff --git a/internal/hub/systems/system.go b/internal/hub/systems/system.go index 469c1e06..58206b93 100644 --- a/internal/hub/systems/system.go +++ b/internal/hub/systems/system.go @@ -153,6 +153,7 @@ func (sys *System) update() error { // ensure deprecated fields from older agents are migrated to current fields migrateDeprecatedFields(data, !sys.detailsFetched.Load()) + sys.data = data // create system records _, err = sys.createRecords(data) @@ -704,11 +705,9 @@ func (sys *System) ensureSSHTransport() error { } // fetchDataFromAgent attempts to fetch data from the agent, prioritizing WebSocket if available. +// Each fetch decodes into a new struct: CBOR leaves fields the agent omits +// untouched, and real-time and regular updates may fetch concurrently. func (sys *System) fetchDataFromAgent(options common.DataRequestOptions) (*system.CombinedData, error) { - if sys.data == nil { - sys.data = &system.CombinedData{} - } - if sys.WsConn != nil && sys.WsConn.IsConnected() { wsData, err := sys.fetchDataViaWebSocket(options) if err == nil { @@ -744,11 +743,11 @@ func (sys *System) fetchDataViaWebSocket(options common.DataRequestOptions) (*sy ctx, cancel := context.WithTimeout(context.Background(), wsDataRequestTimeout) defer cancel() wsTransport := transport.NewWebSocketTransport(sys.WsConn) - err := wsTransport.Request(ctx, common.GetData, options, sys.data) - if err != nil { + data := &system.CombinedData{} + if err := wsTransport.Request(ctx, common.GetData, options, data); err != nil { return nil, err } - return sys.data, nil + return data, nil } // FetchContainerInfoFromAgent fetches container info from the agent @@ -810,9 +809,8 @@ func MakeStableHashId(strings ...string) string { } // fetchDataViaSSH handles fetching data using SSH. -// This function encapsulates the original SSH logic. -// It updates sys.data directly upon successful fetch. func (sys *System) fetchDataViaSSH(options common.DataRequestOptions) (*system.CombinedData, error) { + data := &system.CombinedData{} err := sys.runSSHOperation(4*time.Second, 1, func(session *ssh.Session) (bool, error) { stdout, err := session.StdoutPipe() if err != nil { @@ -823,7 +821,8 @@ func (sys *System) fetchDataViaSSH(options common.DataRequestOptions) (*system.C return false, err } - *sys.data = system.CombinedData{} + // reset in case of retry after a partial decode + *data = system.CombinedData{} if sys.agentVersion.GTE(beszel.MinVersionAgentResponse) && stdinErr == nil { req := common.HubRequest[any]{Action: common.GetData, Data: options} @@ -832,7 +831,7 @@ func (sys *System) fetchDataViaSSH(options common.DataRequestOptions) (*system.C var resp common.AgentResponse if decErr := cbor.NewDecoder(stdout).Decode(&resp); decErr == nil && resp.SystemData != nil { - *sys.data = *resp.SystemData + *data = *resp.SystemData if err := session.Wait(); err != nil { return false, err } @@ -842,9 +841,9 @@ func (sys *System) fetchDataViaSSH(options common.DataRequestOptions) (*system.C var decodeErr error if sys.agentVersion.GTE(beszel.MinVersionCbor) { - decodeErr = cbor.NewDecoder(stdout).Decode(sys.data) + decodeErr = cbor.NewDecoder(stdout).Decode(data) } else { - decodeErr = json.NewDecoder(stdout).Decode(sys.data) + decodeErr = json.NewDecoder(stdout).Decode(data) } if decodeErr != nil { @@ -861,7 +860,7 @@ func (sys *System) fetchDataViaSSH(options common.DataRequestOptions) (*system.C return nil, err } - return sys.data, nil + return data, nil } // runSSHOperation establishes an SSH session and executes the provided operation. diff --git a/internal/hub/systems/ws_fresh_data_test.go b/internal/hub/systems/ws_fresh_data_test.go new file mode 100644 index 00000000..b2f386b9 --- /dev/null +++ b/internal/hub/systems/ws_fresh_data_test.go @@ -0,0 +1,92 @@ +//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" + esystem "github.com/henrygd/beszel/internal/entities/system" + "github.com/henrygd/beszel/internal/hub/ws" + "github.com/lxzan/gws" + "github.com/stretchr/testify/require" +) + +// sequenceDataClient answers each GetData request with the next queued payload. +type sequenceDataClient struct { + gws.BuiltinEventHandler + responses chan esystem.CombinedData +} + +func (c *sequenceDataClient) OnMessage(conn *gws.Conn, message *gws.Message) { + defer message.Close() + var req common.HubRequest[cbor.RawMessage] + if err := cbor.Unmarshal(message.Bytes(), &req); err != nil || req.Action != common.GetData { + return + } + data, _ := cbor.Marshal(<-c.responses) + response, _ := cbor.Marshal(common.AgentResponse{Id: req.Id, Data: data}) + _ = conn.WriteMessage(gws.OpcodeBinary, response) +} + +// Fields the agent omits must not carry over from a previous response. +func TestFetchDataViaWebSocketDoesNotRetainOmittedFields(t *testing.T) { + 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, semver.MustParse("0.20.0")) + conn.Session().Store("wsConn", wsConn) + connections <- wsConn + conn.ReadLoop() + })) + t.Cleanup(server.Close) + + client := &sequenceDataClient{responses: make(chan esystem.CombinedData, 2)} + client.responses <- esystem.CombinedData{ + Details: &esystem.Details{Hostname: "host"}, + SystemdServicesUpdated: true, + Stats: esystem.Stats{Batteries: map[string]uint8{"BAT0": 50, "BAT1": 60}}, + Info: esystem.Info{WiFi: map[string]esystem.WiFi{"wlan0": {SSID: "home"}}}, + } + client.responses <- esystem.CombinedData{ + Stats: esystem.Stats{Batteries: map[string]uint8{"BAT0": 40}}, + } + 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() + + var sys *System + select { + case wsConn := <-connections: + sys = &System{WsConn: wsConn} + case <-time.After(3 * time.Second): + t.Fatal("websocket connection was not established") + } + + first, err := sys.fetchDataFromAgent(common.DataRequestOptions{}) + require.NoError(t, err) + require.NotNil(t, first.Details) + require.Len(t, first.Stats.Batteries, 2) + + second, err := sys.fetchDataFromAgent(common.DataRequestOptions{}) + require.NoError(t, err) + require.Nil(t, second.Details) + require.False(t, second.SystemdServicesUpdated) + require.Equal(t, map[string]uint8{"BAT0": 40}, second.Stats.Batteries) + require.Empty(t, second.Info.WiFi) + + // the first result is not mutated by the second fetch + require.Len(t, first.Stats.Batteries, 2) +} diff --git a/internal/hub/transport/transport.go b/internal/hub/transport/transport.go index 1a956b49..8f99940f 100644 --- a/internal/hub/transport/transport.go +++ b/internal/hub/transport/transport.go @@ -35,14 +35,6 @@ func UnmarshalResponse(resp common.AgentResponse, action common.WebSocketAction, } // Try generic Data field first (0.19+) if len(resp.Data) > 0 { - // Wi-Fi maps are complete snapshots. CBOR otherwise merges entries into - // reused destinations, retaining disconnected interfaces and old RSSI. - if action == common.GetData { - if data, ok := dest.(*system.CombinedData); ok { - data.Info.WiFi = nil - data.Stats.WiFi = nil - } - } if err := cbor.Unmarshal(resp.Data, dest); err != nil { return fmt.Errorf("failed to unmarshal generic response data: %w", err) } diff --git a/internal/hub/transport/wifi_test.go b/internal/hub/transport/wifi_test.go deleted file mode 100644 index a3a71c84..00000000 --- a/internal/hub/transport/wifi_test.go +++ /dev/null @@ -1,43 +0,0 @@ -package transport - -import ( - "github.com/fxamacker/cbor/v2" - "github.com/henrygd/beszel/internal/common" - "github.com/henrygd/beszel/internal/entities/system" - "github.com/stretchr/testify/require" - "testing" -) - -func TestWiFiSequentialResponseSnapshots(t *testing.T) { - signal := -50.0 - var decoded system.CombinedData - for _, snapshot := range []map[string]system.WiFi{ - {"wlan0": {SSID: "home", Signal: &signal}, "wlan1": {Signal: &signal}}, - {"wlan0": {SSID: "home"}}, {}, nil, - {"wlan1": {SSID: "new", Signal: &signal}}, - } { - signals := make(map[string]int8) - for id, reading := range snapshot { - if reading.Signal != nil { - signals[id] = int8(*reading.Signal) - } - } - payload, err := cbor.Marshal(system.CombinedData{Info: system.Info{WiFi: snapshot}, Stats: system.Stats{WiFi: signals}}) - require.NoError(t, err) - require.NoError(t, UnmarshalResponse(common.AgentResponse{Data: payload}, common.GetData, &decoded)) - require.Len(t, decoded.Info.WiFi, len(snapshot)) - require.Len(t, decoded.Stats.WiFi, len(signals)) - for id, want := range snapshot { - require.Equal(t, want, decoded.Info.WiFi[id]) - } - for id, want := range signals { - require.Equal(t, want, decoded.Stats.WiFi[id]) - } - } - // An older generic-response agent may omit both fields entirely. - payload, err := cbor.Marshal(map[int]any{0: map[int]any{}, 1: map[int]any{}}) - require.NoError(t, err) - require.NoError(t, UnmarshalResponse(common.AgentResponse{Data: payload}, common.GetData, &decoded)) - require.Empty(t, decoded.Info.WiFi) - require.Empty(t, decoded.Stats.WiFi) -}