meshstream/server/server.go

package server

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

	"github.com/dpup/prefab"

	"meshstream/mqtt"
)

// Config holds server configuration
type Config struct {
	Host   string
	Port   string
	Broker *mqtt.Broker // The MQTT message broker
}

// Server encapsulates the HTTP server functionality
type Server struct {
	config Config
	server *prefab.Server
	// Channel to signal shutdown to active connections
	shutdown chan struct{}
	// Atomic flag to indicate if server is shutting down
	isShuttingDown atomic.Bool
}

// New creates a new server instance
func New(config Config) *Server {
	if config.Broker == nil {
		log.Println("Warning: Server created without a broker, streaming will not work")
	}
	
	return &Server{
		config:   config,
		shutdown: make(chan struct{}),
	}
}

// Start initializes and starts the web server
func (s *Server) Start() error {
	// Get port as integer
	port, err := strconv.Atoi(s.config.Port)
	if err != nil {
		return err
	}

	// Create a new prefab server
	s.server = prefab.New(
		prefab.WithHost(s.config.Host),
		prefab.WithPort(port),
		prefab.WithHTTPHandlerFunc("/api/status", s.handleStatus),
		prefab.WithHTTPHandlerFunc("/api/stream", s.handleStream),
		prefab.WithStaticFiles("/", "./server/static"),
	)

	// Start the server
	log.Printf("Starting server on %s:%s", s.config.Host, s.config.Port)
	return s.server.Start()
}

// Stop shuts down the server
func (s *Server) Stop() error {
	// Set the shutdown flag first to prevent new connections from starting streams
	s.isShuttingDown.Store(true)
	
	// Signal all active connections to close
	close(s.shutdown)
	
	// Then shut down the HTTP server
	if s.server != nil {
		return s.server.Shutdown()
	}
	return nil
}

// handleStatus is a placeholder API endpoint that returns server status
func (s *Server) handleStatus(w http.ResponseWriter, r *http.Request) {
	status := map[string]interface{}{
		"status": "ok",
		"message": "Meshtastic Stream API is running",
	}
	
	w.Header().Set("Content-Type", "application/json")
	json.NewEncoder(w).Encode(status)
}

// handleStream handles Server-Sent Events streaming of MQTT messages
func (s *Server) handleStream(w http.ResponseWriter, r *http.Request) {
	// Check if the server is shutting down
	if s.isShuttingDown.Load() {
		http.Error(w, "Server is shutting down", http.StatusServiceUnavailable)
		return
	}
	
	// Check if broker is available
	if s.config.Broker == nil {
		http.Error(w, "MQTT broker not available", http.StatusServiceUnavailable)
		return
	}

	// 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", "*")

	// Make sure that the writer supports flushing
	flusher, ok := w.(http.Flusher)
	if !ok {
		http.Error(w, "Streaming not supported", http.StatusInternalServerError)
		return
	}

	// Subscribe to the broker with a buffer size of 10
	packetChan := s.config.Broker.Subscribe(10)
	
	// Signal when the client disconnects
	notify := r.Context().Done()
	
	// Send an initial message
	fmt.Fprintf(w, "event: info\ndata: Connected to Meshtastic stream\n\n")
	flusher.Flush()

	// Stream messages to the client
	for {
		select {
		case <-notify:
			// Client disconnected, unsubscribe and return
			log.Println("Client disconnected, unsubscribing from broker")
			s.config.Broker.Unsubscribe(packetChan)
			return
			
		case <-s.shutdown:
			// Server is shutting down, send a message to client and close
			log.Println("Server shutting down, closing SSE connection")
			fmt.Fprintf(w, "event: info\ndata: Server shutting down, connection closed\n\n")
			flusher.Flush()
			s.config.Broker.Unsubscribe(packetChan)
			return
			
		case packet, ok := <-packetChan:
			if !ok {
				// Channel closed, probably shutting down
				log.Println("Packet channel closed, ending stream")
				return
			}
			
			if packet == nil {
				continue
			}
			
			// Create a serializable wrapper for the packet
			// That includes both the entire packet and some extra fields for convenience
			packetWrapper := struct {
				*mqtt.Packet
				ReceivedAt int64  `json:"received_at"`
				PortString string `json:"port_string"`
			}{
				Packet:     packet,
				ReceivedAt: time.Now().Unix(),
				PortString: packet.PortNum.String(),
			}
			
			// Convert the entire packet to JSON
			data, err := json.Marshal(packetWrapper)
			
			if err != nil {
				log.Printf("Error marshaling packet to JSON: %v", err)
				continue
			}
			
			// Send the event
			fmt.Fprintf(w, "event: message\ndata: %s\n\n", data)
			flusher.Flush()
		}
	}
}