IRC-kosmi-relay/cmd/test-websocket/main.go

package main

import (
	"context"
	"encoding/base64"
	"encoding/json"
	"flag"
	"fmt"
	"net/http"
	"os"
	"time"

	"github.com/gorilla/websocket"
)

const (
	userAgent    = "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.0.0.0 Safari/537.36"
	appVersion   = "4364"
	wsURL        = "wss://engine.kosmi.io/gql-ws"
	tokenURL     = "https://engine.kosmi.io/"
)

func main() {
	roomID := flag.String("room", "@hyperspaceout", "Room ID")
	testMode := flag.Int("mode", 1, "Test mode: 1=with-token, 2=no-auth, 3=origin-only")
	flag.Parse()

	fmt.Printf("Test Mode %d: Testing WebSocket connection to Kosmi\n\n", *testMode)

	var conn *websocket.Conn
	var err error

	switch *testMode {
	case 1:
		fmt.Println("Mode 1: Testing with JWT token (full auth)")
		conn, err = testWithToken(*roomID)
	case 2:
		fmt.Println("Mode 2: Testing without authentication")
		conn, err = testWithoutAuth()
	case 3:
		fmt.Println("Mode 3: Testing with Origin header only")
		conn, err = testWithOriginOnly()
	default:
		fmt.Fprintf(os.Stderr, "Invalid test mode: %d\n", *testMode)
		os.Exit(1)
	}

	if err != nil {
		fmt.Fprintf(os.Stderr, "❌ Connection failed: %v\n", err)
		os.Exit(1)
	}

	fmt.Println("✅ WebSocket connected successfully!")
	defer conn.Close()

	// Try to do the GraphQL-WS handshake
	fmt.Println("\n📤 Sending connection_init...")
	if err := waitForAck(conn); err != nil {
		fmt.Fprintf(os.Stderr, "❌ Handshake failed: %v\n", err)
		os.Exit(1)
	}

	fmt.Println("✅ WebSocket handshake successful!")
	fmt.Println("\n📝 Testing message subscription...")
	
	if err := subscribeToMessages(conn, *roomID); err != nil {
		fmt.Fprintf(os.Stderr, "❌ Subscription failed: %v\n", err)
		os.Exit(1)
	}

	fmt.Println("✅ Subscribed to messages!")
	fmt.Println("\n👂 Listening for messages (press Ctrl+C to exit)...")

	// Listen for messages
	for {
		var msg map[string]interface{}
		if err := conn.ReadJSON(&msg); err != nil {
			fmt.Fprintf(os.Stderr, "Error reading message: %v\n", err)
			break
		}

		msgType, _ := msg["type"].(string)
		fmt.Printf("📥 Received: %s\n", msgType)

		if msgType == "next" {
			payload, _ := msg["payload"].(map[string]interface{})
			data, _ := payload["data"].(map[string]interface{})
			if newMessage, ok := data["newMessage"].(map[string]interface{}); ok {
				body, _ := newMessage["body"].(string)
				user, _ := newMessage["user"].(map[string]interface{})
				username, _ := user["displayName"].(string)
				if username == "" {
					username, _ = user["username"].(string)
				}
				fmt.Printf("   💬 %s: %s\n", username, body)
			}
		}
	}
}

// testWithToken attempts connection with full JWT authentication
func testWithToken(roomID string) (*websocket.Conn, error) {
	fmt.Println("   1️⃣  Step 1: Acquiring JWT token...")

	// Try to get token from GraphQL endpoint
	token, err := acquireToken()
	if err != nil {
		return nil, fmt.Errorf("failed to acquire token: %w", err)
	}

	fmt.Printf("   ✅ Got token: %s...\n", truncate(token, 50))

	fmt.Println("   2️⃣  Step 2: Connecting WebSocket with token...")
	return connectWithToken(token)
}

