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 }