test_sse_server.go

package main

import (
	"encoding/json"
	"fmt"
	"log"
	"net/http"
	"sync"
	"time"
)

// SSE Client structure
type SSEClient struct {
	UserID int
	Send   chan []byte
}

var (
	sseClients      = make(map[chan []byte]int)
	sseClientsMutex sync.Mutex
	newClients      = make(chan SSEClient)
	deadClients     = make(chan chan []byte)
	messages        = make(chan []byte)
)

// SSE Handler - handles client connections
func SSEHandler(w http.ResponseWriter, r *http.Request) {
	// Set headers for SSE
	w.Header().Set("Content-Type", "text/event-stream")
	w.Header().Set("Cache-Control", "no-cache")
	w.Header().Set("Connection", "keep-alive")
	w.Header().Set("Access-Control-Allow-Origin", "*")

	// Create a channel for this client
	clientChan := make(chan []byte, 10)

	// Simulate user ID (in real app, get from session)
	userID := 1

	// Register new client
	newClients <- SSEClient{
		UserID: userID,
		Send:   clientChan,
	}

	// Remove client when done
	defer func() {
		deadClients <- clientChan
	}()

	// Send initial connection message
	initialMsg := map[string]interface{}{
		"message": "Connected to SSE",
		"time":    time.Now().Format(time.RFC3339),
	}
	data, _ := json.Marshal(initialMsg)
	fmt.Fprintf(w, "data: %s\n\n", data)
	w.(http.Flusher).Flush()

	log.Printf("New SSE client connected (userID: %d)", userID)

	// Listen for messages to send to this client
	notify := r.Context().Done()
	for {
		select {
		case msg := <-clientChan:
			fmt.Fprintf(w, "data: %s\n\n", msg)
			w.(http.Flusher).Flush()
		case <-notify:
			log.Printf("SSE client disconnected (userID: %d)", userID)
			return
		}
	}
}

// BroadcastHandler - HTTP endpoint to broadcast messages
func BroadcastHandler(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodPost {
		http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
		return
	}

	var payload struct {
		Event  string                 `json:"event"`
		UserID int                    `json:"user_id"`
		Data   map[string]interface{} `json:"data"`
	}

	if err := json.NewDecoder(r.Body).Decode(&payload); err != nil {
		http.Error(w, "Invalid JSON", http.StatusBadRequest)
		return
	}

	// Create SSE message
	sseMessage := map[string]interface{}{
		"event": payload.Event,
		"data":  payload.Data,
	}

	msgBytes, err := json.Marshal(sseMessage)
	if err != nil {
		http.Error(w, "Failed to marshal message", http.StatusInternalServerError)
		return
	}

	// Send to specific user or broadcast to all
	sseClientsMutex.Lock()
	defer sseClientsMutex.Unlock()

	sentCount := 0
	for clientChan, clientUserID := range sseClients {
		if payload.UserID == 0 || clientUserID == payload.UserID {
			select {
			case clientChan <- msgBytes:
				sentCount++
			default:
				log.Printf("Warning: Client channel full, message dropped")
			}
		}
	}

	log.Printf("Broadcast: event=%s, userID=%d, sentTo=%d clients", payload.Event, payload.UserID, sentCount)

	w.Header().Set("Content-Type", "application/json")
	json.NewEncoder(w).Encode(map[string]interface{}{
		"success": true,
		"sent_to": sentCount,
	})
}

// SSE event loop - manages all clients
func SSEEventLoop() {
	log.Println("Starting SSE event loop...")
	for {
		select {
		case client := <-newClients:
			sseClientsMutex.Lock()
			sseClients[client.Send] = client.UserID
			sseClientsMutex.Unlock()
			log.Printf("Client added, total clients: %d", len(sseClients))

		case clientChan := <-deadClients:
			sseClientsMutex.Lock()
			delete(sseClients, clientChan)
			close(clientChan)
			sseClientsMutex.Unlock()
			log.Printf("Client removed, total clients: %d", len(sseClients))

		case msg := <-messages:
			sseClientsMutex.Lock()
			for clientChan := range sseClients {
				select {
				case clientChan <- msg:
				default:
					log.Println("Warning: Client channel full")
				}
			}
			sseClientsMutex.Unlock()
		}
	}
}

func main() {
	// Start SSE event loop
	go SSEEventLoop()

	// Routes
	http.HandleFunc("/sse", SSEHandler)
	http.HandleFunc("/internal/broadcast", BroadcastHandler)

	// Test page
	http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		html := `<!DOCTYPE html>
<html>
<head>
    <title>SSE Test</title>
</head>
<body>
    <h1>SSE Connection Test</h1>
    <div id="status">Connecting...</div>
    <div id="messages"></div>
    
    <script>
        const eventSource = new EventSource('/sse');
        const statusDiv = document.getElementById('status');
        const messagesDiv = document.getElementById('messages');
        
        eventSource.onopen = function() {
            statusDiv.innerHTML = '<span style="color: green;">✓ Connected to SSE</span>';
            console.log('SSE connection established');
        };
        
        eventSource.onmessage = function(event) {
            console.log('SSE message received:', event.data);
            const data = JSON.parse(event.data);
            const msgDiv = document.createElement('div');
            msgDiv.style.padding = '10px';
            msgDiv.style.margin = '5px 0';
            msgDiv.style.border = '1px solid #ccc';
            msgDiv.innerHTML = '<strong>' + (data.event || 'message') + ':</strong> ' + JSON.stringify(data, null, 2);
            messagesDiv.appendChild(msgDiv);
            
            // Handle logout event
            if (data.event === 'logout') {
                alert('Logout event received: ' + data.data.message);
            }
        };
        
        eventSource.onerror = function(error) {
            statusDiv.innerHTML = '<span style="color: red;">✗ SSE connection error</span>';
            console.error('SSE error:', error);
        };
    </script>
    
    <hr>
    <h2>Send Test Broadcast</h2>
    <button onclick="sendBroadcast()">Send Logout Event</button>
    
    <script>
        function sendBroadcast() {
            fetch('/internal/broadcast', {
                method: 'POST',
                headers: {'Content-Type': 'application/json'},
                body: JSON.stringify({
                    event: 'logout',
                    user_id: 1,
                    data: {
                        user_id: 1,
                        message: 'Your account has been deleted'
                    }
                })
            })
            .then(response => response.json())
            .then(data => {
                console.log('Broadcast sent:', data);
                alert('Broadcast sent to ' + data.sent_to + ' clients');
            })
            .catch(error => console.error('Broadcast error:', error));
        }
    </script>
</body>
</html>`
		w.Header().Set("Content-Type", "text/html")
		fmt.Fprint(w, html)
	})

	log.Println("Starting SSE test server on http://localhost:4000")
	log.Println("Visit http://localhost:4000 in your browser")
	log.Println("Or test with curl:")
	log.Println("  curl -N http://localhost:4000/sse")
	log.Fatal(http.ListenAndServe(":4000", nil))
}