feat: use a debug logger for agent

agent/cmd/hp_agent/hp_agent.go

@@ -2,39 +2,30 @@ package main
 
 import (
 	_ "github.com/joho/godotenv/autoload"
-	"github.com/tale/headplane/agent/config"
+	"github.com/tale/headplane/agent/internal/config"
-	"github.com/tale/headplane/agent/tsnet"
+	"github.com/tale/headplane/agent/internal/hpagent"
-	"github.com/tale/headplane/agent/hpagent"
+	"github.com/tale/headplane/agent/internal/tsnet"
-	"log"
+	"github.com/tale/headplane/agent/internal/util"
 )
 
 func main() {
+	log := util.GetLogger()
 	cfg, err := config.Load()
 	if err != nil {
-		log.Fatalf("Failed to load configuration: %s", err)
+		log.Fatal("Failed to load config: %s", err)
 	}
 
-	agent := tsnet.NewAgent(
+	log.SetDebug(cfg.Debug)
-		cfg.Hostname,
+	agent := tsnet.NewAgent(cfg)
-		cfg.TSControlURL,
-		cfg.TSAuthKey,
-		cfg.Debug,
-	)
 
-	agent.StartAndFetchID()
+	agent.Connect()
 	defer agent.Shutdown()
 
-	ws, err := hpagent.NewSocket(
+	ws, err := hpagent.NewSocket(agent, cfg)
-		agent,
-		cfg.HPControlURL,
-		cfg.HPAuthKey,
-		cfg.Debug,
-	)
-
 	if err != nil {
-		log.Fatalf("Failed to create websocket: %s", err)
+		log.Fatal("Failed to create websocket: %s", err)
 	}
 
 	defer ws.StopListening()
-	ws.StartListening()
+	ws.FollowMaster()
 }

agent/internal/config/config.go

@@ -1,10 +1,6 @@
 package config
 
