jsonrpc2/stream.go

package jsonrpc2

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

// An ObjectStream is a bidirectional stream of JSON-RPC 2.0 objects.
type ObjectStream interface {
	// WriteObject writes a JSON-RPC 2.0 object to the stream.
	WriteObject(obj interface{}) error

	// ReadObject reads the next JSON-RPC 2.0 object from the stream
	// and stores it in the value pointed to by v.
	ReadObject(v interface{}) error

	io.Closer
}

// A bufferedObjectStream is an ObjectStream that uses a buffered
// io.ReadWriteCloser to send and receive objects.
type bufferedObjectStream struct {
	conn io.Closer // all writes should go through w, all reads through r
	w    *bufio.Writer
	r    *bufio.Reader

	codec ObjectCodec

	mu sync.Mutex
}

// NewBufferedStream creates a buffered stream from a network
// connection (or other similar interface). The underlying
// objectStream is used to produce the bytes to write to the stream
// for the JSON-RPC 2.0 objects.
func NewBufferedStream(conn io.ReadWriteCloser, codec ObjectCodec) ObjectStream {
	switch v := codec.(type) {
	case PlainObjectCodec:
		v.decoder = json.NewDecoder(conn)
		v.encoder = json.NewEncoder(conn)
		codec = v
	}
	return &bufferedObjectStream{
		conn:  conn,
		w:     bufio.NewWriter(conn),
		r:     bufio.NewReader(conn),
		codec: codec,
	}
}

// WriteObject implements ObjectStream.
func (t *bufferedObjectStream) WriteObject(obj interface{}) error {
	t.mu.Lock()
	defer t.mu.Unlock()
	if err := t.codec.WriteObject(t.w, obj); err != nil {
		return err
	}
	return t.w.Flush()
}

// ReadObject implements ObjectStream.
func (t *bufferedObjectStream) ReadObject(v interface{}) error {
	return t.codec.ReadObject(t.r, v)
}

// Close implements ObjectStream.
func (t *bufferedObjectStream) Close() error {
	return t.conn.Close()
}

// An ObjectCodec specifies how to encode and decode a JSON-RPC 2.0
// object in a stream.
type ObjectCodec interface {
	// WriteObject writes a JSON-RPC 2.0 object to the stream.
	WriteObject(stream io.Writer, obj interface{}) error

	// ReadObject reads the next JSON-RPC 2.0 object from the stream
	// and stores it in the value pointed to by v.
	ReadObject(stream *bufio.Reader, v interface{}) error
}

// VarintObjectCodec reads/writes JSON-RPC 2.0 objects with a varint
// header that encodes the byte length.
type VarintObjectCodec struct{}

// WriteObject implements ObjectCodec.
func (VarintObjectCodec) WriteObject(stream io.Writer, obj interface{}) error {
	data, err := json.Marshal(obj)
	if err != nil {
		return err
	}
	var buf [binary.MaxVarintLen64]byte
	b := binary.PutUvarint(buf[:], uint64(len(data)))
	if _, err := stream.Write(buf[:b]); err != nil {
		return err
	}
	if _, err := stream.Write(data); err != nil {
		return err
	}
	return nil
}

// ReadObject implements ObjectCodec.
func (VarintObjectCodec) ReadObject(stream *bufio.Reader, v interface{}) error {
	b, err := binary.ReadUvarint(stream)
	if err != nil {
		return err
	}
	return json.NewDecoder(io.LimitReader(stream, int64(b))).Decode(v)
}

// VSCodeObjectCodec reads/writes JSON-RPC 2.0 objects with
// Content-Length and Content-Type headers, as specified by
// https://github.com/Microsoft/language-server-protocol/blob/master/protocol.md#base-protocol.
type VSCodeObjectCodec struct{}

// WriteObject implements ObjectCodec.
func (VSCodeObjectCodec) WriteObject(stream io.Writer, obj interface{}) error {
	data, err := json.Marshal(obj)
	if err != nil {
		return err
	}
	if _, err := fmt.Fprintf(stream, "Content-Length: %d\r\n\r\n", len(data)); err != nil {
		return err
	}
	if _, err := stream.Write(data); err != nil {
		return err
	}
	return nil
}

// ReadObject implements ObjectCodec.
func (VSCodeObjectCodec) ReadObject(stream *bufio.Reader, v interface{}) error {
	var contentLength uint64
	for {
		line, err := stream.ReadString('\r')
		if err != nil {
			return err
		}
		b, err := stream.ReadByte()
		if err != nil {
			return err
		}
		if b != '\n' {
			return fmt.Errorf(`jsonrpc2: line endings must be \r\n`)
		}
		if line == "\r" {
			break
		}
		if strings.HasPrefix(line, "Content-Length: ") {
			line = strings.TrimPrefix(line, "Content-Length: ")
			line = strings.TrimSpace(line)
			var err error
			contentLength, err = strconv.ParseUint(line, 10, 32)
			if err != nil {
				return err
			}
		}
	}
	if contentLength == 0 {
		return fmt.Errorf("jsonrpc2: no Content-Length header found")
	}
	return json.NewDecoder(io.LimitReader(stream, int64(contentLength))).Decode(v)
}

// DEPRECATED: use NewPlainObjectStream
// PlainObjectCodec reads/writes plain JSON-RPC 2.0 objects without a header.
type PlainObjectCodec struct {
	decoder *json.Decoder
	encoder *json.Encoder
}

// WriteObject implements ObjectCodec.
func (c PlainObjectCodec) WriteObject(stream io.Writer, v interface{}) error {
	if c.encoder != nil {
		return c.encoder.Encode(v)
	}
	return json.NewEncoder(stream).Encode(v)
}

// ReadObject implements ObjectCodec.
func (c PlainObjectCodec) ReadObject(stream *bufio.Reader, v interface{}) error {
	if c.decoder != nil {
		return c.decoder.Decode(v)
	}
	return json.NewDecoder(stream).Decode(v)
}

type plainObjectStream struct {
	conn    io.Closer
	decoder *json.Decoder
	encoder *json.Encoder
	mu      sync.Mutex
}

// plainObjectStream reads/writes plain JSON-RPC 2.0 objects without a header.
func NewPlainObjectStream(conn io.ReadWriteCloser) ObjectStream {
	os := &plainObjectStream{conn: conn}
	os.encoder = json.NewEncoder(conn)
	os.decoder = json.NewDecoder(conn)
	return os
}

func (os *plainObjectStream) ReadObject(v interface{}) error {
	return os.decoder.Decode(v)
}

func (os *plainObjectStream) WriteObject(v interface{}) error {
	os.mu.Lock()
	defer os.mu.Unlock()
	return os.encoder.Encode(v)
}

func (os *plainObjectStream) Close() error {
	return os.conn.Close()
}