From 377641265ba5c0e400d3ae79c901fc1561f99f3c Mon Sep 17 00:00:00 2001
From: Aarnav Tale <aarnavtale@icloud.com>
Date: Wed, 15 Jan 2025 15:14:37 +0530
Subject: [PATCH] fix: require auth for agent ws

---
 agent/cmd/hp_agent/hp_agent.go | 8 +++++++-
 agent/config/config.go         | 3 +++
 agent/config/preflight.go      | 5 +++++
 agent/hpagent/websocket.go     | 5 ++++-
 4 files changed, 19 insertions(+), 2 deletions(-)

diff --git a/agent/cmd/hp_agent/hp_agent.go b/agent/cmd/hp_agent/hp_agent.go
index f2b46e3..06a6498 100644
--- a/agent/cmd/hp_agent/hp_agent.go
+++ b/agent/cmd/hp_agent/hp_agent.go
@@ -24,7 +24,13 @@ func main() {
 	agent.StartAndFetchID()
 	defer agent.Shutdown()
 
-	ws, err := hpagent.NewSocket(agent, cfg.HPControlURL, cfg.Debug)
+	ws, err := hpagent.NewSocket(
+		agent,
+		cfg.HPControlURL,
+		cfg.HPAuthKey,
+		cfg.Debug,
+	)
+
 	if err != nil {
 		log.Fatalf("Failed to create websocket: %s", err)
 	}
diff --git a/agent/config/config.go b/agent/config/config.go
index f49797e..6ee7d15 100644
--- a/agent/config/config.go
+++ b/agent/config/config.go
@@ -12,6 +12,7 @@ type Config struct {
 	TSControlURL string
 	TSAuthKey    string
 	HPControlURL string
+	HPAuthKey    string
 }
 
 const (
@@ -20,6 +21,7 @@ const (
 	TSControlURLEnv = "HP_AGENT_TS_SERVER"
 	TSAuthKeyEnv    = "HP_AGENT_TS_AUTHKEY"
 	HPControlURLEnv = "HP_AGENT_HP_SERVER"
+	HPAuthKeyEnv	= "HP_AGENT_HP_AUTHKEY"
 )
 
 // Load reads the agent configuration from environment variables.
@@ -30,6 +32,7 @@ func Load() (*Config, error) {
 		TSControlURL: os.Getenv(TSControlURLEnv),
 		TSAuthKey:    os.Getenv(TSAuthKeyEnv),
 		HPControlURL: os.Getenv(HPControlURLEnv),
+		HPAuthKey:    os.Getenv(HPAuthKeyEnv),
 	}
 
 	if os.Getenv(DebugEnv) == "true" {
diff --git a/agent/config/preflight.go b/agent/config/preflight.go
index 7d9c451..a21b9e2 100644
--- a/agent/config/preflight.go
+++ b/agent/config/preflight.go
@@ -24,6 +24,10 @@ func validateRequired(config *Config) error {
 		return fmt.Errorf("%s is required", TSAuthKeyEnv)
 	}
 
+	if config.HPAuthKey == "" {
+		return fmt.Errorf("%s is required", HPAuthKeyEnv)
+	}
+
 	return nil
 }
 
@@ -34,6 +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)
 	resp, err := http.Get(testURL)
 	if err != nil {
diff --git a/agent/hpagent/websocket.go b/agent/hpagent/websocket.go
index 927f99b..06c9056 100644
--- a/agent/hpagent/websocket.go
+++ b/agent/hpagent/websocket.go
@@ -16,7 +16,7 @@ type Socket struct {
 }
 
 // Creates a new websocket connection to the Headplane server.
-func NewSocket(agent *tsnet.TSAgent, controlURL string, debug bool) (*Socket, error) {
+func NewSocket(agent *tsnet.TSAgent, controlURL, authKey string, debug bool) (*Socket, error) {
 	wsURL, err := httpToWs(controlURL)
 	if err != nil {
 		return nil, err
@@ -25,6 +25,9 @@ func NewSocket(agent *tsnet.TSAgent, controlURL string, debug bool) (*Socket, er
 	headers := http.Header{}
 	headers.Add("X-Headplane-TS-Node-ID", agent.ID)
 
+	auth := fmt.Sprintf("Bearer %s", authKey)
+	headers.Add("Authorization", auth)
+
 	log.Printf("dialing websocket at %s", wsURL)
 	ws, _, err := websocket.DefaultDialer.Dial(wsURL, headers)
 	if err != nil {