WIP: Add websocket support

This commit is contained in:
Owen Schwartz
2024-11-10 22:04:36 -05:00
parent b348c768c7
commit 2e5531b4a5
6 changed files with 515 additions and 151 deletions

202
websocket/client.go Normal file
View File

@@ -0,0 +1,202 @@
package websocket
import (
"bytes"
"encoding/json"
"fmt"
"log"
"net/http"
"net/url"
"sync"
"github.com/gorilla/websocket"
)
type Client struct {
conn *websocket.Conn
config *Config
baseURL string
handlers map[string]MessageHandler
done chan struct{}
handlersMux sync.RWMutex
}
type ClientOption func(*Client)
type MessageHandler func(message WSMessage)
// WithBaseURL sets the base URL for the client
func WithBaseURL(url string) ClientOption {
return func(c *Client) {
c.baseURL = url
}
}
// NewClient creates a new Newt client
func NewClient(newtID, secret string, opts ...ClientOption) (*Client, error) {
config := &Config{
NewtID: newtID,
Secret: secret,
}
client := &Client{
config: config,
baseURL: "http://localhost:3000", // default value
handlers: make(map[string]MessageHandler),
done: make(chan struct{}),
}
// Apply options
for _, opt := range opts {
opt(client)
}
// Load existing config if available
if err := client.loadConfig(); err != nil {
return nil, fmt.Errorf("failed to load config: %w", err)
}
return client, nil
}
// Connect establishes the WebSocket connection
func (c *Client) Connect() error {
// Get token for authentication
token, err := c.getToken()
if err != nil {
return fmt.Errorf("failed to get token: %w", err)
}
// 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("ws://%s/ws", c.baseURL[7:]) // Remove http:// prefix
u, err := url.Parse(wsURL)
if err != nil {
return fmt.Errorf("failed to parse WebSocket URL: %w", err)
}
q := u.Query()
q.Set("token", token)
u.RawQuery = q.Encode()
conn, _, err := websocket.DefaultDialer.Dial(u.String(), nil)
if err != nil {
return fmt.Errorf("failed to connect to WebSocket: %w", err)
}
c.conn = conn
go c.readPump()
return nil
}
// Close closes the WebSocket connection
func (c *Client) Close() error {
close(c.done)
if c.conn != nil {
return c.conn.Close()
}
return nil
}
// SendMessage sends a message through the WebSocket connection
func (c *Client) SendMessage(messageType string, data interface{}) error {
if c.conn == nil {
return fmt.Errorf("not connected")
}
msg := WSMessage{
Type: messageType,
Data: data,
}
return c.conn.WriteJSON(msg)
}
// RegisterHandler registers a handler for a specific message type
func (c *Client) RegisterHandler(messageType string, handler MessageHandler) {
c.handlersMux.Lock()
defer c.handlersMux.Unlock()
c.handlers[messageType] = handler
}
// readPump pumps messages from the WebSocket connection
func (c *Client) readPump() {
defer c.conn.Close()
for {
select {
case <-c.done:
return
default:
var msg WSMessage
err := c.conn.ReadJSON(&msg)
if err != nil {
log.Printf("read error: %v", err)
return
}
c.handlersMux.RLock()
if handler, ok := c.handlers[msg.Type]; ok {
handler(msg)
}
c.handlersMux.RUnlock()
}
}
}
func (c *Client) getToken() (string, error) {
// If we already have a token, try to use it
if c.config.Token != "" {
tokenCheckData := map[string]interface{}{
"newtId": c.config.NewtID,
"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))
if err != nil {
return "", err
}
defer resp.Body.Close()
var tokenResp TokenResponse
if err := json.NewDecoder(resp.Body).Decode(&tokenResp); err != nil {
return "", err
}
if tokenResp.Success && tokenResp.Message == "Token session already valid" {
return c.config.Token, nil
}
}
// Get a new token
tokenData := map[string]interface{}{
"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))
if err != nil {
return "", err
}
defer resp.Body.Close()
var tokenResp TokenResponse
if err := json.NewDecoder(resp.Body).Decode(&tokenResp); err != nil {
return "", err
}
if !tokenResp.Success {
return "", fmt.Errorf("failed to get token: %s", tokenResp.Message)
}
return tokenResp.Data.Token, nil
}

57
websocket/config.go Normal file
View File

@@ -0,0 +1,57 @@
package websocket
import (
"encoding/json"
"io/ioutil"
"log"
"os"
"path/filepath"
"runtime"
)
func getConfigPath() string {
var configDir string
switch runtime.GOOS {
case "darwin":
configDir = filepath.Join(os.Getenv("HOME"), "Library", "Application Support", "newt-client")
case "windows":
configDir = filepath.Join(os.Getenv("APPDATA"), "newt-client")
default: // linux and others
configDir = filepath.Join(os.Getenv("HOME"), ".config", "newt-client")
}
if err := os.MkdirAll(configDir, 0755); err != nil {
log.Printf("Failed to create config directory: %v", err)
}
return filepath.Join(configDir, "config.json")
}
func (c *Client) loadConfig() error {
configPath := getConfigPath()
data, err := ioutil.ReadFile(configPath)
if err != nil {
if os.IsNotExist(err) {
return nil
}
return err
}
var config Config
if err := json.Unmarshal(data, &config); err != nil {
return err
}
// Only update token from saved config
c.config.Token = config.Token
return nil
}
func (c *Client) saveConfig() error {
configPath := getConfigPath()
data, err := json.MarshalIndent(c.config, "", " ")
if err != nil {
return err
}
return ioutil.WriteFile(configPath, data, 0644)
}

View File

@@ -1,60 +0,0 @@
package websocket
import (
"encoding/json"
"fmt"
"log"
"net/http"
"github.com/gorilla/websocket"
)
func connectWebSocket(url, token string) error {
// Create custom header with the auth token
header := http.Header{}
header.Add("Sec-WebSocket-Protocol", token)
// Create dialer with default options
dialer := websocket.Dialer{
EnableCompression: true,
}
// Connect to WebSocket server
conn, resp, err := dialer.Dial(url, header)
if err != nil {
log.Printf("Dial failed: %v", err)
if resp != nil {
log.Printf("HTTP Response Status: %s", resp.Status)
}
return err
}
defer conn.Close()
log.Printf("Connected to WebSocket server")
// Message handling loop
for {
// Read message
messageType, message, err := conn.ReadMessage()
if err != nil {
log.Printf("Read error: %v", err)
return err
}
// Handle text messages (JSON expected)
if messageType == websocket.TextMessage {
// Create a map to store the JSON data
var jsonData map[string]interface{}
// Unmarshal the JSON message
if err := json.Unmarshal(message, &jsonData); err != nil {
log.Printf("JSON parsing error: %v", err)
// Continue reading messages even if one fails to parse
continue
}
// Print the parsed JSON message
fmt.Printf("Received message: %+v\n", jsonData)
}
}
}

20
websocket/types.go Normal file
View File

@@ -0,0 +1,20 @@
package websocket
type Config struct {
NewtID string `json:"newtId"`
Secret string `json:"secret"`
Token string `json:"token"`
}
type TokenResponse struct {
Data struct {
Token string `json:"token"`
} `json:"data"`
Success bool `json:"success"`
Message string `json:"message"`
}
type WSMessage struct {
Type string `json:"type"`
Data interface{} `json:"data"`
}