iris/host/livereload.go

package host

import (
	"bufio"
	"crypto/sha1"
	"encoding/base64"
	"encoding/binary"
	"io"
	"net"
	"net/http"
	"sync"
)

// LiveReload manages WebSocket connections for browser reload notifications.
type LiveReload struct {
	mu      sync.RWMutex
	clients map[*wsConn]bool
}

// NewLiveReload creates a new LiveReload manager.
func NewLiveReload() *LiveReload {
	return &LiveReload{
		clients: make(map[*wsConn]bool),
	}
}

// ServeHTTP handles WebSocket upgrade requests.
func (lr *LiveReload) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	conn, err := lr.upgrade(w, r)
	if err != nil {
		http.Error(w, "WebSocket upgrade failed", http.StatusBadRequest)
		return
	}

	lr.addClient(conn)
	defer lr.removeClient(conn)

	// Keep connection alive and handle pings
	conn.readLoop()
}

// NotifyReload sends a reload message to all connected clients.
func (lr *LiveReload) NotifyReload() {
	lr.mu.RLock()
	defer lr.mu.RUnlock()

	for conn := range lr.clients {
		conn.send([]byte("reload"))
	}
}

// ClientCount returns the number of connected clients.
func (lr *LiveReload) ClientCount() int {
	lr.mu.RLock()
	defer lr.mu.RUnlock()
	return len(lr.clients)
}

func (lr *LiveReload) addClient(conn *wsConn) {
	lr.mu.Lock()
	defer lr.mu.Unlock()
	lr.clients[conn] = true
}

func (lr *LiveReload) removeClient(conn *wsConn) {
	lr.mu.Lock()
	defer lr.mu.Unlock()
	delete(lr.clients, conn)
	conn.close()
}

func (lr *LiveReload) upgrade(w http.ResponseWriter, r *http.Request) (*wsConn, error) {
	if r.Header.Get("Upgrade") != "websocket" {
		return nil, io.ErrUnexpectedEOF
	}

	key := r.Header.Get("Sec-WebSocket-Key")
	if key == "" {
		return nil, io.ErrUnexpectedEOF
	}

	// Compute accept key
	h := sha1.New()
	h.Write([]byte(key + "258EAFA5-E914-47DA-95CA-C5AB0DC85B11"))
	acceptKey := base64.StdEncoding.EncodeToString(h.Sum(nil))

	// Hijack the connection
	hijacker, ok := w.(http.Hijacker)
	if !ok {
		return nil, io.ErrUnexpectedEOF
	}

	conn, bufrw, err := hijacker.Hijack()
	if err != nil {
		return nil, err
	}

	// Send upgrade response
	bufrw.WriteString("HTTP/1.1 101 Switching Protocols\r\n")
	bufrw.WriteString("Upgrade: websocket\r\n")
	bufrw.WriteString("Connection: Upgrade\r\n")
	bufrw.WriteString("Sec-WebSocket-Accept: " + acceptKey + "\r\n")
	bufrw.WriteString("\r\n")
	bufrw.Flush()

	return &wsConn{
		conn: conn,
		rw:   bufrw,
	}, nil
}

// wsConn is a minimal WebSocket connection.
type wsConn struct {
	conn   net.Conn
	rw     *bufio.ReadWriter
	mu     sync.Mutex
	closed bool
}

func (c *wsConn) send(data []byte) error {
	c.mu.Lock()
	defer c.mu.Unlock()

	if c.closed {
		return io.ErrClosedPipe
	}

	// Write text frame
	frame := makeFrame(1, data) // 1 = text frame
	_, err := c.conn.Write(frame)
	return err
}

func (c *wsConn) close() {
	c.mu.Lock()
	defer c.mu.Unlock()

	if c.closed {
		return
	}
	c.closed = true

	// Send close frame
	frame := makeFrame(8, nil) // 8 = close frame
	c.conn.Write(frame)
	c.conn.Close()
}

func (c *wsConn) readLoop() {
	buf := make([]byte, 1024)
	for {
		_, err := c.rw.Read(buf)
		if err != nil {
			return
		}
		// We don't process incoming messages, just keep connection alive
	}
}

func makeFrame(opcode byte, data []byte) []byte {
	length := len(data)

	var frame []byte
	if length < 126 {
		frame = make([]byte, 2+length)
		frame[0] = 0x80 | opcode // FIN + opcode
		frame[1] = byte(length)
		copy(frame[2:], data)
	} else if length < 65536 {
		frame = make([]byte, 4+length)
		frame[0] = 0x80 | opcode
		frame[1] = 126
		binary.BigEndian.PutUint16(frame[2:], uint16(length))
		copy(frame[4:], data)
	} else {
		frame = make([]byte, 10+length)
		frame[0] = 0x80 | opcode
		frame[1] = 127
		binary.BigEndian.PutUint64(frame[2:], uint64(length))
		copy(frame[10:], data)
	}

	return frame
}