server.go

// Package wrapper implements the [ws-wrapper] protocol over any WebSocket
// transport, adding named events, request/response (RPC), named channels, and
// async streaming to a raw WebSocket connection.
//
// # Core types
//
//   - [Server] — accepts inbound WebSocket connections and dispatches events to
//     registered handlers. Register handlers with [Server.On] or [Server.Once];
//     scope them to a named channel with [Server.Of].
//   - [Client] — represents a single WebSocket connection (server-side or
//     standalone). Handlers registered on a [Client] take precedence over the
//     matching [Server] handler for that client.
//   - [AnonymousChannel] — a request-scoped sub-channel returned by a handler
//     to stream events or exchange multiple messages with the requesting client
//     over a single WebSocket connection.
//
// # Event handlers
//
// Handler functions may accept an optional leading [context.Context] argument.
// Use [ClientFromContext] to retrieve the [*Client] that sent the event.
//
// # Transports
//
// This package has no WebSocket library dependency of its own. Wrap any library
// by implementing the [Conn] interface, or use one of the ready-made adapters:
//
//   - github.com/bminer/ws-server-wrapper-go/adapters/coder   ([coder/websocket])
//   - github.com/bminer/ws-server-wrapper-go/adapters/gorilla ([gorilla/websocket])
//
// [ws-wrapper]: https://github.com/bminer/ws-wrapper
// [coder/websocket]: https://github.com/coder/websocket
// [gorilla/websocket]: https://github.com/gorilla/websocket
package wrapper

import (
	"fmt"
	"sync"
)

// Server represents a server that accepts WebSocket connections, handles
// inbound messages, and sends messages to connected clients. Rather than
// listening for connections itself, the Accept method takes any connection that
// implements the Conn interface. This allows the server to be used with any
// WebSocket library. Various adapter libraries are available in the adapters
// subdirectory.
type Server struct {
	ServerChannel  // the "main" server channel with no name
	clientsMu      sync.Mutex
	clients        map[*Client]struct{} // set to nil when server is closed
	handlersMu     sync.Mutex
	handlers       map[handlerName]any
	handlersOnce   map[handlerName]any
	handlerCtxFunc HandlerContextFunc
}

// NewServer creates a new server.
func NewServer() *Server {
	s := &Server{
		clients:      make(map[*Client]struct{}),
		handlers:     make(map[handlerName]any),
		handlersOnce: make(map[handlerName]any),
	}
	// set reference back to client, so channel methods work properly
	s.ServerChannel.server = s
	return s
}

// Accept adds a new client connection to the server. conn only needs to
// implement the Conn interface. If the server has been closed, Accept will
// return an error and close conn.
func (s *Server) Accept(conn Conn) error {
	s.clientsMu.Lock()

	if s.clients == nil {
		s.clientsMu.Unlock()
		conn.Close(StatusGoingAway, "server is closed")
		return fmt.Errorf("server is closed and cannot accept connections")
	}

	// Add client to the set
	client := NewClient(nil)
	client.server = s
	s.clients[client] = struct{}{}
	s.clientsMu.Unlock()

	// Finally, bind the connection to the client
	// This emits the "open" event and starts reading inbound messages
	client.Bind(conn)
	return nil
}

// Close closes the server and all connected clients. Returns the first error
// encountered while closing clients.
func (s *Server) Close() error {
	var clientErr error

	s.clientsMu.Lock()
	clients := s.clients
	s.clients = nil // stop accepting new connections
	s.clientsMu.Unlock()

	// Close all client connections
	for client := range clients {
		if err := client.close(
			StatusGoingAway, "server is closing", false, true,
		); err != nil && clientErr == nil {
			clientErr = err
		}
	}

	return clientErr
}

// Of returns a named channel that scopes events and requests to name. Handlers
// registered on the returned channel only fire for events received on that
// named channel, and events emitted on it are sent only to the named channel
// on the remote end.
func (s *Server) Of(name string) *ServerChannel {
	return &ServerChannel{
		name:   name,
		server: s,
	}
}

// SetHandlerContext sets the handler context function. This function is called
// before every event handler is called and is used to modify the context
// passed to every event handler. It can be used to implement timeouts for
// individual event handlers, for example.
func (s *Server) SetHandlerContext(f HandlerContextFunc) {
	s.handlersMu.Lock()
	s.handlerCtxFunc = f
	s.handlersMu.Unlock()
}

// emitOpen calls the "open" and "connect" event handlers on the main channel
func (s *Server) emitOpen(c *Client) bool {
	return emitReserved(
		func(f any) bool {
			if f, ok := f.(OpenHandler); ok {
				f(c)
				return true
			}
			return false
		},
		&s.handlersMu, s.handlers, s.handlersOnce,
		"open", "connect",
	)
}

// emitError calls the "error" event handler on the main channel
func (s *Server) emitError(c *Client, err error) bool {
	return emitReserved(
		func(f any) bool {
			if f, ok := f.(ErrorHandler); ok {
				f(c, err)
				return true
			}
			return false
		},
		&s.handlersMu, s.handlers, s.handlersOnce,
		"error",
	)
}

// emitClose calls the "close" and "disconnect" event handlers on the main
// channel
func (s *Server) emitClose(
	c *Client, status StatusCode, reason string, userClosed bool,
) bool {
	return emitReserved(
		func(f any) bool {
			switch f := f.(type) {
			case CloseHandler:
				f(c, status, reason, userClosed)
				return true
			case CloseHandlerOld:
				f(c, status, reason)
				return true
			}
			return false
		},
		&s.handlersMu, s.handlers, s.handlersOnce,
		"close", "disconnect",
	)
}