-import (
+import "os"
-	"os"
-
-	_ "github.com/joho/godotenv/autoload"
-)
 
 // Config represents the configuration for the agent.
 type Config struct {

agent/internal/config/preflight.go

@@ -38,8 +38,7 @@ func validateTSReady(config *Config) error {
 		testURL = testURL[:len(testURL)-1]
 	}
 
-	// TODO: Consequences of switching to /health (headscale only)
+	testURL = fmt.Sprintf("%s/health", testURL)
-	testURL = fmt.Sprintf("%s/key?v=109", testURL)
 	resp, err := http.Get(testURL)
 	if err != nil {
 		return fmt.Errorf("Failed to connect to TS control server: %s", err)

agent/internal/hpagent/handler.go

@@ -2,38 +2,41 @@ package hpagent
 
 import (
 	"encoding/json"
-	"log"
 	"sync"
+
+	"github.com/tale/headplane/agent/internal/util"
 	"tailscale.com/tailcfg"
 )
 
 // Represents messages from the Headplane master
 type RecvMessage struct {
-	NodeIDs []string `json:omitempty`
+	NodeIDs []string
 }
 
 // Starts listening for messages from the Headplane master
-func (s *Socket) StartListening() {
+func (s *Socket) FollowMaster() {
+	log := util.GetLogger()
+
 	for {
 		_, message, err := s.ReadMessage()
 		if err != nil {
-			log.Printf("error reading message: %v", err)
+			log.Error("Error reading message: %s", err)
 			return
 		}
 
 		var msg RecvMessage
 		err = json.Unmarshal(message, &msg)
 		if err != nil {
-			log.Printf("error unmarshalling message: %v", err)
+			log.Error("Unable to unmarshal message: %s", err)
+			log.Debug("Full Error: %v", err)
 			continue
 		}
 
-		if s.Debug {
+		log.Debug("Recieved message from master: %v", message)
-			log.Printf("got message: %s", message)
-		}
 
 		if len(msg.NodeIDs) == 0 {
-			log.Printf("got a message with no node IDs? %s", message)
+			log.Debug("Message recieved had no node IDs")
+			log.Debug("Full message: %s", message)
 			continue
 		}
 
@@ -48,11 +51,12 @@ func (s *Socket) StartListening() {
 				defer wg.Done()
 				result, err := s.Agent.GetStatusForPeer(nodeID)
 				if err != nil {
-					log.Printf("error getting status: %v", err)
+					log.Error("Unable to get status for node %s: %s", nodeID, err)
 					return
 				}
 
 				if result == nil {
+					log.Debug("No status for node %s", nodeID)
 					return
 				}
 
@@ -65,15 +69,12 @@ func (s *Socket) StartListening() {
 		wg.Wait()
 
 		// 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.Printf("error sending status: %v", err)
+			log.Error("Error sending status: %s", err)
 			return
 		}
-
-		if s.Debug {
-			log.Printf("sent status: %s", results)
-		}
 	}
 }
 

agent/internal/hpagent/websocket.go

@@ -2,46 +2,50 @@ package hpagent
 
 import (
 	"fmt"
-	"log"
 	"net/http"
 	"net/url"
 
 	"github.com/gorilla/websocket"
-	"github.com/tale/headplane/agent/tsnet"
+	"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
-	Debug bool
 	Agent *tsnet.TSAgent
 }
 
 // Creates a new websocket connection to the Headplane server.
-func NewSocket(agent *tsnet.TSAgent, controlURL, authKey string, debug bool) (*Socket, error) {
+func NewSocket(agent *tsnet.TSAgent, cfg *config.Config) (*Socket, error) {
-	wsURL, err := httpToWs(controlURL)
+	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)
-	auth := fmt.Sprintf("Bearer %s", authKey)
 	headers.Add("Authorization", auth)
 
-	log.Printf("dialing websocket at %s", wsURL)
+	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, debug, agent}, nil
+	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
 	}
 

agent/internal/tsnet/peers.go

@@ -2,9 +2,11 @@ package tsnet
 
 import (
 	"context"
+	"encoding/hex"
 	"fmt"
-	"log"
 	"strings"
+
+	"github.com/tale/headplane/agent/internal/util"
 	"tailscale.com/tailcfg"
 	"tailscale.com/types/key"
 
@@ -13,26 +15,41 @@ import (
 
 // Returns the raw hostinfo for a peer based on node ID.
 func (s *TSAgent) GetStatusForPeer(id string) (*tailcfg.HostinfoView, error) {
+	log := util.GetLogger()
+
 	if !strings.HasPrefix(id, "nodekey:") {
+		log.Debug("Node ID with missing prefix: %s", id)
 		return nil, fmt.Errorf("invalid node ID: %s", id)
 	}
 
-	if s.Debug {
+	log.Debug("Querying status of peer: %s", id)
-		log.Printf("querying peer state for %s", id)
-	}
-
 	status, err := s.Lc.Status(context.Background())
 	if err != nil {
+		log.Debug("Failed to get status: %s", err)
 		return nil, fmt.Errorf("failed to get status: %w", err)
 	}
 
-	nodeKey, err := key.ParseNodePublicUntyped(mem.S(id[8:]))
+	// We need to convert from 64 char hex to 32 byte raw.
+	bytes, err := hex.DecodeString(id[8:])
+	if err != nil {
+		log.Debug("Failed to decode hex: %s", err)
+		return nil, fmt.Errorf("failed to decode hex: %w", err)
+	}
+
+	raw := mem.B(bytes)
+	if raw.Len() != 32 {
+		log.Debug("Invalid node ID length: %d", raw.Len())
+		return nil, fmt.Errorf("invalid node ID length: %d", raw.Len())
+	}
+
+	nodeKey := key.NodePublicFromRaw32(raw)
 	peer := status.Peer[nodeKey]
 	if peer == nil {
 		// Check if we are on Self.
 		if status.Self.PublicKey == nodeKey {
 			peer = status.Self
 		} else {
+			log.Debug("Peer not found in status: %s", id)
 			return nil, nil
 		}
 	}
@@ -40,8 +57,10 @@ func (s *TSAgent) GetStatusForPeer(id string) (*tailcfg.HostinfoView, error) {
 	ip := peer.TailscaleIPs[0].String()
 	whois, err := s.Lc.WhoIs(context.Background(), ip)
 	if err != nil {
+		log.Debug("Failed to get whois: %s", err)
 		return nil, fmt.Errorf("failed to get whois: %w", err)
 	}
 
+	log.Debug("Got whois for peer %s: %v", id, whois)
 	return &whois.Node.Hostinfo, nil
 }

agent/internal/tsnet/server.go

@@ -2,10 +2,8 @@ package tsnet
 
 import (
 	"context"
-	"fmt"
+	"github.com/tale/headplane/agent/internal/config"
-	"log"
+	"github.com/tale/headplane/agent/internal/util"
-	"os"
-
 	"tailscale.com/client/tailscale"
 	"tailscale.com/tsnet"
 )
@@ -15,43 +13,41 @@ type TSAgent struct {
 	*tsnet.Server
 	Lc *tailscale.LocalClient
 	ID string
-	Debug bool
 }
 
 // Creates a new tsnet agent and returns an instance of the server.
-func NewAgent(hostname, controlURL, authKey string, debug bool) *TSAgent {
+func NewAgent(cfg *config.Config) *TSAgent {
-	s := &tsnet.Server{
+	server := &tsnet.Server{
-		Hostname:   hostname,
+		Hostname:   cfg.Hostname,
-		ControlURL: controlURL,
+		ControlURL: cfg.TSControlURL,
-		AuthKey:    authKey,
+		AuthKey:    cfg.TSAuthKey,
 		Logf:       func(string, ...interface{}) {}, // Disabled by default
 	}
 
-	if debug {
+	if cfg.Debug {
-		s.Logf = log.New(
+		log := util.GetLogger()
-			os.Stderr,
+		server.Logf = log.Debug
-			fmt.Sprintf("[DBG:%s] ", hostname),
-			log.LstdFlags,
-		).Printf
 	}
 
-	return &TSAgent{s, nil, "", debug}
+	return &TSAgent{server, nil, ""}
 }
 
 // Starts the tsnet agent and sets the node ID.
-func (s *TSAgent) StartAndFetchID() {
+func (s *TSAgent) Connect() {
+	log := util.GetLogger()
+
 	// Waits until the agent is up and running.
 	status, err := s.Up(context.Background())
 	if err != nil {
-		log.Fatalf("Failed to start agent: %v", err)
+		log.Fatal("Failed to connect to Tailnet: %s", err)
 	}
 
 	s.Lc, err = s.LocalClient()
 	if err != nil {
-		log.Fatalf("Failed to create local client: %v", err)
+		log.Fatal("Failed to initialize local Tailscale client: %s", err)
 	}
 
-	log.Printf("Agent running with ID: %s", status.Self.PublicKey)
+	log.Info("Connected to Tailnet (PublicKey: %s)", status.Self.PublicKey)
 	s.ID = string(status.Self.ID)
 }
 

agent/internal/util/logger.go

@@ -0,0 +1,66 @@
+package util
+
+import (
+	"log"
+	"os"
+	"sync"
+)
+
+type Logger struct {
+	debug *log.Logger
+	info  *log.Logger
+	error *log.Logger
+}
+
+var lock = &sync.Mutex{}
+var logger *Logger
+
+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
+	return &Logger{
+		debug: nil,
+		info:  log.New(os.Stdout, "[INFO] ", log.LstdFlags),
+		error: log.New(os.Stderr, "[ERROR] ", log.LstdFlags),
+	}
+}
+
+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 (logger *Logger) Info(fmt string, v ...any) {
+	logger.info.Printf(fmt, v...)
+}
+
+func (logger *Logger) Debug(fmt string, v ...any) {
+	if logger.debug != nil {
+		logger.debug.Printf(fmt, v...)
+	}
+}
+
+func (logger *Logger) Error(fmt string, v ...any) {
+	logger.error.Printf(fmt, v...)
+}
+
+func (logger *Logger) Fatal(fmt string, v ...any) {
+	logger.error.Fatalf(fmt, v...)
+}