refactor hub requestJsonFromAgent

This commit is contained in:
Henry Dollman
2024-08-18 18:30:44 -04:00
parent 0566433aa1
commit b7934931cf

View File

@@ -5,10 +5,8 @@ import (
"beszel/internal/entities/system" "beszel/internal/entities/system"
"beszel/internal/records" "beszel/internal/records"
"beszel/site" "beszel/site"
"bytes"
"crypto/ed25519" "crypto/ed25519"
"encoding/pem" "encoding/pem"
"errors"
"fmt" "fmt"
"log" "log"
"net/http" "net/http"
@@ -211,7 +209,7 @@ func (h *Hub) startSystemUpdateTicker() {
func (h *Hub) updateSystems() { func (h *Hub) updateSystems() {
records, err := h.app.Dao().FindRecordsByFilter( records, err := h.app.Dao().FindRecordsByFilter(
"2hz5ncl8tizk5nx", // collection "2hz5ncl8tizk5nx", // systems collection
"status != 'paused'", // filter "status != 'paused'", // filter
"updated", // sort "updated", // sort
-1, // limit -1, // limit
@@ -259,9 +257,9 @@ func (h *Hub) updateSystem(record *models.Record) {
h.connectionLock.Unlock() h.connectionLock.Unlock()
} }
// get system stats from agent // get system stats from agent
systemData, err := requestJson(client) var systemData system.CombinedData
if err != nil { if err := requestJsonFromAgent(client, &systemData); err != nil {
if err.Error() == "retry" { if err.Error() == "bad client" {
// if previous connection was closed, try again // if previous connection was closed, try again
h.app.Logger().Error("Existing SSH connection closed. Retrying...", "host", record.GetString("host"), "port", record.GetString("port")) h.app.Logger().Error("Existing SSH connection closed. Retrying...", "host", record.GetString("host"), "port", record.GetString("port"))
h.deleteSystemConnection(record) h.deleteSystemConnection(record)
@@ -300,7 +298,7 @@ func (h *Hub) updateSystem(record *models.Record) {
} }
} }
// set system to status down and close connection // set system to specified status and save record
func (h *Hub) updateSystemStatus(record *models.Record, status string) { func (h *Hub) updateSystemStatus(record *models.Record, status string) {
if record.GetString("status") != status { if record.GetString("status") != status {
record.Set("status", status) record.Set("status", status)
@@ -353,34 +351,32 @@ func (h *Hub) createSSHClientConfig() error {
return nil return nil
} }
func requestJson(client *ssh.Client) (system.CombinedData, error) { func requestJsonFromAgent(client *ssh.Client, systemData *system.CombinedData) error {
session, err := client.NewSession() session, err := client.NewSession()
if err != nil { if err != nil {
return system.CombinedData{}, errors.New("retry") return fmt.Errorf("bad client")
} }
defer session.Close() defer session.Close()
// Create a buffer to capture the output stdout, err := session.StdoutPipe()
var outputBuffer bytes.Buffer if err != nil {
session.Stdout = &outputBuffer return err
}
if err := session.Shell(); err != nil { if err := session.Shell(); err != nil {
return system.CombinedData{}, err return err
} }
err = session.Wait() if err := json.NewDecoder(stdout).Decode(systemData); err != nil {
if err != nil { return err
return system.CombinedData{}, err
} }
// Unmarshal the output into our struct // wait for the session to complete
var systemData system.CombinedData if err := session.Wait(); err != nil {
err = json.Unmarshal(outputBuffer.Bytes(), &systemData) return err
if err != nil {
return system.CombinedData{}, err
} }
return systemData, nil return nil
} }
func (h *Hub) getSSHKey() ([]byte, error) { func (h *Hub) getSSHKey() ([]byte, error) {