335 lines
9.6 KiB
Go
335 lines
9.6 KiB
Go
|
package main
|
||
|
|
|
||
|
|
import (
|
||
|
|
"encoding/json"
|
||
|
|
"fmt"
|
||
|
|
"log"
|
||
|
|
"sync"
|
||
|
|
"time"
|
||
|
|
|
||
|
|
"github.com/gorilla/websocket"
|
||
|
|
)
|
||
|
|
|
||
|
|
const (
|
||
|
|
kosmiWebSocketURL = "wss://engine.kosmi.io/gql-ws"
|
||
|
|
)
|
||
|
|
|
||
|
|
// KosmiClient represents a WebSocket client for Kosmi
|
||
|
|
type KosmiClient struct {
|
||
|
|
conn *websocket.Conn
|
||
|
|
token string
|
||
|
|
roomID string
|
||
|
|
mu sync.Mutex
|
||
|
|
msgID int
|
||
|
|
onMessage func(Message)
|
||
|
|
closeChan chan struct{}
|
||
|
|
closedOnce sync.Once
|
||
|
|
}
|
||
|
|
|
||
|
|
// Message represents an incoming chat message
|
||
|
|
type Message struct {
|
||
|
|
ID string
|
||
|
|
Username string
|
||
|
|
DisplayName string
|
||
|
|
Body string
|
||
|
|
Time string
|
||
|
|
}
|
||
|
|
|
||
|
|
// GraphQL message types
|
||
|
|
type gqlMessage struct {
|
||
|
|
ID string `json:"id,omitempty"`
|
||
|
|
Type string `json:"type"`
|
||
|
|
Payload json.RawMessage `json:"payload,omitempty"`
|
||
|
|
}
|
||
|
|
|
||
|
|
// NewKosmiClient creates a new Kosmi client
|
||
|
|
func NewKosmiClient(token, roomID string) *KosmiClient {
|
||
|
|
return &KosmiClient{
|
||
|
|
token: token,
|
||
|
|
roomID: roomID,
|
||
|
|
closeChan: make(chan struct{}),
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
// SetMessageHandler sets the callback for incoming messages
|
||
|
|
func (c *KosmiClient) SetMessageHandler(handler func(Message)) {
|
||
|
|
c.onMessage = handler
|
||
|
|
}
|
||
|
|
|
||
|
|
// Connect establishes the WebSocket connection and performs the handshake
|
||
|
|
func (c *KosmiClient) Connect() error {
|
||
|
|
log.Println("Connecting to Kosmi WebSocket...")
|
||
|
|
|
||
|
|
// Establish WebSocket connection
|
||
|
|
conn, _, err := websocket.DefaultDialer.Dial(kosmiWebSocketURL, nil)
|
||
|
|
if err != nil {
|
||
|
|
return fmt.Errorf("failed to connect to WebSocket: %w", err)
|
||
|
|
}
|
||
|
|
c.conn = conn
|
||
|
|
|
||
|
|
// Step 1: Send connection_init with token
|
||
|
|
log.Println("Sending connection_init...")
|
||
|
|
initMsg := gqlMessage{
|
||
|
|
Type: "connection_init",
|
||
|
|
Payload: json.RawMessage(fmt.Sprintf(`{"token":"%s"}`, c.token)),
|
||
|
|
}
|
||
|
|
if err := c.conn.WriteJSON(initMsg); err != nil {
|
||
|
|
return fmt.Errorf("failed to send connection_init: %w", err)
|
||
|
|
}
|
||
|
|
|
||
|
|
// Step 2: Wait for connection_ack (synchronous)
|
||
|
|
var ackMsg gqlMessage
|
||
|
|
if err := c.conn.ReadJSON(&ackMsg); err != nil {
|
||
|
|
return fmt.Errorf("failed to read connection_ack: %w", err)
|
||
|
|
}
|
||
|
|
if ackMsg.Type != "connection_ack" {
|
||
|
|
return fmt.Errorf("expected connection_ack, got: %s", ackMsg.Type)
|
||
|
|
}
|
||
|
|
log.Println("✓ Connection acknowledged")
|
||
|
|
|
||
|
|
// Step 3: Send ExtendedCurrentUserQuery (synchronous)
|
||
|
|
log.Println("Querying current user...")
|
||
|
|
userQueryMsg := gqlMessage{
|
||
|
|
ID: c.nextMsgID(),
|
||
|
|
Type: "subscribe",
|
||
|
|
Payload: json.RawMessage(`{"query":"query ExtendedCurrentUserQuery { currentUser { id connectionId user { id displayName username isAnonymous avatarUrl email } } }","operationName":"ExtendedCurrentUserQuery","variables":{},"extensions":{}}`),
|
||
|
|
}
|
||
|
|
if err := c.conn.WriteJSON(userQueryMsg); err != nil {
|
||
|
|
return fmt.Errorf("failed to send ExtendedCurrentUserQuery: %w", err)
|
||
|
|
}
|
||
|
|
|
||
|
|
// Read user query response (synchronous)
|
||
|
|
var userResponse gqlMessage
|
||
|
|
if err := c.conn.ReadJSON(&userResponse); err != nil {
|
||
|
|
return fmt.Errorf("failed to read user query response: %w", err)
|
||
|
|
}
|
||
|
|
if userResponse.Type == "next" {
|
||
|
|
var userData struct {
|
||
|
|
Data struct {
|
||
|
|
CurrentUser struct {
|
||
|
|
User struct {
|
||
|
|
DisplayName string `json:"displayName"`
|
||
|
|
Username string `json:"username"`
|
||
|
|
IsAnonymous bool `json:"isAnonymous"`
|
||
|
|
} `json:"user"`
|
||
|
|
} `json:"currentUser"`
|
||
|
|
} `json:"data"`
|
||
|
|
}
|
||
|
|
if err := json.Unmarshal(userResponse.Payload, &userData); err == nil {
|
||
|
|
log.Printf("✓ Logged in as: %s (@%s) [Anonymous: %v]",
|
||
|
|
userData.Data.CurrentUser.User.DisplayName,
|
||
|
|
userData.Data.CurrentUser.User.Username,
|
||
|
|
userData.Data.CurrentUser.User.IsAnonymous)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
// Read complete message for user query
|
||
|
|
var completeMsg gqlMessage
|
||
|
|
if err := c.conn.ReadJSON(&completeMsg); err != nil {
|
||
|
|
return fmt.Errorf("failed to read complete message: %w", err)
|
||
|
|
}
|
||
|
|
|
||
|
|
// Step 4: Send JoinRoom mutation (synchronous)
|
||
|
|
log.Printf("Joining room: %s...", c.roomID)
|
||
|
|
joinRoomMsg := gqlMessage{
|
||
|
|
ID: c.nextMsgID(),
|
||
|
|
Type: "subscribe",
|
||
|
|
Payload: json.RawMessage(fmt.Sprintf(`{"query":"mutation JoinRoom($id: String!, $disconnectOtherConnections: Boolean) { joinRoom(id: $id, disconnectOtherConnections: $disconnectOtherConnections) { ok __typename } }","operationName":"JoinRoom","variables":{"id":"%s","disconnectOtherConnections":false},"extensions":{}}`, c.roomID)),
|
||
|
|
}
|
||
|
|
if err := c.conn.WriteJSON(joinRoomMsg); err != nil {
|
||
|
|
return fmt.Errorf("failed to send JoinRoom: %w", err)
|
||
|
|
}
|
||
|
|
|
||
|
|
// Read join room response (synchronous)
|
||
|
|
var joinResponse gqlMessage
|
||
|
|
if err := c.conn.ReadJSON(&joinResponse); err != nil {
|
||
|
|
return fmt.Errorf("failed to read join room response: %w", err)
|
||
|
|
}
|
||
|
|
if joinResponse.Type == "next" {
|
||
|
|
var joinData struct {
|
||
|
|
Data struct {
|
||
|
|
JoinRoom struct {
|
||
|
|
Ok bool `json:"ok"`
|
||
|
|
} `json:"joinRoom"`
|
||
|
|
} `json:"data"`
|
||
|
|
}
|
||
|
|
if err := json.Unmarshal(joinResponse.Payload, &joinData); err == nil {
|
||
|
|
if joinData.Data.JoinRoom.Ok {
|
||
|
|
log.Println("✓ Successfully joined room")
|
||
|
|
} else {
|
||
|
|
log.Println("⚠ Join room returned ok=false")
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
// Read complete message for join room
|
||
|
|
if err := c.conn.ReadJSON(&completeMsg); err != nil {
|
||
|
|
return fmt.Errorf("failed to read join complete message: %w", err)
|
||
|
|
}
|
||
|
|
|
||
|
|
// Step 5: Subscribe to NewMessageSubscription (synchronous)
|
||
|
|
log.Println("Subscribing to new messages...")
|
||
|
|
subscribeMsg := gqlMessage{
|
||
|
|
ID: c.nextMsgID(),
|
||
|
|
Type: "subscribe",
|
||
|
|
Payload: json.RawMessage(fmt.Sprintf(`{"query":"subscription NewMessageSubscription($roomId: String!, $channelId: String) { newMessage(roomId: $roomId, channelId: $channelId) { id user { id displayName username } body time } }","operationName":"NewMessageSubscription","variables":{"roomId":"%s","channelId":null},"extensions":{}}`, c.roomID)),
|
||
|
|
}
|
||
|
|
if err := c.conn.WriteJSON(subscribeMsg); err != nil {
|
||
|
|
return fmt.Errorf("failed to send NewMessageSubscription: %w", err)
|
||
|
|
}
|
||
|
|
|
||
|
|
log.Println("✓ Subscription active")
|
||
|
|
log.Println("Connection setup complete!")
|
||
|
|
|
||
|
|
// Step 6: Start listening for messages in goroutine
|
||
|
|
go c.listenForMessages()
|
||
|
|
|
||
|
|
return nil
|
||
|
|
}
|
||
|
|
|
||
|
|
// listenForMessages continuously reads messages from the WebSocket
|
||
|
|
func (c *KosmiClient) listenForMessages() {
|
||
|
|
defer func() {
|
||
|
|
if r := recover(); r != nil {
|
||
|
|
log.Printf("Panic in listenForMessages: %v", r)
|
||
|
|
}
|
||
|
|
}()
|
||
|
|
|
||
|
|
for {
|
||
|
|
select {
|
||
|
|
case <-c.closeChan:
|
||
|
|
log.Println("Message listener shutting down...")
|
||
|
|
return
|
||
|
|
default:
|
||
|
|
}
|
||
|
|
|
||
|
|
var msg gqlMessage
|
||
|
|
if err := c.conn.ReadJSON(&msg); err != nil {
|
||
|
|
select {
|
||
|
|
case <-c.closeChan:
|
||
|
|
// Expected error during shutdown
|
||
|
|
return
|
||
|
|
default:
|
||
|
|
log.Printf("Error reading message: %v", err)
|
||
|
|
return
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
// Handle different message types
|
||
|
|
switch msg.Type {
|
||
|
|
case "next":
|
||
|
|
c.handleNextMessage(msg.Payload)
|
||
|
|
case "error":
|
||
|
|
log.Printf("GraphQL error: %s", string(msg.Payload))
|
||
|
|
case "complete":
|
||
|
|
log.Printf("Subscription completed: %s", msg.ID)
|
||
|
|
case "ka":
|
||
|
|
// Keep-alive, ignore
|
||
|
|
default:
|
||
|
|
log.Printf("Unknown message type: %s", msg.Type)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
// handleNextMessage processes a "next" type message (data payload)
|
||
|
|
func (c *KosmiClient) handleNextMessage(payload json.RawMessage) {
|
||
|
|
var response struct {
|
||
|
|
Data struct {
|
||
|
|
NewMessage struct {
|
||
|
|
ID string `json:"id"`
|
||
|
|
User struct {
|
||
|
|
ID string `json:"id"`
|
||
|
|
DisplayName string `json:"displayName"`
|
||
|
|
Username string `json:"username"`
|
||
|
|
} `json:"user"`
|
||
|
|
Body string `json:"body"`
|
||
|
|
Time interface{} `json:"time"` // Can be string or number
|
||
|
|
} `json:"newMessage"`
|
||
|
|
} `json:"data"`
|
||
|
|
}
|
||
|
|
|
||
|
|
if err := json.Unmarshal(payload, &response); err != nil {
|
||
|
|
log.Printf("Failed to parse message payload: %v", err)
|
||
|
|
return
|
||
|
|
}
|
||
|
|
|
||
|
|
// Check if this is a new message
|
||
|
|
if response.Data.NewMessage.ID != "" {
|
||
|
|
// Convert time to string
|
||
|
|
timeStr := ""
|
||
|
|
switch t := response.Data.NewMessage.Time.(type) {
|
||
|
|
case string:
|
||
|
|
timeStr = t
|
||
|
|
case float64:
|
||
|
|
timeStr = time.Unix(int64(t), 0).Format(time.RFC3339)
|
||
|
|
}
|
||
|
|
|
||
|
|
msg := Message{
|
||
|
|
ID: response.Data.NewMessage.ID,
|
||
|
|
Username: response.Data.NewMessage.User.Username,
|
||
|
|
DisplayName: response.Data.NewMessage.User.DisplayName,
|
||
|
|
Body: response.Data.NewMessage.Body,
|
||
|
|
Time: timeStr,
|
||
|
|
}
|
||
|
|
|
||
|
|
// Call the message handler if set
|
||
|
|
if c.onMessage != nil {
|
||
|
|
c.onMessage(msg)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
// SendMessage sends a message to the room
|
||
|
|
func (c *KosmiClient) SendMessage(text string) error {
|
||
|
|
c.mu.Lock()
|
||
|
|
defer c.mu.Unlock()
|
||
|
|
|
||
|
|
sendMsg := gqlMessage{
|
||
|
|
ID: c.nextMsgID(),
|
||
|
|
Type: "subscribe",
|
||
|
|
Payload: json.RawMessage(fmt.Sprintf(`{"query":"mutation SendMessage2($roomId: String!, $body: String!, $channelId: String, $replyToMessageId: String) { sendMessage(roomId: $roomId, body: $body, channelId: $channelId, replyToMessageId: $replyToMessageId) { ok __typename } }","operationName":"SendMessage2","variables":{"roomId":"%s","body":%s,"channelId":null,"replyToMessageId":null},"extensions":{}}`, c.roomID, jsonEscape(text))),
|
||
|
|
}
|
||
|
|
|
||
|
|
if err := c.conn.WriteJSON(sendMsg); err != nil {
|
||
|
|
return fmt.Errorf("failed to send message: %w", err)
|
||
|
|
}
|
||
|
|
|
||
|
|
return nil
|
||
|
|
}
|
||
|
|
|
||
|
|
// Close gracefully closes the WebSocket connection
|
||
|
|
func (c *KosmiClient) Close() error {
|
||
|
|
var err error
|
||
|
|
c.closedOnce.Do(func() {
|
||
|
|
log.Println("Closing Kosmi client...")
|
||
|
|
close(c.closeChan)
|
||
|
|
|
||
|
|
// Give the listener goroutine a moment to exit
|
||
|
|
time.Sleep(100 * time.Millisecond)
|
||
|
|
|
||
|
|
if c.conn != nil {
|
||
|
|
// Send close message
|
||
|
|
closeMsg := websocket.FormatCloseMessage(websocket.CloseNormalClosure, "")
|
||
|
|
c.conn.WriteMessage(websocket.CloseMessage, closeMsg)
|
||
|
|
|
||
|
|
// Close the connection
|
||
|
|
err = c.conn.Close()
|
||
|
|
}
|
||
|
|
log.Println("✓ Client closed")
|
||
|
|
})
|
||
|
|
return err
|
||
|
|
}
|
||
|
|
|
||
|
|
// nextMsgID generates the next message ID
|
||
|
|
func (c *KosmiClient) nextMsgID() string {
|
||
|
|
c.msgID++
|
||
|
|
return fmt.Sprintf("%d", c.msgID)
|
||
|
|
}
|
||
|
|
|
||
|
|
// jsonEscape properly escapes a string for JSON
|
||
|
|
func jsonEscape(s string) string {
|
||
|
|
b, _ := json.Marshal(s)
|
||
|
|
return string(b)
|
||
|
|
}
|
||
|
|
|