server.go

package jsonrpc2

import (
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"io"
	"net"
	"net/http"
	"os"
	"strings"
	"sync"
)

// Server dispatches JSON-RPC 2.0 requests to a Registry.
//
// A single Server may serve many connections concurrently; each connection
// runs its own dispatch loop but shares the Registry.
type Server struct {
	reg *Registry
	mws []Middleware
}

// NewServer returns a Server with an empty Registry.
func NewServer() *Server {
	return &Server{reg: NewRegistry()}
}

// NewServerWithRegistry returns a Server backed by an existing Registry.
func NewServerWithRegistry(r *Registry) *Server {
	if r == nil {
		r = NewRegistry()
	}
	return &Server{reg: r}
}

// Register forwards to the underlying Registry.
func (s *Server) Register(rcvr any) error { return s.reg.Register(rcvr) }

// RegisterName forwards to the underlying Registry.
func (s *Server) RegisterName(name string, rcvr any) error {
	return s.reg.RegisterName(name, rcvr)
}

// Use appends middlewares to the server's chain. Middlewares registered first
// run outermost. Use is not safe to call concurrently with request handling —
// call it during setup, before serving.
func (s *Server) Use(mws ...Middleware) {
	s.mws = append(s.mws, mws...)
}

// Accept serves connections from ln until ln.Accept returns an error.
func (s *Server) Accept(ln net.Listener) error {
	for {
		conn, err := ln.Accept()
		if err != nil {
			return err
		}
		go s.ServeConn(conn)
	}
}

// ServeConn serves a single connection until it is closed or a fatal error
// occurs on the codec. It wraps the connection in a StreamCodec.
func (s *Server) ServeConn(rwc io.ReadWriteCloser) {
	s.ServeCodec(NewStreamCodec(rwc))
}

// ServeCodec drives the dispatch loop over an arbitrary Codec.
func (s *Server) ServeCodec(codec Codec) {
	defer codec.Close()
	for {
		var raw json.RawMessage
		if err := codec.ReadMessage(&raw); err != nil {
			if errors.Is(err, io.EOF) || errors.Is(err, io.ErrClosedPipe) {
				return
			}
			// Any other error at the frame boundary means we cannot recover
			// the stream — reply with a parse error (id: null) if possible
			// and terminate the connection.
			resp := NewErrorResponse(NullID(), ErrParse(err.Error()))
			_ = codec.WriteMessage(resp)
			return
		}
		go func(raw json.RawMessage) {
			defer safeRecover()
			resp := s.Dispatch(context.Background(), raw)
			if resp != nil {
				_ = codec.WriteMessage(json.RawMessage(resp))
			}
		}(raw)
	}
}

// Dispatch processes one top-level JSON message (single request or batch) and
// returns the raw bytes of the response to send back. A nil return means the
// caller should send no response (Notification, or a batch of all Notifications).
//
// Dispatch is transport-agnostic: HTTP handlers, custom codecs, or in-process
// harnesses can drive it directly without going through ServeCodec.
func (s *Server) Dispatch(ctx context.Context, raw []byte) []byte {
	trimmed := bytes.TrimSpace(raw)
	if len(trimmed) == 0 {
		buf, _ := json.Marshal(NewErrorResponse(NullID(), ErrInvalidRequest("empty message")))
		return buf
	}
	if trimmed[0] == '[' {
		return s.dispatchBatch(ctx, raw)
	}
	resp := s.handleSingle(ctx, raw)
	if resp == nil {
		return nil
	}
	buf, err := json.Marshal(resp)
	if err != nil {
		fallback, _ := json.Marshal(NewErrorResponse(NullID(), ErrInternal(err.Error())))
		return fallback
	}
	return buf
}

// handleSingle dispatches a single request. It returns nil if the request
// was a valid Notification (no response is to be sent).
//
// Flow:
//  1. JSON decode + envelope validate (errors here bypass middleware — we
//     may not even have a valid method name to log).
//  2. Middleware chain wraps the terminal handler (lookup + reflect.Call).
//  3. Notification requests suppress the response even if the chain returns
//     one.
func (s *Server) handleSingle(ctx context.Context, raw json.RawMessage) *Response {
	var req Request
	if err := json.Unmarshal(raw, &req); err != nil {
		return NewErrorResponse(NullID(), ErrParse(err.Error()))
	}
	if verr := req.Validate(); verr != nil {
		if req.ID.IsNotification() {
			return nil
		}
		return NewErrorResponse(req.ID, verr)
	}
	terminal := s.terminalHandler()
	h := chain(s.mws, terminal)
	resp, err := h(ctx, &req)
	if req.ID.IsNotification() {
		return nil
	}
	if err != nil {
		// A middleware returned a non-JSON-RPC error. Wrap it as -32603.
		return NewErrorResponse(req.ID, ErrInternal(err.Error()))
	}
	return resp
}

// terminalHandler returns the innermost Handler in the middleware chain:
// method lookup + reflect.Call + response encoding.
func (s *Server) terminalHandler() Handler {
	return func(ctx context.Context, req *Request) (*Response, error) {
		svc, m := s.reg.Lookup(req.Method)
		if svc == nil || m == nil {
			if req.ID.IsNotification() {
				return nil, nil
			}
			return NewErrorResponse(req.ID, ErrMethodNotFound(req.Method)), nil
		}
		result, jerr := svc.call(ctx, m, req.Params)
		if req.ID.IsNotification() {
			return nil, nil
		}
		if jerr != nil {
			return NewErrorResponse(req.ID, jerr), nil
		}
		return NewResultResponse(req.ID, result), nil
	}
}

