feat: switch away from websocket to stdout messaging for agent

agent/cmd/hp_agent/hp_agent.go

@@ -8,6 +8,11 @@ import (
 	"github.com/tale/headplane/agent/internal/util"
 )
 
+type Register struct {
+	Type string
+	ID   string
+}
+
 func main() {
 	log := util.GetLogger()
 	cfg, err := config.Load()
@@ -21,11 +26,10 @@ func main() {
 	agent.Connect()
 	defer agent.Shutdown()
 
-	ws, err := hpagent.NewSocket(agent, cfg)
-	if err != nil {
-		log.Fatal("Failed to create websocket: %s", err)
-	}
+	log.Msg(&Register{
+		Type: "register",
+		ID:   agent.ID,
+	})
 
-	defer ws.StopListening()
-	ws.FollowMaster()
+	hpagent.FollowMaster(agent)
 }

agent/internal/config/config.go

@@ -8,8 +8,7 @@ type Config struct {
 	Hostname     string
 	TSControlURL string
 	TSAuthKey    string
-	HPControlURL string
-	HPAuthKey    string
+	WorkDir      string
 }
 
 const (
@@ -17,8 +16,7 @@ const (
 	HostnameEnv     = "HEADPLANE_AGENT_HOSTNAME"
 	TSControlURLEnv = "HEADPLANE_AGENT_TS_SERVER"
 	TSAuthKeyEnv    = "HEADPLANE_AGENT_TS_AUTHKEY"
-	HPControlURLEnv = "HEADPLANE_AGENT_HP_SERVER"
-	HPAuthKeyEnv    = "HEADPLANE_AGENT_HP_AUTHKEY"
+	WorkDirEnv      = "HEADPLANE_AGENT_WORK_DIR"
 )
 
 // Load reads the agent configuration from environment variables.
@@ -28,8 +26,7 @@ func Load() (*Config, error) {
 		Hostname:     os.Getenv(HostnameEnv),
 		TSControlURL: os.Getenv(TSControlURLEnv),
 		TSAuthKey:    os.Getenv(TSAuthKeyEnv),
-		HPControlURL: os.Getenv(HPControlURLEnv),
-		HPAuthKey:    os.Getenv(HPAuthKeyEnv),
+		WorkDir:      os.Getenv(WorkDirEnv),
 	}
 
 	if os.Getenv(DebugEnv) == "true" {
@@ -44,9 +41,5 @@ func Load() (*Config, error) {
 		return nil, err
 	}
 
-	if err := validateHPReady(c); err != nil {
-		return nil, err
-	}
-
 	return c, nil
 }

agent/internal/config/preflight.go

@@ -16,16 +16,12 @@ func validateRequired(config *Config) error {
 		return fmt.Errorf("%s is required", TSControlURLEnv)
 	}
 
-	if config.HPControlURL == "" {
-		return fmt.Errorf("%s is required", HPControlURLEnv)
-	}
-
 	if config.TSAuthKey == "" {
 		return fmt.Errorf("%s is required", TSAuthKeyEnv)
 	}
 
-	if config.HPAuthKey == "" {
-		return fmt.Errorf("%s is required", HPAuthKeyEnv)
+	if config.WorkDir == "" {
+		return fmt.Errorf("%s is required", WorkDirEnv)
 	}
 
 	return nil
@@ -50,23 +46,3 @@ func validateTSReady(config *Config) error {
 
 	return nil
 }
-
-// Pings the Headplane server to make sure it's up and running
-func validateHPReady(config *Config) error {
-	testURL := config.HPControlURL
-	if strings.HasSuffix(testURL, "/") {
-		testURL = testURL[:len(testURL)-1]
-	}
-
-	testURL = fmt.Sprintf("%s/healthz", testURL)
-	resp, err := http.Get(testURL)
-	if err != nil {
-		return fmt.Errorf("Failed to connect to HP control server: %s", err)
-	}
-
-	if resp.StatusCode != 200 {
-		return fmt.Errorf("Failed to connect to HP control server: %s", resp.Status)
-	}
-
-	return nil
-}

agent/internal/hpagent/handler.go

@@ -1,9 +1,12 @@
 package hpagent
 
 import (
+	"bufio"
 	"encoding/json"
+	"os"
 	"sync"
 
+	"github.com/tale/headplane/agent/internal/tsnet"
 	"github.com/tale/headplane/agent/internal/util"
 	"tailscale.com/tailcfg"
 )
@@ -13,30 +16,32 @@ type RecvMessage struct {
 	NodeIDs []string
 }
 
-// Starts listening for messages from the Headplane master
-func (s *Socket) FollowMaster() {
-	log := util.GetLogger()
+type SendMessage struct {
+	Type string
+	Data any
+}
 
-	for {
-		_, message, err := s.ReadMessage()
-		if err != nil {
-			log.Error("Error reading message: %s", err)
-			return
-		}
+// Starts listening for messages from stdin
+func FollowMaster(agent *tsnet.TSAgent) {
+	log := util.GetLogger()
+	scanner := bufio.NewScanner(os.Stdin)
+
+	for scanner.Scan() {
+		line := scanner.Bytes()
 
 		var msg RecvMessage
-		err = json.Unmarshal(message, &msg)
+		err := json.Unmarshal(line, &msg)
 		if err != nil {
 			log.Error("Unable to unmarshal message: %s", err)
 			log.Debug("Full Error: %v", err)
 			continue
 		}
 
-		log.Debug("Recieved message from master: %v", message)
+		log.Debug("Recieved message from master: %v", line)
 
 		if len(msg.NodeIDs) == 0 {
 			log.Debug("Message recieved had no node IDs")
-			log.Debug("Full message: %s", message)
+			log.Debug("Full message: %s", line)
 			continue
 		}
 
@@ -49,7 +54,7 @@ func (s *Socket) FollowMaster() {
 			wg.Add(1)
 			go func(nodeID string) {
 				defer wg.Done()
-				result, err := s.Agent.GetStatusForPeer(nodeID)
+				result, err := agent.GetStatusForPeer(nodeID)
 				if err != nil {
 					log.Error("Unable to get status for node %s: %s", nodeID, err)
 					return
@@ -70,15 +75,13 @@ func (s *Socket) FollowMaster() {
 
 		// Send the results back to the Headplane master
 		log.Debug("Sending status back to master: %v", results)
-		err = s.SendStatus(results)
-		if err != nil {
-			log.Error("Error sending status: %s", err)
-			return
-		}
+		log.Msg(&SendMessage{
+			Type: "status",
+			Data: results,
+		})
+	}
+
+	if err := scanner.Err(); err != nil {
+		log.Fatal("Error reading from stdin: %s", err)
 	}
 }
-
-// Stops listening for messages from the Headplane master
-func (s *Socket) StopListening() {
-	s.Close()
-}

agent/internal/hpagent/sender.go

@@ -1,11 +0,0 @@
-package hpagent
-
-import (
-	"tailscale.com/tailcfg"
-)
-
-// Sends the status to the Headplane master
-func (s *Socket) SendStatus(status map[string]*tailcfg.HostinfoView) error {
-	err := s.WriteJSON(status)
-	return err
-}

agent/internal/hpagent/websocket.go

@@ -1,67 +0,0 @@
-package hpagent
-
-import (
-	"fmt"
-	"net/http"
-	"net/url"
-
-	"github.com/gorilla/websocket"
-	"github.com/tale/headplane/agent/internal/config"
-	"github.com/tale/headplane/agent/internal/tsnet"
-	"github.com/tale/headplane/agent/internal/util"
-)
-
-type Socket struct {
-	*websocket.Conn
-	Agent *tsnet.TSAgent
-}
-
-// Creates a new websocket connection to the Headplane server.
-func NewSocket(agent *tsnet.TSAgent, cfg *config.Config) (*Socket, error) {
-	log := util.GetLogger()
-
-	wsURL, err := httpToWs(cfg.HPControlURL)
-	if err != nil {
-		return nil, err
-	}
-
-	headers := http.Header{}
-	headers.Add("X-Headplane-Tailnet-ID", agent.ID)
-	auth := fmt.Sprintf("Bearer %s", cfg.HPAuthKey)
-	headers.Add("Authorization", auth)
-
-	log.Info("Dialing WebSocket with master: %s", wsURL)
-	ws, _, err := websocket.DefaultDialer.Dial(wsURL, headers)
-	if err != nil {
-		log.Debug("Failed to dial WebSocket: %s", err)
-		return nil, err
-	}
-
-	return &Socket{ws, agent}, nil
-}
-
-// We need to convert the control URL to a websocket URL
-func httpToWs(controlURL string) (string, error) {
-	log := util.GetLogger()
-	u, err := url.Parse(controlURL)
-	if err != nil {
-		log.Debug("Failed to parse control URL: %s", err)
-		return "", err
-	}
-
-	if u.Scheme == "http" {
-		u.Scheme = "ws"
-	} else if u.Scheme == "https" {
-		u.Scheme = "wss"
-	} else {
-		return "", fmt.Errorf("unsupported scheme: %s", u.Scheme)
-	}
-
-	// We also need to append /_dial to the path
-	if u.Path[len(u.Path)-1] != '/' {
-		u.Path += "/"
-	}
-
-	u.Path += "_dial"
-	return u.String(), nil
-}

agent/internal/tsnet/server.go

@@ -2,6 +2,9 @@ package tsnet
 
 import (
 	"context"
+	"os"
+	"path/filepath"
+
 	"github.com/tale/headplane/agent/internal/config"
 	"github.com/tale/headplane/agent/internal/util"
 	"tailscale.com/client/tailscale"
@@ -17,15 +20,27 @@ type TSAgent struct {
 
 // Creates a new tsnet agent and returns an instance of the server.
 func NewAgent(cfg *config.Config) *TSAgent {
+	log := util.GetLogger()
+
+	dir, err := filepath.Abs(cfg.WorkDir)
+	if err != nil {
+		log.Fatal("Failed to get absolute path: %s", err)
+	}
+
+	if err := os.MkdirAll(dir, 0700); err != nil {
+		log.Fatal("Cannot create agent work directory: %s", err)
+	}
+
 	server := &tsnet.Server{
+		Dir:        dir,
 		Hostname:   cfg.Hostname,
 		ControlURL: cfg.TSControlURL,
 		AuthKey:    cfg.TSAuthKey,
-		Logf:       func(string, ...interface{}) {}, // Disabled by default
+		Logf:       func(string, ...any) {}, // Disabled by default
+		UserLogf:   log.Info,
 	}
 
 	if cfg.Debug {
-		log := util.GetLogger()
 		server.Logf = log.Debug
 	}
 
@@ -47,8 +62,13 @@ func (s *TSAgent) Connect() {
 		log.Fatal("Failed to initialize local Tailscale client: %s", err)
 	}
 
+	id, err := status.Self.PublicKey.MarshalText()
+	if err != nil {
+		log.Fatal("Failed to marshal public key: %s", err)
+	}
+
 	log.Info("Connected to Tailnet (PublicKey: %s)", status.Self.PublicKey)
-	s.ID = string(status.Self.ID)
+	s.ID = string(id)
 }
 
 // Shuts down the tsnet agent.

agent/internal/util/logger.go

@@ -1,66 +1,117 @@
 package util
 
 import (
-	"log"
+	"encoding/json"
+	"fmt"
 	"os"
+	"strings"
 	"sync"
+	"time"
 )
 
-type Logger struct {
-	debug *log.Logger
-	info  *log.Logger
-	error *log.Logger
+type LogLevel string
+
+const (
+	LevelInfo  LogLevel = "info"
+	LevelDebug LogLevel = "debug"
+	LevelError LogLevel = "error"
+	LevelFatal LogLevel = "fatal"
+	LevelMsg   LogLevel = "msg"
+)
+
+type LogMessage struct {
+	Level   LogLevel
+	Time    string
+	Message any
 }
 
-var lock = &sync.Mutex{}
-var logger *Logger
+type Logger struct {
+	debugEnabled bool
+	encoder      *json.Encoder
+	pool         *sync.Pool
+}
+
+var logger = NewLogger()
 
 func GetLogger() *Logger {
-	if logger == nil {
-		lock.Lock()
-		defer lock.Unlock()
-		if logger == nil {
-			logger = NewLogger()
-		}
-	}
-
 	return logger
 }
 
 func NewLogger() *Logger {
-	// Create a new Logger for stdout and stderr
-	// Errors still go to both stdout and stderr
+	enc := json.NewEncoder(os.Stdout)
+	enc.SetEscapeHTML(false)
+
 	return &Logger{
-		debug: nil,
-		info:  log.New(os.Stdout, "[INFO] ", log.LstdFlags),
-		error: log.New(os.Stderr, "[ERROR] ", log.LstdFlags),
+		encoder: enc,
+		pool: &sync.Pool{
+			New: func() any {
+				return &LogMessage{}
+			},
+		},
 	}
 }
 
-func (logger *Logger) SetDebug(debug bool) {
-	if debug {
-		logger.Info("Enabling Debug logging for headplane-agent")
-		logger.Info("Be careful, this will spam a lot of information")
-		logger.debug = log.New(os.Stdout, "[DEBUG] ", log.LstdFlags)
-	} else {
-		logger.debug = nil
+func (l *Logger) SetDebug(enabled bool) {
+	if enabled {
+		l.debugEnabled = true
+		l.Info("Enabling Debug logging for headplane-agent")
+		l.Info("Be careful, this will spam a lot of information")
 	}
 }
 
-func (logger *Logger) Info(fmt string, v ...any) {
-	logger.info.Printf(fmt, v...)
-}
+func (l *Logger) log(level LogLevel, format string, v ...any) {
+	msg := fmt.Sprintf(format, v...)
+	timestamp := time.Now().Format(time.RFC3339)
 
-func (logger *Logger) Debug(fmt string, v ...any) {
-	if logger.debug != nil {
-		logger.debug.Printf(fmt, v...)
+	// Manually construct compact JSON line for performance
+	line := `{"Level":"` + string(level) +
+		`","Time":"` + timestamp +
+		`","Message":"` + escapeString(msg) + `"}` + "\n"
+
+	if level == LevelError || level == LevelFatal {
+		os.Stderr.WriteString(line)
+	}
+
+	// Always write to stdout but also write to stderr for errors
+	os.Stdout.WriteString(line)
+	if level == LevelFatal {
+		os.Exit(1)
 	}
 }
 
-func (logger *Logger) Error(fmt string, v ...any) {
-	logger.error.Printf(fmt, v...)
+func (l *Logger) Debug(format string, v ...any) {
+	if l.debugEnabled {
+		l.log(LevelDebug, format, v...)
+	}
 }
 
-func (logger *Logger) Fatal(fmt string, v ...any) {
-	logger.error.Fatalf(fmt, v...)
+func (l *Logger) Info(format string, v ...any)  { l.log(LevelInfo, format, v...) }
+func (l *Logger) Error(format string, v ...any) { l.log(LevelError, format, v...) }
+func (l *Logger) Fatal(format string, v ...any) { l.log(LevelFatal, format, v...) }
+
+func (l *Logger) Msg(obj any) {
+	entry := l.pool.Get().(*LogMessage)
+	defer l.pool.Put(entry)
+
+	entry.Level = LevelMsg
+	entry.Time = time.Now().Format(time.RFC3339)
+	entry.Message = obj
+
+	// Because the encoder is tied to STDOUT we get a message
+	_ = l.encoder.Encode(entry)
+
+	// Reset the entry for reuse
+	entry.Level = ""
+	entry.Time = ""
+	entry.Message = nil
+}
+
+func escapeString(s string) string {
+	replacer := strings.NewReplacer(
+		`"`, `\"`,
+		`\`, `\\`,
+		"\n", `\n`,
+		"\t", `\t`,
+	)
+	return replacer.Replace(s)
 }