diff --git a/agent/cmd/hp_agent/hp_agent.go b/agent/cmd/hp_agent/hp_agent.go index 06a6498..61ffc08 100644 --- a/agent/cmd/hp_agent/hp_agent.go +++ b/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/tsnet" - "github.com/tale/headplane/agent/hpagent" - "log" + "github.com/tale/headplane/agent/internal/config" + "github.com/tale/headplane/agent/internal/hpagent" + "github.com/tale/headplane/agent/internal/tsnet" + "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( - cfg.Hostname, - cfg.TSControlURL, - cfg.TSAuthKey, - cfg.Debug, - ) + log.SetDebug(cfg.Debug) + agent := tsnet.NewAgent(cfg) - agent.StartAndFetchID() + agent.Connect() defer agent.Shutdown() - ws, err := hpagent.NewSocket( - agent, - cfg.HPControlURL, - cfg.HPAuthKey, - cfg.Debug, - ) - + ws, err := hpagent.NewSocket(agent, cfg) 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() } diff --git a/agent/config/config.go b/agent/internal/config/config.go similarity index 95% rename from agent/config/config.go rename to agent/internal/config/config.go index a0cde3c..64f14a3 100644 --- a/agent/config/config.go +++ b/agent/internal/config/config.go @@ -1,10 +1,6 @@ package config -import ( - "os" - - _ "github.com/joho/godotenv/autoload" -) +import "os" // Config represents the configuration for the agent. type Config struct { diff --git a/agent/config/preflight.go b/agent/internal/config/preflight.go similarity index 93% rename from agent/config/preflight.go rename to agent/internal/config/preflight.go index a21b9e2..3f74958 100644 --- a/agent/config/preflight.go +++ b/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/key?v=109", testURL) + testURL = fmt.Sprintf("%s/health", testURL) resp, err := http.Get(testURL) if err != nil { return fmt.Errorf("Failed to connect to TS control server: %s", err) diff --git a/agent/hpagent/handler.go b/agent/internal/hpagent/handler.go similarity index 64% rename from agent/hpagent/handler.go rename to agent/internal/hpagent/handler.go index 3244e48..8523bf1 100644 --- a/agent/hpagent/handler.go +++ b/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.Printf("got message: %s", message) - } + log.Debug("Recieved message from master: %v", 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) - } } } diff --git a/agent/hpagent/sender.go b/agent/internal/hpagent/sender.go similarity index 100% rename from agent/hpagent/sender.go rename to agent/internal/hpagent/sender.go diff --git a/agent/hpagent/websocket.go b/agent/internal/hpagent/websocket.go similarity index 63% rename from agent/hpagent/websocket.go rename to agent/internal/hpagent/websocket.go index 1434441..faffe71 100644 --- a/agent/hpagent/websocket.go +++ b/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) { - wsURL, err := httpToWs(controlURL) +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", authKey) + auth := fmt.Sprintf("Bearer %s", cfg.HPAuthKey) 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 } diff --git a/agent/tsnet/peers.go b/agent/internal/tsnet/peers.go similarity index 52% rename from agent/tsnet/peers.go rename to agent/internal/tsnet/peers.go index 5011547..9628dca 100644 --- a/agent/tsnet/peers.go +++ b/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.Printf("querying peer state for %s", id) - } - + log.Debug("Querying status of peer: %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 } diff --git a/agent/tsnet/server.go b/agent/internal/tsnet/server.go similarity index 53% rename from agent/tsnet/server.go rename to agent/internal/tsnet/server.go index cf90e8b..f41c73a 100644 --- a/agent/tsnet/server.go +++ b/agent/internal/tsnet/server.go @@ -2,10 +2,8 @@ package tsnet import ( "context" - "fmt" - "log" - "os" - + "github.com/tale/headplane/agent/internal/config" + "github.com/tale/headplane/agent/internal/util" "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 { - s := &tsnet.Server{ - Hostname: hostname, - ControlURL: controlURL, - AuthKey: authKey, +func NewAgent(cfg *config.Config) *TSAgent { + server := &tsnet.Server{ + Hostname: cfg.Hostname, + ControlURL: cfg.TSControlURL, + AuthKey: cfg.TSAuthKey, Logf: func(string, ...interface{}) {}, // Disabled by default } - if debug { - s.Logf = log.New( - os.Stderr, - fmt.Sprintf("[DBG:%s] ", hostname), - log.LstdFlags, - ).Printf + if cfg.Debug { + log := util.GetLogger() + server.Logf = log.Debug } - 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) } diff --git a/agent/internal/util/logger.go b/agent/internal/util/logger.go new file mode 100644 index 0000000..070e061 --- /dev/null +++ b/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...) +}