feat: create initial agent for localapi

This commit is contained in:
Aarnav Tale
2024-12-30 13:48:49 +05:30
parent 6156d78907
commit 719d7a5e76
10 changed files with 675 additions and 0 deletions
+34
View File
@@ -0,0 +1,34 @@
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"
)
func main() {
cfg, err := config.Load()
if err != nil {
log.Fatalf("Failed to load configuration: %s", err)
}
agent := tsnet.NewAgent(
cfg.Hostname,
cfg.TSControlURL,
cfg.TSAuthKey,
cfg.Debug,
)
agent.StartAndFetchID()
defer agent.Shutdown()
ws, err := hpagent.NewSocket(agent, cfg.HPControlURL, cfg.Debug)
if err != nil {
log.Fatalf("Failed to create websocket: %s", err)
}
defer ws.StopListening()
ws.StartListening()
}
+52
View File
@@ -0,0 +1,52 @@
package config
import (
_ "github.com/joho/godotenv/autoload"
"os"
)
// Config represents the configuration for the agent.
type Config struct {
Debug bool
Hostname string
TSControlURL string
TSAuthKey string
HPControlURL string
}
const (
DebugEnv = "HP_AGENT_DEBUG"
HostnameEnv = "HP_AGENT_HOSTNAME"
TSControlURLEnv = "HP_AGENT_TS_SERVER"
TSAuthKeyEnv = "HP_AGENT_TS_AUTHKEY"
HPControlURLEnv = "HP_AGENT_HP_SERVER"
)
// 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),
}
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
}
+68
View File
@@ -0,0 +1,68 @@
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)
}
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/key?v=109", 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
}
+83
View File
@@ -0,0 +1,83 @@
package hpagent
import (
"encoding/json"
"log"
"sync"
"tailscale.com/tailcfg"
)
// Represents messages from the Headplane master
type RecvMessage struct {
NodeIDs []string `json:omitempty`
}
// Starts listening for messages from the Headplane master
func (s *Socket) StartListening() {
for {
_, message, err := s.ReadMessage()
if err != nil {
log.Printf("error reading message: %v", err)
return
}
var msg RecvMessage
err = json.Unmarshal(message, &msg)
if err != nil {
log.Printf("error unmarshalling message: %v", err)
continue
}
if s.Debug {
log.Printf("got message: %s", message)
}
if len(msg.NodeIDs) == 0 {
log.Printf("got a message with no node IDs? %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.Printf("error getting status: %v", err)
return
}
if result == nil {
return
}
mu.Lock()
results[nodeID] = result
mu.Unlock()
}(nodeID)
}
wg.Wait()
// Send the results back to the Headplane master
err = s.SendStatus(results)
if err != nil {
log.Printf("error sending status: %v", err)
return
}
if s.Debug {
log.Printf("sent status: %s", results)
}
}
}
// 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
}
+53
View File
@@ -0,0 +1,53 @@
package hpagent
import (
"fmt"
"github.com/gorilla/websocket"
"github.com/tale/headplane/agent/tsnet"
"log"
"net/http"
"net/url"
)
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 string, debug bool) (*Socket, error) {
wsURL, err := httpToWs(controlURL)
if err != nil {
return nil, err
}
headers := http.Header{}
headers.Add("X-Headplane-TS-Node-ID", agent.ID)
log.Printf("dialing websocket at %s", wsURL)
ws, _, err := websocket.DefaultDialer.Dial(wsURL, headers)
if err != nil {
return nil, err
}
return &Socket{ws, debug, agent}, nil
}
// We need to convert the control URL to a websocket URL
func httpToWs(controlURL string) (string, error) {
u, err := url.Parse(controlURL)
if err != nil {
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)
}
return u.String(), nil
}
+42
View File
@@ -0,0 +1,42 @@
package tsnet
import (
"context"
"fmt"
"log"
"strings"
"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) {
if !strings.HasPrefix(id, "nodekey:") {
return nil, fmt.Errorf("invalid node ID")
}
if s.Debug {
log.Printf("querying peer state for %s", id)
}
status, err := s.Lc.Status(context.Background())
if err != nil {
return nil, fmt.Errorf("failed to get status: %w", err)
}
nodeKey, err := key.ParseNodePublicUntyped(mem.S(id[8:]))
peer := status.Peer[nodeKey]
if peer == nil {
return nil, nil
}
ip := peer.TailscaleIPs[0].String()
whois, err := s.Lc.WhoIs(context.Background(), ip)
if err != nil {
return nil, fmt.Errorf("failed to get whois: %w", err)
}
return &whois.Node.Hostinfo, nil
}
+61
View File
@@ -0,0 +1,61 @@
package tsnet
import (
"context"
"fmt"
"log"
"os"
"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
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,
Logf: func(string, ...interface{}) {}, // Disabled by default
}
if debug {
s.Logf = log.New(
os.Stderr,
fmt.Sprintf("[DBG:%s] ", hostname),
log.LstdFlags,
).Printf
}
return &TSAgent{s, nil, "", debug}
}
// Starts the tsnet agent and sets the node ID.
func (s *TSAgent) StartAndFetchID() {
// 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)
}
s.Lc, err = s.LocalClient()
if err != nil {
log.Fatalf("Failed to create local client: %v", err)
}
log.Printf("Agent running with ID: %s", status.Self.PublicKey)
s.ID = string(status.Self.ID)
}
// Shuts down the tsnet agent.
func (s *TSAgent) Shutdown() {
s.Close()
}