SSH Connect back related to #555 and discussion #51

beszel/Makefile

@@ -63,5 +63,13 @@ dev-agent:
 		go run beszel/cmd/agent; \
 	fi
 
+dev-agent-connect-back:
+	@if command -v entr >/dev/null 2>&1; then \
+		find ./cmd/agent/*.go ./internal/agent/*.go | entr -r go run beszel/cmd/agent client; \
+	else \
+		go run beszel/cmd/agent client; \
+	fi
+
+
 # KEY="..." make -j dev
 dev: dev-server dev-hub dev-agent

beszel/cmd/agent/agent.go

@@ -11,14 +11,19 @@ import (
 
 func main() {
 	// handle flags / subcommands
+	isClient := false
 	if len(os.Args) > 1 {
 		switch os.Args[1] {
 		case "-v":
 			fmt.Println(beszel.AppName+"-agent", beszel.Version)
+			os.Exit(0)
 		case "update":
 			agent.Update()
+			os.Exit(0)
+		case "client":
+			isClient = true
 		}
-		os.Exit(0)
+
 	}
 
 	// Try to get the key from the KEY environment variable.
@@ -38,15 +43,44 @@ func main() {
 		}
 	}
 
-	addr := ":45876"
-	// TODO: change env var to ADDR
-	if portEnvVar, exists := agent.GetEnv("PORT"); exists {
-		// allow passing an address in the form of "127.0.0.1:45876"
-		if !strings.Contains(portEnvVar, ":") {
-			portEnvVar = ":" + portEnvVar
+	addr := ":45877"
+
+	envAddr := ""
+	addrEnvVar, specifiedByAddr := agent.GetEnv("ADDR")
+	// Legacy from when PORT was used
+	portEnvVar, specifiedByPort := agent.GetEnv("PORT")
+
+	if specifiedByAddr {
+		envAddr = addrEnvVar
+	} else if specifiedByPort {
+		envAddr = portEnvVar
+	}
+
+	if specifiedByAddr || specifiedByPort {
+		if len(envAddr) == 0 && isClient {
+			log.Fatal("No address specified for client to connect to, ADDR was empty")
+		}
+
+		// allow passing an address in the form of "127.0.0.1:45877"
+		if !strings.Contains(envAddr, ":") && !isClient {
+			envAddr = ":" + envAddr
+		} else if isClient {
+			// set the default port if non is specified for clients
+			envAddr = envAddr + ":45877"
+		}
+
+		addr = envAddr
+
+	} else if isClient {
+		log.Fatal("No address specified for client to connect to (use ADDR env)")
+	}
+
+	if isClient {
+		_, exists := agent.GetEnv("CONNECTION_KEY")
+		if !exists {
+			log.Fatal("Started in client mode without CONNECTION_KEY specified")
 		}
-		addr = portEnvVar
 	}
 
-	agent.NewAgent().Run(pubKey, addr)
+	agent.NewAgent(isClient).Run(pubKey, addr)
 }

beszel/internal/agent/agent.go

@@ -25,12 +25,14 @@ type Agent struct {
 	sensorsWhitelist map[string]struct{}        // List of sensors to monitor
 	systemInfo       system.Info                // Host system info
 	gpuManager       *GPUManager                // Manages GPU data
+	client           bool                       // Connect to server rather than open port
 }
 
-func NewAgent() *Agent {
+func NewAgent(client bool) *Agent {
 	newAgent := &Agent{
 		sensorsContext: context.Background(),
 		fsStats:        make(map[string]*system.FsStats),
+		client:         client,
 	}
 	newAgent.memCalc, _ = GetEnv("MEM_CALC")
 	return newAgent
@@ -97,7 +99,11 @@ func (a *Agent) Run(pubKey []byte, addr string) {
 		slog.Debug("Stats", "data", a.gatherStats())
 	}
 
-	a.startServer(pubKey, addr)
+	if a.client {
+		a.startClient(pubKey, addr)
+	} else {
+		a.startServer(pubKey, addr)
+	}
 }
 
 func (a *Agent) gatherStats() system.CombinedData {

beszel/internal/agent/link.go

@@ -0,0 +1,227 @@
+package agent
+
+import (
+	"beszel"
+	"beszel/internal/entities/system"
+	"crypto/ed25519"
+	"crypto/rand"
+	"crypto/x509"
+	"encoding/json"
+	"encoding/pem"
+	"fmt"
+	"io"
+	"log/slog"
+	"net"
+	"os"
+	"time"
+
+	// We are using this for random stats collection timing
+	insecureRandom "math/rand"
+
+	sshServer "github.com/gliderlabs/ssh"
+	"golang.org/x/crypto/ssh"
+)
+
+func (a *Agent) writeStats(output io.Writer) (system.CombinedData, error) {
+	stats := a.gatherStats()
+	return stats, json.NewEncoder(output).Encode(stats)
+}
+
+func (a *Agent) startServer(pubKey []byte, addr string) {
+
+	sshServer.Handle(func(s sshServer.Session) {
+		if stats, err := a.writeStats(s); err != nil {
+			slog.Error("Error encoding stats", "err", err, "stats", stats)
+			s.Exit(1)
+			return
+		}
+
+		s.Exit(0)
+	})
+
+	slog.Info("Starting SSH server", "address", addr)
+	if err := sshServer.ListenAndServe(addr, nil, sshServer.NoPty(),
+		sshServer.PublicKeyAuth(func(ctx sshServer.Context, key sshServer.PublicKey) bool {
+			allowed, _, _, _, _ := sshServer.ParseAuthorizedKey(pubKey)
+			return sshServer.KeysEqual(key, allowed)
+		}),
+	); err != nil {
+		slog.Error("Error starting SSH server", "err", err)
+		os.Exit(1)
+	}
+}
+
+func (a *Agent) startClient(pubKey []byte, addr string) {
+
+	signer, err := createOrLoadKey("id_ed25519")
+	if err != nil {
+		slog.Error("Failed to load private key: ", "err", err)
+		os.Exit(1)
+	}
+
+	slog.Info("Public Key: ", "fingerprint", ssh.FingerprintSHA256(signer.PublicKey()))
+
+	allowed, _, _, _, err := sshServer.ParseAuthorizedKey(pubKey)
+	if err != nil {
+		slog.Error("Failed to parse server public key: ", "err", err)
+		os.Exit(1)
+	}
+
+	hostname, err := os.Hostname()
+	if err != nil {
+		slog.Warn("failed to get host hostname", "err", err)
+		hostname = "unknown_hostname"
+	}
+
+	apiKey, _ := GetEnv("CONNECTION_KEY")
+
+	c := &ssh.ClientConfig{
+		User: fmt.Sprintf("%s@%s", hostname, apiKey),
+		Auth: []ssh.AuthMethod{
+			ssh.PublicKeys(signer),
+		},
+		HostKeyCallback: ssh.FixedHostKey(allowed),
+		ClientVersion:   "SSH-2.0-" + beszel.AppName + "-" + beszel.Version,
+		Timeout:         10 * time.Second,
+	}
+
+	lastBackoff := 1
+	backoff := func() {
+		currentBackoff := 2 * time.Second * time.Duration(lastBackoff)
+		slog.Info("Waiting to reconnect: ", "backoff", currentBackoff)
+		time.Sleep(currentBackoff)
+		if lastBackoff < 10 {
+			lastBackoff++
+		}
+	}
+
+	var currentConn ssh.Conn
+
+	for {
+		if currentConn != nil {
+			currentConn.Close()
+		}
+		slog.Info("Connecting to beszel SSH server", "address", addr)
+
+		conn, err := net.DialTimeout("tcp", addr, c.Timeout)
+		if err != nil {
+			slog.Error("Failed to connect to server: ", "addr", addr, "err", err)
+			backoff()
+			continue
+		}
+
+		clientConn, chans, reqs, err := ssh.NewClientConn(conn, addr, c)
+		if err != nil {
+			conn.Close()
+			slog.Error("Failed to join to server: ", "addr", addr, "err", err)
+			backoff()
+			continue
+		}
+		currentConn = clientConn
+
+		// Reset backoff on first solid connection
+		lastBackoff = 1
+
+		go ssh.DiscardRequests(reqs)
+		go a.discardChannels(chans)
+
+	Inner:
+		for {
+
+			dataChan, req, err := clientConn.OpenChannel("stats", nil)
+			if err != nil {
+				slog.Debug("failed to send statistics, server did not allow opening of stats channel", "err", err)
+				a.randomSleep(15, 30)
+				break Inner
+			}
+
+			go ssh.DiscardRequests(req)
+
+			if stats, err := a.writeStats(dataChan); err != nil {
+				slog.Error("Error writing stats", "err", err, "stats", stats)
+				a.exit(dataChan, 1)
+				break Inner
+			}
+
+			a.exit(dataChan, 0)
+
+			// Make sure that all clients arent syncing up and sending massive amounts of data at once
+			a.randomSleep(15, 30)
+		}
+
+	}
+
+}
+
+func (a *Agent) exit(channel ssh.Channel, code int) error {
+	status := struct{ Status uint32 }{uint32(code)}
+	_, err := channel.SendRequest("exit-status", false, ssh.Marshal(&status))
+	if err != nil {
+		return err
+	}
+	return channel.Close()
+}
+
+func (a *Agent) discardChannels(chans <-chan ssh.NewChannel) {
+	for c := range chans {
+		c.Reject(ssh.Prohibited, "Clients do not support connections")
+	}
+}
+
+func (a *Agent) randomSleep(min, max int) {
+	variance := insecureRandom.Intn(max)
+	time.Sleep(time.Duration(variance)*time.Second + time.Duration(min))
+}
+
+// https://github.com/NHAS/reverse_ssh/blob/main/internal/server/server.go
+func createOrLoadKey(privateKeyPath string) (ssh.Signer, error) {
+
+	//If we have already created a private key (or there is one in the current directory) dont overwrite/create another one
+	if _, err := os.Stat(privateKeyPath); os.IsNotExist(err) {
+
+		privateKeyPem, err := generatePrivateKey()
+		if err != nil {
+			return nil, fmt.Errorf("unable to generate private key, and no private key specified: %s", err)
+		}
+
+		err = os.WriteFile(privateKeyPath, privateKeyPem, 0600)
+		if err != nil {
+			return nil, fmt.Errorf("unable to write private key to disk: %s", err)
+		}
+	}
+
+	privateBytes, err := os.ReadFile(privateKeyPath)
+	if err != nil {
+		return nil, fmt.Errorf("failed to load private key (%s): %s", privateKeyPath, err)
+	}
+
+	private, err := ssh.ParsePrivateKey(privateBytes)
+	if err != nil {
+		return nil, fmt.Errorf("failed to parse private key: %s", err)
+	}
+
+	return private, nil
+}
+
+// https://github.com/NHAS/reverse_ssh/blob/71420af670aebbe632f35ce8428cbfbd21dc5f53/internal/global.go#L44
+func generatePrivateKey() ([]byte, error) {
+	_, priv, err := ed25519.GenerateKey(rand.Reader)
+	if err != nil {
+		return nil, err
+	}
+
+	// Convert a generated ed25519 key into a PEM block so that the ssh library can ingest it, bit round about tbh
+	bytes, err := x509.MarshalPKCS8PrivateKey(priv)
+	if err != nil {
+		return nil, err
+	}
+
+	privatePem := pem.EncodeToMemory(
+		&pem.Block{
+			Type:  "PRIVATE KEY",
+			Bytes: bytes,
+		},
+	)
+
+	return privatePem, nil
+}

beszel/internal/hub/config.go

@@ -104,6 +104,7 @@ func (h *Hub) syncSystemsWithConfig() error {
 			existingSystem.Set("name", sysConfig.Name)
 			existingSystem.Set("users", sysConfig.Users)
 			existingSystem.Set("port", sysConfig.Port)
+
 			if err := h.Save(existingSystem); err != nil {
 				return err
 			}