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() agent.StartAndFetchID()
defer agent.Shutdown() 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 { if err != nil {
log.Fatalf("Failed to create websocket: %s", err) log.Fatalf("Failed to create websocket: %s", err)
} }

View File

@ -12,6 +12,7 @@ type Config struct {
TSControlURL string TSControlURL string
TSAuthKey string TSAuthKey string
HPControlURL string HPControlURL string
HPAuthKey string
} }
const ( const (
@ -20,6 +21,7 @@ const (
TSControlURLEnv = "HP_AGENT_TS_SERVER" TSControlURLEnv = "HP_AGENT_TS_SERVER"
TSAuthKeyEnv = "HP_AGENT_TS_AUTHKEY" TSAuthKeyEnv = "HP_AGENT_TS_AUTHKEY"
HPControlURLEnv = "HP_AGENT_HP_SERVER" HPControlURLEnv = "HP_AGENT_HP_SERVER"
HPAuthKeyEnv = "HP_AGENT_HP_AUTHKEY"
) )
// Load reads the agent configuration from environment variables. // Load reads the agent configuration from environment variables.
@ -30,6 +32,7 @@ func Load() (*Config, error) {
TSControlURL: os.Getenv(TSControlURLEnv), TSControlURL: os.Getenv(TSControlURLEnv),
TSAuthKey: os.Getenv(TSAuthKeyEnv), TSAuthKey: os.Getenv(TSAuthKeyEnv),
HPControlURL: os.Getenv(HPControlURLEnv), HPControlURL: os.Getenv(HPControlURLEnv),
HPAuthKey: os.Getenv(HPAuthKeyEnv),
} }
if os.Getenv(DebugEnv) == "true" { if os.Getenv(DebugEnv) == "true" {

View File

@ -24,6 +24,10 @@ func validateRequired(config *Config) error {
return fmt.Errorf("%s is required", TSAuthKeyEnv) return fmt.Errorf("%s is required", TSAuthKeyEnv)
} }
if config.HPAuthKey == "" {
return fmt.Errorf("%s is required", HPAuthKeyEnv)
}
return nil return nil
} }
@ -34,6 +38,7 @@ func validateTSReady(config *Config) error {
testURL = testURL[:len(testURL)-1] 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/key?v=109", testURL)
resp, err := http.Get(testURL) resp, err := http.Get(testURL)
if err != nil { if err != nil {

View File

@ -16,7 +16,7 @@ type Socket struct {
} }
// Creates a new websocket connection to the Headplane server. // 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) wsURL, err := httpToWs(controlURL)
if err != nil { if err != nil {
return nil, err return nil, err
@ -25,6 +25,9 @@ func NewSocket(agent *tsnet.TSAgent, controlURL string, debug bool) (*Socket, er
headers := http.Header{} headers := http.Header{}
headers.Add("X-Headplane-TS-Node-ID", agent.ID) 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) log.Printf("dialing websocket at %s", wsURL)
ws, _, err := websocket.DefaultDialer.Dial(wsURL, headers) ws, _, err := websocket.DefaultDialer.Dial(wsURL, headers)
if err != nil { if err != nil {