codec.go

package jsonrpc2

import (
	"bufio"
	"encoding/json"
	"fmt"
	"io"
	"strconv"
	"strings"
	"sync"
)

// Codec reads and writes framed JSON-RPC messages over a byte stream.
//
// The default implementation (StreamCodec) uses json.Decoder / json.Encoder,
// which relies on the JSON syntax itself for framing — no length prefix or
// delimiter is required between messages. This matches the framing used by
// Go's own net/rpc/jsonrpc package.
//
// ContentLengthCodec is a second implementation that uses HTTP-style
// Content-Length headers for framing, matching LSP / DAP / MCP over stdio.
type Codec interface {
	// ReadMessage decodes the next JSON value from the stream into v.
	// v is typically *json.RawMessage so the caller can decide whether the
	// payload is a single request or a batch.
	ReadMessage(v any) error

	// WriteMessage encodes v as one JSON value and flushes it to the stream.
	// It is safe for concurrent callers.
	WriteMessage(v any) error

	// Close closes the underlying stream if it implements io.Closer.
	Close() error
}

// ==== StreamCodec: JSON syntax framing (TCP, generic streams) ============

// StreamCodec is a Codec built on top of a bidirectional byte stream.
type StreamCodec struct {
	rwc io.ReadWriteCloser
	dec *json.Decoder
	enc *json.Encoder

	writeMu sync.Mutex
}

// NewStreamCodec wraps a byte stream with a JSON codec.
func NewStreamCodec(rwc io.ReadWriteCloser) *StreamCodec {
	return &StreamCodec{
		rwc: rwc,
		dec: json.NewDecoder(rwc),
		enc: json.NewEncoder(rwc),
	}
}

// ReadMessage decodes the next JSON value into v.
// It is NOT safe to call ReadMessage concurrently from multiple goroutines
// on the same codec — reads are inherently sequential over a stream.
func (c *StreamCodec) ReadMessage(v any) error {
	return c.dec.Decode(v)
}

// WriteMessage encodes v as a JSON value and writes it to the stream.
// A mutex serializes concurrent writers so that messages do not interleave.
func (c *StreamCodec) WriteMessage(v any) error {
	c.writeMu.Lock()
	defer c.writeMu.Unlock()
	return c.enc.Encode(v)
}

// Close closes the underlying stream.
func (c *StreamCodec) Close() error {
	return c.rwc.Close()
}

// ==== ContentLengthCodec: LSP-style Content-Length framing (stdio) =======

// DefaultMaxMessageSize is the default upper bound on Content-Length codec
// message bodies. Callers can override with (*ContentLengthCodec).MaxMessageSize.
const DefaultMaxMessageSize = 16 << 20 // 16 MiB

// maxHeaderBytes bounds the total header block size to protect against a
// malicious peer that streams headers without ever sending the blank
// separator line.
const maxHeaderBytes = 8 << 10 // 8 KiB

// maxHeaderLines bounds the number of header lines in one frame.
const maxHeaderLines = 100

// ContentLengthCodec frames JSON-RPC messages using an HTTP-like header block:
//
//	Content-Length: <n>\r\n
//	<optional other headers>\r\n
//	\r\n
//	<n bytes of JSON body>
//
// This is the framing used by LSP, DAP, and MCP over stdio. Non-Content-Length
// headers are read and ignored. Only UTF-8 bodies are supported; specifying a
// non-UTF-8 charset in Content-Type is a decode error.
type ContentLengthCodec struct {
	r io.ReadCloser
	w io.WriteCloser

	br *bufio.Reader

	writeMu sync.Mutex

	// MaxMessageSize bounds a single body length. Zero uses DefaultMaxMessageSize.
	MaxMessageSize int
}

