fix: require auth for agent ws

This commit is contained in:
Aarnav Tale 2025-01-15 15:14:37 +05:30
parent 5bc313e5fe
commit 377641265b
No known key found for this signature in database
4 changed files with 19 additions and 2 deletions

View File

@ -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)
}

View File

@ -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" {

View File

@ -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 {

View File

@ -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 {