// dispatchBatch processes every element of a batch concurrently and returns
// the encoded JSON array of responses (omitting Notification responses).
// Returns nil if the batch is all notifications (transport should send nothing).
//
// Behavior for edge cases:
//   - []       → single -32600 response
//   - all notifications → nil (no response)
func (s *Server) dispatchBatch(ctx context.Context, raw []byte) []byte {
	var elems []json.RawMessage
	if err := json.Unmarshal(raw, &elems); err != nil {
		buf, _ := json.Marshal(NewErrorResponse(NullID(), ErrParse(err.Error())))
		return buf
	}
	if len(elems) == 0 {
		buf, _ := json.Marshal(NewErrorResponse(NullID(), ErrInvalidRequest("empty batch")))
		return buf
	}
	responses := make([]*Response, len(elems))
	var wg sync.WaitGroup
	wg.Add(len(elems))
	for i, e := range elems {
		go func(i int, e json.RawMessage) {
			defer wg.Done()
			defer safeRecover()
			responses[i] = s.handleSingle(ctx, e)
		}(i, e)
	}
	wg.Wait()

	out := make([]*Response, 0, len(responses))
	for _, r := range responses {
		if r != nil {
			out = append(out, r)
		}
	}
	if len(out) == 0 {
		return nil
	}
	buf, err := json.Marshal(out)
	if err != nil {
		fallback, _ := json.Marshal(NewErrorResponse(NullID(), ErrInternal(err.Error())))
		return fallback
	}
	return buf
}

// ==== stdio transport ====================================================

// ServeStdio serves JSON-RPC requests over os.Stdin / os.Stdout using
// Content-Length framing. This is the typical entry point for LSP/MCP-style
// service processes that are spawned by a parent process.
func (s *Server) ServeStdio() {
	s.ServeCodec(NewContentLengthCodec(os.Stdin, os.Stdout))
}

// ==== HTTP transport =====================================================

// HTTPHandler returns an http.Handler that services JSON-RPC 2.0 requests
// over HTTP POST. Users mount it under any path they choose, on their own
// http.ServeMux, gin, echo, or other router.
//
// Behavior:
//   - Only POST is accepted; other methods return 405.
//   - Content-Type must be application/json (or start with it) — otherwise 415.
//   - Request body is capped at DefaultMaxMessageSize (16 MiB); oversized
//     requests return 413 Request Entity Too Large.
//   - Empty body returns a -32700 Parse error inside a 200 response.
//   - A Notification (or an all-Notification batch) returns 204 No Content.
//   - Anything else returns 200 with a JSON body.
//
// HTTP is unidirectional (client → server); server-initiated notifications
// and long-lived subscriptions are NOT supported over this transport. Use the
// TCP or stdio codecs if you need bidirectional communication semantics.
func (s *Server) HTTPHandler() http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.Method != http.MethodPost {
			w.Header().Set("Allow", "POST")
			http.Error(w, "method not allowed", http.StatusMethodNotAllowed)
			return
		}
		ct := r.Header.Get("Content-Type")
		if ct != "" && !strings.HasPrefix(strings.ToLower(ct), "application/json") {
			http.Error(w, "unsupported media type", http.StatusUnsupportedMediaType)
			return
		}
		r.Body = http.MaxBytesReader(w, r.Body, DefaultMaxMessageSize)
		body, err := io.ReadAll(r.Body)
		if err != nil {
			// MaxBytesReader triggers ErrBodyTooLarge; anything else is a
			// mid-stream read failure. Both are best-signaled as a JSON-RPC
			// parse error inside a 200 response, but oversized bodies deserve
			// a 413 to help HTTP-level clients.
			var maxErr *http.MaxBytesError
			if errors.As(err, &maxErr) {
				http.Error(w, "request body too large", http.StatusRequestEntityTooLarge)
				return
			}
			writeHTTPJSON(w, mustMarshalErrResp(err.Error()))
			return
		}
		resp := s.Dispatch(r.Context(), body)
		if resp == nil {
			w.WriteHeader(http.StatusNoContent)
			return
		}
		writeHTTPJSON(w, resp)
	})
}

func writeHTTPJSON(w http.ResponseWriter, body []byte) {
	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(http.StatusOK)
	_, _ = w.Write(body)
}

func mustMarshalErrResp(msg string) []byte {
	b, err := json.Marshal(NewErrorResponse(NullID(), ErrParse(msg)))
	if err != nil {
		b, _ = json.Marshal(NewErrorResponse(NullID(), ErrInternal(err.Error())))
	}
	return b
}

// safeRecover is a last-resort panic handler for dispatch goroutines. Handler
// panics are already recovered inside registry.call, and users can insert
// Recover() middleware for their own middleware panics — this exists to
// guarantee that a panic in server plumbing itself (an unexpected json.Marshal
// failure, a middleware panic without Recover(), etc.) does not take down the
// whole process. The panic is swallowed silently because there's no request
// context left to send a JSON-RPC error to.
func safeRecover() {
	_ = recover()
}