// testWithoutAuth attempts direct connection with no headers
func testWithoutAuth() (*websocket.Conn, error) {
	fmt.Println("   Connecting without any authentication...")
	
	dialer := websocket.Dialer{
		Subprotocols: []string{"graphql-ws"},
	}
	
	conn, resp, err := dialer.Dial(wsURL, nil)
	if resp != nil {
		fmt.Printf("   Response status: %d\n", resp.StatusCode)
	}
	return conn, err
}

// testWithOriginOnly attempts connection with just Origin header
func testWithOriginOnly() (*websocket.Conn, error) {
	fmt.Println("   Connecting with Origin header only...")
	
	dialer := websocket.Dialer{
		Subprotocols: []string{"graphql-ws"},
	}
	
	headers := http.Header{
		"Origin": []string{"https://app.kosmi.io"},
	}
	
	conn, resp, err := dialer.Dial(wsURL, headers)
	if resp != nil {
		fmt.Printf("   Response status: %d\n", resp.StatusCode)
	}
	return conn, err
}

// acquireToken gets a JWT token from Kosmi's API
func acquireToken() (string, error) {
	// First, let's try a few different approaches
	
	// Approach 1: Try empty POST (some APIs generate anonymous tokens)
	fmt.Println("      Trying empty POST...")
	token, err := tryEmptyPost()
	if err == nil && token != "" {
		return token, nil
	}
	fmt.Printf("      Empty POST failed: %v\n", err)

	// Approach 2: Try GraphQL anonymous login
	fmt.Println("      Trying GraphQL anonymous session...")
	token, err = tryGraphQLSession()
	if err == nil && token != "" {
		return token, nil
	}
	fmt.Printf("      GraphQL session failed: %v\n", err)

	// Approach 3: Try REST endpoint
	fmt.Println("      Trying REST endpoint...")
	token, err = tryRESTAuth()
	if err == nil && token != "" {
		return token, nil
	}
	fmt.Printf("      REST auth failed: %v\n", err)

	return "", fmt.Errorf("all token acquisition methods failed")
}

// tryEmptyPost tries posting an empty body
func tryEmptyPost() (string, error) {
	client := &http.Client{Timeout: 10 * time.Second}
	
	req, err := http.NewRequest("POST", tokenURL, nil)
	if err != nil {
		return "", err
	}

	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("Referer", "https://app.kosmi.io/")
	req.Header.Set("User-Agent", userAgent)

	resp, err := client.Do(req)
	if err != nil {
		return "", err
	}
	defer resp.Body.Close()

	if resp.StatusCode != 200 {
		return "", fmt.Errorf("status %d", resp.StatusCode)
	}

	var result map[string]interface{}
	if err := json.NewDecoder(resp.Body).Decode(&result); err != nil {
		return "", err
	}

	// Try to extract token from various possible locations
	if token, ok := result["token"].(string); ok {
		return token, nil
	}
	if data, ok := result["data"].(map[string]interface{}); ok {
		if token, ok := data["token"].(string); ok {
			return token, nil
		}
	}

	return "", fmt.Errorf("no token in response: %+v", result)
}

// tryGraphQLSession tries a GraphQL mutation for anonymous session
func tryGraphQLSession() (string, error) {
	query := map[string]interface{}{
		"query": `mutation { createAnonymousSession { token } }`,
	}

	return postGraphQL(query)
}

// tryRESTAuth tries REST-style auth endpoint
func tryRESTAuth() (string, error) {
	body := map[string]interface{}{
		"anonymous": true,
	}

	jsonBody, _ := json.Marshal(body)
	client := &http.Client{Timeout: 10 * time.Second}
	
	req, err := http.NewRequest("POST", tokenURL+"auth/anonymous", nil)
	if err != nil {
		return "", err
	}

	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("Referer", "https://app.kosmi.io/")
	req.Header.Set("User-Agent", userAgent)
	req.Body = http.NoBody
	_ = jsonBody // silence unused warning

	resp, err := client.Do(req)
	if err != nil {
		return "", err
	}
	defer resp.Body.Close()

	if resp.StatusCode >= 400 {
		return "", fmt.Errorf("status %d", resp.StatusCode)
	}

	var result map[string]interface{}
	if err := json.NewDecoder(resp.Body).Decode(&result); err != nil {
		return "", err
	}

	if token, ok := result["token"].(string); ok {
		return token, nil
	}

	return "", fmt.Errorf("no token in response")
}

