feat: use a debug logger for agent

This commit is contained in:
Aarnav Tale
2025-04-08 14:51:28 -04:00
parent 3dca401837
commit 608fcacdb1
9 changed files with 151 additions and 79 deletions
+52
View File
@@ -0,0 +1,52 @@
package config
import "os"
// Config represents the configuration for the agent.
type Config struct {
Debug bool
Hostname string
TSControlURL string
TSAuthKey string
HPControlURL string
HPAuthKey string
}
const (
DebugEnv = "HEADPLANE_AGENT_DEBUG"
HostnameEnv = "HEADPLANE_AGENT_HOSTNAME"
TSControlURLEnv = "HEADPLANE_AGENT_TS_SERVER"
TSAuthKeyEnv = "HEADPLANE_AGENT_TS_AUTHKEY"
HPControlURLEnv = "HEADPLANE_AGENT_HP_SERVER"
HPAuthKeyEnv = "HEADPLANE_AGENT_HP_AUTHKEY"
)
// Load reads the agent configuration from environment variables.
func Load() (*Config, error) {
c := &Config{
Debug: false,
Hostname: os.Getenv(HostnameEnv),
TSControlURL: os.Getenv(TSControlURLEnv),
TSAuthKey: os.Getenv(TSAuthKeyEnv),
HPControlURL: os.Getenv(HPControlURLEnv),
HPAuthKey: os.Getenv(HPAuthKeyEnv),
}
if os.Getenv(DebugEnv) == "true" {
c.Debug = true
}
if err := validateRequired(c); err != nil {
return nil, err
}
if err := validateTSReady(c); err != nil {
return nil, err
}
if err := validateHPReady(c); err != nil {
return nil, err
}
return c, nil
}
+72
View File
@@ -0,0 +1,72 @@
package config
import (
"fmt"
"net/http"
"strings"
)
// Checks to make sure all required environment variables are set
func validateRequired(config *Config) error {
if config.Hostname == "" {
return fmt.Errorf("%s is required", HostnameEnv)
}
if config.TSControlURL == "" {
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)
}
return nil
}
// Pings the Tailscale control server to make sure it's up and running
func validateTSReady(config *Config) error {
testURL := config.TSControlURL
if strings.HasSuffix(testURL, "/") {
testURL = testURL[:len(testURL)-1]
}
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)
}
if resp.StatusCode != 200 {
return fmt.Errorf("Failed to connect to TS control server: %s", resp.Status)
}
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
}
+84
View File
@@ -0,0 +1,84 @@
package hpagent
import (
"encoding/json"
"sync"
"github.com/tale/headplane/agent/internal/util"
"tailscale.com/tailcfg"
)
// Represents messages from the Headplane master
type RecvMessage struct {
NodeIDs []string
}
// Starts listening for messages from the Headplane master
func (s *Socket) FollowMaster() {
log := util.GetLogger()
for {
_, message, err := s.ReadMessage()
if err != nil {
log.Error("Error reading message: %s", err)
return
}
var msg RecvMessage
err = json.Unmarshal(message, &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)
if len(msg.NodeIDs) == 0 {
log.Debug("Message recieved had no node IDs")
log.Debug("Full message: %s", message)
continue
}
// Accumulate the results since we invoke via gofunc
results := make(map[string]*tailcfg.HostinfoView)
mu := sync.Mutex{}
wg := sync.WaitGroup{}
for _, nodeID := range msg.NodeIDs {
wg.Add(1)
go func(nodeID string) {
defer wg.Done()
result, err := s.Agent.GetStatusForPeer(nodeID)
if err != nil {
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
}
mu.Lock()
results[nodeID] = result
mu.Unlock()
}(nodeID)
}
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.Error("Error sending status: %s", err)
return
}
}
}
// Stops listening for messages from the Headplane master
func (s *Socket) StopListening() {
s.Close()
}
+11
View File
@@ -0,0 +1,11 @@
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
}
+67
View File
@@ -0,0 +1,67 @@
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
}
+66
View File
@@ -0,0 +1,66 @@
package tsnet
import (
"context"
"encoding/hex"
"fmt"
"strings"
"github.com/tale/headplane/agent/internal/util"
"tailscale.com/tailcfg"
"tailscale.com/types/key"
"go4.org/mem"
)
// 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)
}
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)
}
// 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
}
}
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
}
+57
View File
@@ -0,0 +1,57 @@
package tsnet
import (
"context"
"github.com/tale/headplane/agent/internal/config"
"github.com/tale/headplane/agent/internal/util"
"tailscale.com/client/tailscale"
"tailscale.com/tsnet"
)
// Wrapper type so we can add methods to the server.
type TSAgent struct {
*tsnet.Server
Lc *tailscale.LocalClient
ID string
}
// Creates a new tsnet agent and returns an instance of the server.
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 cfg.Debug {
log := util.GetLogger()
server.Logf = log.Debug
}
return &TSAgent{server, nil, ""}
}
// Starts the tsnet agent and sets the node ID.
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.Fatal("Failed to connect to Tailnet: %s", err)
}
s.Lc, err = s.LocalClient()
if err != nil {
log.Fatal("Failed to initialize local Tailscale client: %s", err)
}
log.Info("Connected to Tailnet (PublicKey: %s)", status.Self.PublicKey)
s.ID = string(status.Self.ID)
}
// Shuts down the tsnet agent.
func (s *TSAgent) Shutdown() {
s.Close()
}
+66
View File
@@ -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...)
}