// NewContentLengthCodec wraps a reader/writer pair. Typical use is os.Stdin
// and os.Stdout, or the stdin/stdout pipes of an exec.Cmd child process.
//
// If r and w back the same underlying resource (e.g. a *net.TCPConn), pass
// it once as both arguments — Close will be called on each, and idempotent
// Close implementations handle the duplicate call cleanly.
func NewContentLengthCodec(r io.ReadCloser, w io.WriteCloser) *ContentLengthCodec {
	return &ContentLengthCodec{
		r:  r,
		w:  w,
		br: bufio.NewReader(r),
	}
}

// ReadMessage reads one framed message and decodes its JSON body into v.
func (c *ContentLengthCodec) ReadMessage(v any) error {
	var (
		contentLength = -1
		charset       = "utf-8"
		headerBytes   = 0
		headerLines   = 0
	)
	for {
		line, err := c.br.ReadString('\n')
		if err != nil {
			return err
		}
		headerBytes += len(line)
		headerLines++
		if headerBytes > maxHeaderBytes || headerLines > maxHeaderLines {
			return fmt.Errorf("jsonrpc2: header block too large (bytes=%d lines=%d)", headerBytes, headerLines)
		}
		line = strings.TrimRight(line, "\r\n")
		if line == "" {
			break
		}
		colon := strings.IndexByte(line, ':')
		if colon < 0 {
			return fmt.Errorf("jsonrpc2: malformed header line %q", line)
		}
		name := strings.TrimSpace(line[:colon])
		value := strings.TrimSpace(line[colon+1:])
		switch strings.ToLower(name) {
		case "content-length":
			n, err := strconv.Atoi(value)
			if err != nil || n < 0 {
				return fmt.Errorf("jsonrpc2: invalid Content-Length %q", value)
			}
			contentLength = n
		case "content-type":
			// application/vscode-jsonrpc; charset=utf-8
			// RFC 7231 permits LWS after "=", so trim the value after
			// stripping the prefix.
			for _, part := range strings.Split(value, ";") {
				part = strings.TrimSpace(part)
				if strings.HasPrefix(strings.ToLower(part), "charset=") {
					charset = strings.ToLower(strings.TrimSpace(strings.TrimPrefix(part, "charset=")))
				}
			}
		}
	}
	if contentLength < 0 {
		return fmt.Errorf("jsonrpc2: missing Content-Length header")
	}
	limit := c.MaxMessageSize
	if limit <= 0 {
		limit = DefaultMaxMessageSize
	}
	if contentLength > limit {
		return fmt.Errorf("jsonrpc2: message too large: %d > %d", contentLength, limit)
	}
	if charset != "utf-8" && charset != "utf8" {
		return fmt.Errorf("jsonrpc2: unsupported charset %q (only utf-8)", charset)
	}
	body := make([]byte, contentLength)
	if _, err := io.ReadFull(c.br, body); err != nil {
		return err
	}
	return json.Unmarshal(body, v)
}

// WriteMessage encodes v as JSON and writes it with a Content-Length frame.
func (c *ContentLengthCodec) WriteMessage(v any) error {
	body, err := json.Marshal(v)
	if err != nil {
		return err
	}
	c.writeMu.Lock()
	defer c.writeMu.Unlock()
	if _, err := fmt.Fprintf(c.w, "Content-Length: %d\r\n\r\n", len(body)); err != nil {
		return err
	}
	_, err = c.w.Write(body)
	return err
}

// Close closes the reader, then the writer if they refer to different
// underlying objects. When the same value is passed as both r and w (e.g.
// a *net.TCPConn used for both directions), the writer's Close is skipped
// because net.Conn.Close is not idempotent — a second call returns
// "use of closed network connection".
func (c *ContentLengthCodec) Close() error {
	errR := c.r.Close()
	// io.ReadCloser and io.WriteCloser are distinct interface types, so a
	// direct == comparison won't compile. Compare through any().
	if any(c.r) != any(c.w) {
		if errW := c.w.Close(); errW != nil && errR == nil {
			return errW
		}
	}
	return errR
}