From 31692c03e8dd9c70189d174084067d5214bf1143 Mon Sep 17 00:00:00 2001 From: Owen Schwartz Date: Sat, 7 Dec 2024 22:07:02 -0500 Subject: [PATCH] Fix various things --- main.go | 15 ++++---- websocket/client.go | 86 +++++++++++++++++++++++++++++++++++---------- websocket/config.go | 19 +++++++--- websocket/types.go | 7 ++-- 4 files changed, 96 insertions(+), 31 deletions(-) diff --git a/main.go b/main.go index bc22d6f..b65acaf 100644 --- a/main.go +++ b/main.go @@ -185,7 +185,7 @@ func main() { logLevel string ) - flag.StringVar(&endpoint, "endpoint", "http://localhost:3000/api/v1", "Endpoint of your pangolin server") + flag.StringVar(&endpoint, "endpoint", "", "Endpoint of your pangolin server") flag.StringVar(&id, "id", "", "Newt ID") flag.StringVar(&secret, "secret", "", "Newt secret") flag.StringVar(&dns, "dns", "8.8.8.8", "DNS server to use") @@ -204,10 +204,9 @@ func main() { // Create a new client client, err := websocket.NewClient( - // the id and secret from the params - id, - secret, - websocket.WithBaseURL(endpoint), // TODO: save the endpoint in the config file so we dont have to pass it in every time + id, // CLI arg takes precedence + secret, // CLI arg takes precedence + endpoint, ) if err != nil { logger.Fatal("Failed to create client: %v", err) @@ -223,6 +222,8 @@ func main() { // Register handlers for different message types client.RegisterHandler("newt/wg/connect", func(msg websocket.WSMessage) { + logger.Info("Received registration message") + if connected { logger.Info("Already connected! Put I will send a ping anyway...") ping(tnet, wgData.ServerIP) @@ -385,7 +386,7 @@ persistent_keepalive_interval=5`, fixKey(fmt.Sprintf("%s", privateKey)), fixKey( defer client.Close() publicKey := privateKey.PublicKey() - logger.Info("Public key: %s", publicKey) + logger.Debug("Public key: %s", publicKey) // TODO: how to retry? err = client.SendMessage("newt/wg/register", map[string]interface{}{ "publicKey": fmt.Sprintf("%s", publicKey), @@ -394,6 +395,8 @@ persistent_keepalive_interval=5`, fixKey(fmt.Sprintf("%s", privateKey)), fixKey( logger.Info("Failed to send message: %v", err) } + logger.Info("Sent registration message") + // Wait for interrupt signal sigCh := make(chan os.Signal, 1) signal.Notify(sigCh, syscall.SIGINT, syscall.SIGTERM) diff --git a/websocket/client.go b/websocket/client.go index 4c291df..b473f16 100644 --- a/websocket/client.go +++ b/websocket/client.go @@ -6,6 +6,8 @@ import ( "fmt" "net/http" "net/url" + "newt/logger" + "strings" "sync" "github.com/gorilla/websocket" @@ -32,20 +34,21 @@ func WithBaseURL(url string) ClientOption { } // NewClient creates a new Newt client -func NewClient(newtID, secret string, opts ...ClientOption) (*Client, error) { +func NewClient(newtID, secret string, endpoint string, opts ...ClientOption) (*Client, error) { config := &Config{ - NewtID: newtID, - Secret: secret, + NewtID: newtID, + Secret: secret, + Endpoint: endpoint, } client := &Client{ config: config, - baseURL: "https://fossorial.io", // default value + baseURL: endpoint, // default value handlers: make(map[string]MessageHandler), done: make(chan struct{}), } - // Apply options + // Apply options before loading config for _, opt := range opts { opt(client) } @@ -66,31 +69,48 @@ func (c *Client) Connect() error { return fmt.Errorf("failed to get token: %w", err) } + logger.Info("Using token: %s", token) + // Update config with new token and save c.config.Token = token if err := c.saveConfig(); err != nil { return fmt.Errorf("failed to save config: %w", err) } - // Connect to WebSocket - wsURL := fmt.Sprintf("wss://%s/ws", "fossorial.io") // Remove http:// prefix + // Parse the base URL to determine protocol and hostname + baseURL, err := url.Parse(c.baseURL) + if err != nil { + return fmt.Errorf("failed to parse base URL: %w", err) + } + + // Determine WebSocket protocol based on HTTP protocol + wsProtocol := "wss" + if baseURL.Scheme == "http" { + wsProtocol = "ws" + } + + // Create WebSocket URL using the hostname without path + wsURL := fmt.Sprintf("%s://%s/ws", wsProtocol, baseURL.Host) u, err := url.Parse(wsURL) if err != nil { return fmt.Errorf("failed to parse WebSocket URL: %w", err) } + // Add token to query parameters q := u.Query() q.Set("token", token) u.RawQuery = q.Encode() + // Connect to WebSocket conn, _, err := websocket.DefaultDialer.Dial(u.String(), nil) if err != nil { return fmt.Errorf("failed to connect to WebSocket: %w", err) } + logger.Info("Connected to WebSocket") + c.conn = conn go c.readPump() - return nil } @@ -149,6 +169,15 @@ func (c *Client) readPump() { } func (c *Client) getToken() (string, error) { + // Parse the base URL to ensure we have the correct hostname + baseURL, err := url.Parse(c.baseURL) + if err != nil { + return "", fmt.Errorf("failed to parse base URL: %w", err) + } + + // Ensure we have the base URL without trailing slashes + baseEndpoint := strings.TrimRight(baseURL.String(), "/") + // If we already have a token, try to use it if c.config.Token != "" { tokenCheckData := map[string]interface{}{ @@ -156,19 +185,28 @@ func (c *Client) getToken() (string, error) { "secret": c.config.Secret, "token": c.config.Token, } - jsonData, _ := json.Marshal(tokenCheckData) - - resp, err := http.Post(c.baseURL+"/auth/newt/get-token", "application/json", bytes.NewBuffer(jsonData)) + jsonData, err := json.Marshal(tokenCheckData) if err != nil { - return "", err + return "", fmt.Errorf("failed to marshal token check data: %w", err) + } + + // Make request to validate existing token + resp, err := http.Post( + baseEndpoint+"/api/v1/auth/newt/get-token", + "application/json", + bytes.NewBuffer(jsonData), + ) + if err != nil { + return "", fmt.Errorf("failed to check token validity: %w", err) } defer resp.Body.Close() var tokenResp TokenResponse if err := json.NewDecoder(resp.Body).Decode(&tokenResp); err != nil { - return "", err + return "", fmt.Errorf("failed to decode token check response: %w", err) } + // If token is still valid, return it if tokenResp.Success && tokenResp.Message == "Token session already valid" { return c.config.Token, nil } @@ -179,22 +217,34 @@ func (c *Client) getToken() (string, error) { "newtId": c.config.NewtID, "secret": c.config.Secret, } - jsonData, _ := json.Marshal(tokenData) - - resp, err := http.Post(c.baseURL+"/auth/newt/get-token", "application/json", bytes.NewBuffer(jsonData)) + jsonData, err := json.Marshal(tokenData) if err != nil { - return "", err + return "", fmt.Errorf("failed to marshal token request data: %w", err) + } + + // Make request to get new token + resp, err := http.Post( + baseEndpoint+"/api/v1/auth/newt/get-token", + "application/json", + bytes.NewBuffer(jsonData), + ) + if err != nil { + return "", fmt.Errorf("failed to request new token: %w", err) } defer resp.Body.Close() var tokenResp TokenResponse if err := json.NewDecoder(resp.Body).Decode(&tokenResp); err != nil { - return "", err + return "", fmt.Errorf("failed to decode token response: %w", err) } if !tokenResp.Success { return "", fmt.Errorf("failed to get token: %s", tokenResp.Message) } + if tokenResp.Data.Token == "" { + return "", fmt.Errorf("received empty token from server") + } + return tokenResp.Data.Token, nil } diff --git a/websocket/config.go b/websocket/config.go index ae5efa5..794ff1e 100644 --- a/websocket/config.go +++ b/websocket/config.go @@ -27,7 +27,7 @@ func getConfigPath() string { } func (c *Client) loadConfig() error { - if c.config.NewtID != "" && c.config.Secret != "" { + if c.config.NewtID != "" && c.config.Secret != "" && c.config.Endpoint != "" { return nil } @@ -45,9 +45,20 @@ func (c *Client) loadConfig() error { return err } - c.config.Token = config.Token // I think it always needs to get a new token - c.config.NewtID = config.NewtID - c.config.Secret = config.Secret + if c.config.NewtID == "" { + c.config.NewtID = config.NewtID + } + if c.config.Token == "" { + c.config.Token = config.Token + } + if c.config.Secret == "" { + c.config.Secret = config.Secret + } + if c.config.Endpoint == "" { + c.config.Endpoint = config.Endpoint + c.baseURL = config.Endpoint + } + return nil } diff --git a/websocket/types.go b/websocket/types.go index 623a0a0..084465a 100644 --- a/websocket/types.go +++ b/websocket/types.go @@ -1,9 +1,10 @@ package websocket type Config struct { - NewtID string `json:"newtId"` - Secret string `json:"secret"` - Token string `json:"token"` + NewtID string `json:"newtId"` + Secret string `json:"secret"` + Token string `json:"token"` + Endpoint string `json:"endpoint"` } type TokenResponse struct {