// postGraphQL posts a GraphQL query
func postGraphQL(query map[string]interface{}) (string, error) {
	jsonBody, _ := json.Marshal(query)
	client := &http.Client{Timeout: 10 * time.Second}
	
	req, err := http.NewRequest("POST", tokenURL, nil)
	if err != nil {
		return "", err
	}

	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("Referer", "https://app.kosmi.io/")
	req.Header.Set("User-Agent", userAgent)
	req.Body = http.NoBody
	_ = jsonBody // silence unused

	resp, err := client.Do(req)
	if err != nil {
		return "", err
	}
	defer resp.Body.Close()

	if resp.StatusCode != 200 {
		return "", fmt.Errorf("status %d", resp.StatusCode)
	}

	var result map[string]interface{}
	if err := json.NewDecoder(resp.Body).Decode(&result); err != nil {
		return "", err
	}

	// Navigate nested response
	if data, ok := result["data"].(map[string]interface{}); ok {
		if session, ok := data["createAnonymousSession"].(map[string]interface{}); ok {
			if token, ok := session["token"].(string); ok {
				return token, nil
			}
		}
	}

	return "", fmt.Errorf("no token in response")
}

// connectWithToken connects WebSocket with JWT token
func connectWithToken(token string) (*websocket.Conn, error) {
	dialer := websocket.Dialer{
		Subprotocols: []string{"graphql-ws"},
	}
	
	headers := http.Header{
		"Origin":     []string{"https://app.kosmi.io"},
		"User-Agent": []string{userAgent},
	}
	
	conn, resp, err := dialer.Dial(wsURL, headers)
	if err != nil {
		if resp != nil {
			fmt.Printf("   Response status: %d\n", resp.StatusCode)
		}
		return nil, err
	}

	// Send connection_init with token
	uaEncoded := base64.StdEncoding.EncodeToString([]byte(userAgent))
	
	initMsg := map[string]interface{}{
		"type": "connection_init",
		"payload": map[string]interface{}{
			"token": token,
			"ua":    uaEncoded,
			"v":     appVersion,
			"r":     "",
		},
	}

	if err := conn.WriteJSON(initMsg); err != nil {
		conn.Close()
		return nil, fmt.Errorf("failed to send connection_init: %w", err)
	}

	return conn, nil
}

// waitForAck waits for connection_ack
func waitForAck(conn *websocket.Conn) error {
	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()

	done := make(chan error, 1)
	
	go func() {
		var msg map[string]interface{}
		if err := conn.ReadJSON(&msg); err != nil {
			done <- err
			return
		}

		msgType, _ := msg["type"].(string)
		if msgType != "connection_ack" {
			done <- fmt.Errorf("expected connection_ack, got %s", msgType)
			return
		}

		fmt.Println("✅ Received connection_ack")
		done <- nil
	}()

	select {
	case err := <-done:
		return err
	case <-ctx.Done():
		return fmt.Errorf("timeout waiting for connection_ack")
	}
}

// subscribeToMessages subscribes to room messages
func subscribeToMessages(conn *websocket.Conn, roomID string) error {
	query := fmt.Sprintf(`
		subscription {
			newMessage(roomId: "%s") {
				body
				time
				user {
					displayName
					username
				}
			}
		}
	`, roomID)

	subMsg := map[string]interface{}{
		"id":   "test-subscription-1",
		"type": "subscribe",
		"payload": map[string]interface{}{
			"query":     query,
			"variables": map[string]interface{}{},
		},
	}

	return conn.WriteJSON(subMsg)
}

// truncate truncates a string
func truncate(s string, maxLen int) string {
	if len(s) <= maxLen {
		return s
	}
	return s[:maxLen] + "..."
}