initial commit

.travis.yml

@@ -0,0 +1,7 @@
+language: go
+go: 
+ - release
+ - tip
+
+script:
+ - go test -race -v ./...

LICENSE

@@ -0,0 +1,9 @@
+Copyright (c) 2016 Sourcegraph Inc
+
+MIT License
+
+Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.

README.md

@@ -0,0 +1,11 @@
+# jsonrpc2: JSON-RPC 2.0 implementation for Go [![Build Status](https://travis-ci.org/sourcegraph/jsonrpc2.svg)](https://travis-ci.org/sourcegraph/jsonrpc2)
+
+Package jsonrpc2 provides a [Go](https://golang.org) implementation of [JSON-RPC 2.0](http://www.jsonrpc.org/specification).
+
+This package is **experimental** until further notice.
+
+[**Open the code in Sourcegraph**](https://sourcegraph.com/github.com/sourcegraph/jsonrpc2)
+
+## Known issues
+
+* Batch requests and responses are not yet supported. A handler will panic if it receives a batch request. Because of this, you should not expose any server using this package to external, untrusted traffic (yet).

call_opt.go

@@ -0,0 +1,20 @@
+package jsonrpc2
+
+// CallOption is an option that can be provided to (*Conn).Call to
+// configure custom behavior. See Meta.
+type CallOption interface {
+	apply(r *Request) error
+}
+
+type callOptionFunc func(r *Request) error
+
+func (c callOptionFunc) apply(r *Request) error { return c(r) }
+
+// Meta returns a call option which attaches the given meta object to
+// the JSON-RPC 2.0 request (this is a Sourcegraph extension to JSON
+// RPC 2.0 for carrying metadata).
+func Meta(meta interface{}) CallOption {
+	return callOptionFunc(func(r *Request) error {
+		return r.SetMeta(meta)
+	})
+}

conn_opt.go

@@ -0,0 +1,96 @@
+package jsonrpc2
+
+import (
+	"encoding/json"
+	"log"
+	"sync"
+)
+
+// ConnOpt is the type of function that can be passed to NewConn to
+// customize the Conn before it is created.
+type ConnOpt func(*Conn)
+
+// OnRecv causes all requests received on conn to invoke f(req, nil)
+// and all responses to invoke f(req, resp),
+func OnRecv(f func(*Request, *Response)) ConnOpt {
+	return func(c *Conn) { c.onRecv = f }
+}
+
+// OnSend causes all requests sent on conn to invoke f(req, nil) and
+// all responses to invoke f(nil, resp),
+func OnSend(f func(*Request, *Response)) ConnOpt {
+	return func(c *Conn) { c.onSend = f }
+}
+
+// LogMessages causes all messages sent and received on conn to be
+// logged using the provided logger.
+func LogMessages(log *log.Logger) ConnOpt {
+	return func(c *Conn) {
+		// Remember reqs we have received so we can helpfully show the
+		// request method in OnSend for responses.
+		var (
+			mu         sync.Mutex
+			reqMethods = map[uint64]string{}
+		)
+
+		OnRecv(func(req *Request, resp *Response) {
+			switch {
+			case req != nil && resp == nil:
+				mu.Lock()
+				reqMethods[req.ID] = req.Method
+				mu.Unlock()
+
+				params, _ := json.Marshal(req.Params)
+				if req.Notif {
+					log.Printf("--> notif: %s: %s", req.Method, params)
+				} else {
+					log.Printf("--> request #%d: %s: %s", req.ID, req.Method, params)
+				}
+
+			case resp != nil:
+				var method string
+				if req != nil {
+					method = req.Method
+				} else {
+					method = "(no matching request)"
+				}
+				switch {
+				case resp.Result != nil:
+					result, _ := json.Marshal(resp.Result)
+					log.Printf("--> result #%d: %s: %s", resp.ID, method, result)
+				case resp.Error != nil:
+					err, _ := json.Marshal(resp.Error)
+					log.Printf("--> error #%d: %s: %s", resp.ID, method, err)
+				}
+			}
+		})(c)
+		OnSend(func(req *Request, resp *Response) {
+			switch {
+			case req != nil:
+				params, _ := json.Marshal(req.Params)
+				if req.Notif {
+					log.Printf("<-- notif: %s: %s", req.Method, params)
+				} else {
+					log.Printf("<-- request #%d: %s: %s", req.ID, req.Method, params)
+				}
+
+			case resp != nil:
+				mu.Lock()
+				method := reqMethods[resp.ID]
+				delete(reqMethods, resp.ID)
+				mu.Unlock()
+				if method == "" {
+					method = "(no previous request)"
+				}
+
+				if resp.Result != nil {
+					result, _ := json.Marshal(resp.Result)
+					log.Printf("<-- result #%d: %s: %s", resp.ID, method, result)
+				} else {
+					err, _ := json.Marshal(resp.Error)
+					log.Printf("<-- error #%d: %s: %s", resp.ID, method, err)
+				}
+			}
+		})(c)
+	}
+}

handler_with_error.go

@@ -0,0 +1,42 @@
+package jsonrpc2
+
+import (
+	"context"
+	"log"
+)
+
+// HandlerWithError implements Handler by calling the func for each
+// request and handling returned errors and results.
+type HandlerWithError func(context.Context, *Conn, *Request) (result interface{}, err error)
+
+// Handle implements Handler.
+func (h HandlerWithError) Handle(ctx context.Context, conn *Conn, req *Request) {
+	result, err := h(ctx, conn, req)
+	if req.Notif {
+		if err != nil {
+			log.Printf("jsonrpc2 handler: notification %q handling error: %s", req.Method, err)
+		}
+		return
+	}
+
+	resp := &Response{ID: req.ID}
+	if err == nil {
+		if result == nil {
+			result = struct{}{}
+		}
+		err = resp.SetResult(result)
+	}
+	if err != nil {
+		if e, ok := err.(*Error); ok {
+			resp.Error = e
+		} else {
+			resp.Error = &Error{Message: err.Error()}
+		}
+	}
+
+	if !req.Notif {
+		if err := conn.SendResponse(ctx, resp); err != nil {
+			log.Printf("jsonrpc2 handler: sending response %d: %s", resp.ID, err)
+		}
+	}
+}

jsonrpc2.go

@@ -0,0 +1,744 @@
+// Package jsonrpc2 provides a client and server implementation of
+// [JSON-RPC 2.0](http://www.jsonrpc.org/specification).
+package jsonrpc2
+
+import (
+	"bufio"
+	"bytes"
+	"context"
+	"encoding/json"
+	"errors"
+	"fmt"
+	"io"
+	"log"
+	"net"
+	"strconv"
+	"strings"
+	"sync"
+)
+
+// Request represents a JSON-RPC request or
+// notification. See
+// http://www.jsonrpc.org/specification#request_object and
+// http://www.jsonrpc.org/specification#notification.
+type Request struct {
+	Method string           `json:"method"`
+	Params *json.RawMessage `json:"params,omitempty"`
+	ID     uint64           `json:"id"`
+	Meta   *json.RawMessage `json:"meta,omitempty"`
+	Notif  bool             `json:"-"`
+}
+
+// MarshalJSON implements json.Marshaler and adds the "jsonrpc":"2.0"
+// property.
+func (r *Request) MarshalJSON() ([]byte, error) {
+	if r == nil {
+		return nil, errors.New("can't marshal nil *jsonrpc2.Request")
+	}
+	r2 := struct {
+		Method  string           `json:"method"`
+		Params  *json.RawMessage `json:"params,omitempty"`
+		ID      *uint64          `json:"id,omitempty"`
+		Meta    *json.RawMessage `json:"meta,omitempty"`
+		JSONRPC string           `json:"jsonrpc"`
+	}{
+		Method:  r.Method,
+		Params:  r.Params,
+		Meta:    r.Meta,
+		JSONRPC: "2.0",
+	}
+	if !r.Notif {
+		r2.ID = &r.ID
+	}
+	return json.Marshal(r2)
+}
+
+// UnmarshalJSON implements json.Unmarshaler.
+func (r *Request) UnmarshalJSON(data []byte) error {
+	var r2 struct {
+		Method string           `json:"method"`
+		Params *json.RawMessage `json:"params,omitempty"`
+		Meta   *json.RawMessage `json:"meta,omitempty"`
+		ID     *uint64          `json:"id"`
+	}
+	if err := json.Unmarshal(data, &r2); err != nil {
+		return err
+	}
+	r.Method = r2.Method
+	r.Params = r2.Params
+	r.Meta = r2.Meta
+	if r2.ID == nil {
+		r.ID = 0
+		r.Notif = true
+	} else {
+		r.ID = *r2.ID
+		r.Notif = false
+	}
+	return nil
+}
+
+// SetParams sets r.Params to the JSON representation of v. If JSON
+// marshaling fails, it returns an error.
+func (r *Request) SetParams(v interface{}) error {
+	b, err := json.Marshal(v)
+	if err != nil {
+		return err
+	}
+	r.Params = (*json.RawMessage)(&b)
+	return nil
+}
+
+// SetMeta sets r.Meta to the JSON representation of v. If JSON
+// marshaling fails, it returns an error.
+func (r *Request) SetMeta(v interface{}) error {
+	b, err := json.Marshal(v)
+	if err != nil {
+		return err
+	}
+	r.Meta = (*json.RawMessage)(&b)
+	return nil
+}
+
+// Response represents a JSON-RPC response. See
+// http://www.jsonrpc.org/specification#response_object.
+type Response struct {
+	ID     uint64           `json:"id"`
+	Result *json.RawMessage `json:"result,omitempty"`
+	Error  *Error           `json:"error,omitempty"`
+
+	// SPEC NOTE: The spec says "If there was an error in detecting
+	// the id in the Request object (e.g. Parse error/Invalid
+	// Request), it MUST be Null." If we made the ID field nullable,
+	// then we'd have to make it a pointer type. For simplicity, we're
+	// ignoring the case where there was an error in detecting the ID
+	// in the Request object.
+}
+
+// MarshalJSON implements json.Marshaler and adds the "jsonrpc":"2.0"
+// property.
+func (r *Response) MarshalJSON() ([]byte, error) {
+	if r == nil {
+		return nil, errors.New("can't marshal nil *jsonrpc2.Response")
+	}
+	b, err := json.Marshal(*r)
+	if err != nil {
+		return nil, err
+	}
+	b = append(b[:len(b)-1], []byte(`,"jsonrpc":"2.0"}`)...)
+	return b, nil
+}
+
+// SetResult sets r.Result to the JSON representation of v. If JSON
+// marshaling fails, it returns an error.
+func (r *Response) SetResult(v interface{}) error {
+	b, err := json.Marshal(v)
+	if err != nil {
+		return err
+	}
+	r.Result = (*json.RawMessage)(&b)
+	return nil
+}
+
+// Error represents a JSON-RPC response error.
+type Error struct {
+	Code    int64            `json:"code"`
+	Message string           `json:"message"`
+	Data    *json.RawMessage `json:"data"`
+}
+
+// SetError sets e.Error to the JSON representation of v. If JSON
+// marshaling fails, it panics.
+func (e *Error) SetError(v interface{}) {
+	b, err := json.Marshal(v)
+	if err != nil {
+		panic("Error.SetData: " + err.Error())
+	}
+	e.Data = (*json.RawMessage)(&b)
+}
+
+// Error implements the Go error interface.
+func (e *Error) Error() string {
+	return fmt.Sprintf("jsonrpc2: code %v message: %s", e.Code, e.Message)
+}
+
+const (
+	// Errors defined in the JSON-RPC spec. See
+	// http://www.jsonrpc.org/specification#error_object.
+	CodeParseError       = -32700
+	CodeInvalidRequest   = -32600
+	CodeMethodNotFound   = -32601
+	CodeInvalidParams    = -32602
+	CodeInternalError    = -32603
+	codeServerErrorStart = -32099
+	codeServerErrorEnd   = -32000
+)
+
+// Handler handles JSON-RPC requests and notifications.
+type Handler interface {
+	// Handle is called to handle a request.
+	Handle(context.Context, *Conn, *Request)
+}
+
+// Conn is a JSON-RPC client/server connection. The JSON-RPC protocol
+// is symmetric, so a Conn runs on both ends of a client-server
+// connection.
+type Conn struct {
+	conn io.Closer // all writes should go through w, all reads through readMessages
+	w    *bufio.Writer
+
+	h Handler
+
+	mu       sync.Mutex
+	shutdown bool
+	closing  bool
+	seq      uint64
+	pending  map[uint64]*call
+
+	sending sync.Mutex
+
+	disconnect chan struct{}
+
+	// Set by ConnOpt funcs.
+	onRecv func(*Request, *Response)
+	onSend func(*Request, *Response)
+}
+
+// ErrClosed indicates that the JSON-RPC connection is closed (or in
+// the process of closing).
+var ErrClosed = errors.New("jsonrpc2: connection is closed")
+
+// NewConn creates a new JSON-RPC client/server connection using the
+// given ReadWriteCloser (typically a TCP connection or stdio). The
+// JSON-RPC protocol is symmetric, so a Conn runs on both ends of a
+// client-server connection.
+//
+// NewClient consumes conn, so you should call Close on the returned
+// client not on the given conn.
+func NewConn(ctx context.Context, conn io.ReadWriteCloser, h Handler, opt ...ConnOpt) *Conn {
+	c := &Conn{
+		conn:       conn,
+		w:          bufio.NewWriter(conn),
+		h:          h,
+		pending:    map[uint64]*call{},
+		disconnect: make(chan struct{}),
+	}
+	for _, opt := range opt {
+		opt(c)
+	}
+	go c.readMessages(ctx, bufio.NewReader(conn))
+	return c
+}
+
+// Close closes the JSON-RPC connection. The connection may not be
+// used after it has been closed.
+func (c *Conn) Close() error {
+	c.mu.Lock()
+	if c.shutdown || c.closing {
+		c.mu.Unlock()
+		return ErrClosed
+	}
+	c.closing = true
+	c.mu.Unlock()
+	return c.conn.Close()
+}
+
+func (c *Conn) send(ctx context.Context, m *anyMessage, wait bool) (*call, error) {
+	c.sending.Lock()
+	defer c.sending.Unlock()
+
+	c.mu.Lock()
+	if c.shutdown || c.closing {
+		c.mu.Unlock()
+		return nil, ErrClosed
+	}
+
+	// Store requests so we can later associate them with incoming
+	// responses.
+	var cc *call
+	if m.request != nil && wait {
+		cc = &call{request: m.request, seq: c.seq, done: make(chan error)}
+		c.pending[c.seq] = cc // use first seq as call ID for batch
+		for _, req := range m.requests() {
+			req.ID = c.seq
+			c.seq++
+		}
+	}
+	c.mu.Unlock()
+
+	if c.onSend != nil {
+		switch {
+		case m.request != nil:
+			if m.request.batch != nil {
+				panic("batching not yet implemented")
+			}
+			c.onSend(m.request.single, nil)
+		case m.response != nil:
+			if m.response.batch != nil {
+				panic("batching not yet implemented")
+			}
+			c.onSend(nil, m.response.single)
+		}
+	}
+
+	err := marshalHeadersAndBody(c.w, m)
+	if err != nil {
+		c.w.Flush()
+		if cc != nil {
+			c.mu.Lock()
+			delete(c.pending, cc.seq)
+			c.mu.Unlock()
+		}
+		return nil, err
+	}
+	return cc, c.w.Flush()
+}
+
+// Call initiates a JSON-RPC call using the specified method and
+// params, and waits for the response. If the response is successful,
+// its result is stored in result (a pointer to a value that can be
+// JSON-unmarshaled into); otherwise, a non-nil error is returned.
+func (c *Conn) Call(ctx context.Context, method string, params, result interface{}, opts ...CallOption) error {
+	req := &Request{Method: method}
+	if err := req.SetParams(params); err != nil {
+		return err
+	}
+	for _, opt := range opts {
+		if err := opt.apply(req); err != nil {
+			return err
+		}
+	}
+	call, err := c.send(ctx, &anyMessage{request: &requestOrRequestBatch{single: req}}, true)
+	if err != nil {
+		return err
+	}
+	select {
+	case err, ok := <-call.done:
+		if !ok {
+			err = ErrClosed
+		}
+		if err != nil {
+			return err
+		}
+		if result != nil && call.response.single.Result != nil {
+			// TODO(sqs): error handling
+			if err := json.Unmarshal(*call.response.single.Result, result); err != nil {
+				return err
+			}
+		}
+		return nil
+
+	case <-ctx.Done():
+		return ctx.Err()
+	}
+}
+
+// Notify is like Call, but it returns when the notification request
+// is sent (without waiting for a response, because JSON-RPC
+// notifications do not have responses).
+func (c *Conn) Notify(ctx context.Context, method string, params interface{}, opts ...CallOption) error {
+	req := &Request{Method: method, Notif: true}
+	if err := req.SetParams(params); err != nil {
+		return err
+	}
+	for _, opt := range opts {
+		if err := opt.apply(req); err != nil {
+			return err
+		}
+	}
+	_, err := c.send(ctx, &anyMessage{request: &requestOrRequestBatch{single: req}}, false)
+	return err
+}
+
+// Reply sends a successful response with a result.
+func (c *Conn) Reply(ctx context.Context, id uint64, result interface{}) error {
+	resp := &Response{ID: id}
+	if err := resp.SetResult(result); err != nil {
+		return err
+	}
+	_, err := c.send(ctx, &anyMessage{response: &responseOrResponseBatch{single: resp}}, false)
+	return err
+}
+
+// ReplyWithError sends a response with an error.
+func (c *Conn) ReplyWithError(ctx context.Context, id uint64, respErr *Error) error {
+	_, err := c.send(ctx, &anyMessage{response: &responseOrResponseBatch{single: &Response{ID: id, Error: respErr}}}, false)
+	return err
+}
+
+// SendResponse sends resp to the peer. It is lower level than (*Conn).Reply.
+func (c *Conn) SendResponse(ctx context.Context, resp *Response) error {
+	_, err := c.send(ctx, &anyMessage{response: &responseOrResponseBatch{single: resp}}, false)
+	return err
+}
+
+// DisconnectNotify returns a channel that is closed when the
+// underlying connection is disconnected.
+func (c *Conn) DisconnectNotify() <-chan struct{} {
+	return c.disconnect
+}
+
+func (c *Conn) readMessages(ctx context.Context, r *bufio.Reader) {
+	var err error
+	for err == nil {
+		var m anyMessage
+
+		var n uint32
+		n, err = readHeaderContentLength(r)
+		if err == nil {
+			err = json.NewDecoder(io.LimitReader(r, int64(n))).Decode(&m)
+		}
+		if err != nil {
+			break
+		}
+
+		switch {
+		case m.request != nil:
+			switch {
+			case m.request.batch != nil:
+				panic("batching not yet implemented")
+
+			case m.request.single != nil:
+				if c.onRecv != nil {
+					c.onRecv(m.request.single, nil)
+				}
+				go c.h.Handle(ctx, c, m.request.single)
+			}
+
+		case m.response != nil:
+			resp := *m.response
+			if resp := resp.single; resp != nil {
+				seq := resp.ID
+				c.mu.Lock()
+				call := c.pending[seq]
+				delete(c.pending, seq)
+				c.mu.Unlock()
+
+				if call != nil {
+					call.response = &responseOrResponseBatch{single: resp}
+				}
+
+				if c.onRecv != nil {
+					var req *Request
+
+					if call != nil {
+						if call.request.batch != nil {
+							panic("batching not yet implemented")
+						}
+						req = call.request.single
+					}
+					c.onRecv(req, resp)
+				}
+
+				switch {
+				case call == nil:
+					log.Printf("jsonrpc2: ignoring response %d with no corresponding request", seq)
+
+				case resp.Error != nil:
+					call.done <- resp.Error
+					close(call.done)
+
+				default:
+					call.done <- nil
+					close(call.done)
+				}
+			} else {
+				panic("batches are not yet implemented") // TODO(sqs): support batches
+			}
+		}
+	}
+
+	c.sending.Lock()
+	c.mu.Lock()
+	c.shutdown = true
+	closing := c.closing
+	if err == io.EOF {
+		if closing {
+			err = ErrClosed
+		} else {
+			err = io.ErrUnexpectedEOF
+		}
+	}
+	for _, call := range c.pending {
+		call.done <- err
+		close(call.done)
+	}
+	c.mu.Unlock()
+	c.sending.Unlock()
+	if err != io.ErrUnexpectedEOF && !closing {
+		log.Println("jsonrpc2: protocol error:", err)
+	}
+	close(c.disconnect)
+}
+
+// Server is a JSON-RPC server.
+type Server struct{}
+
+// Serve starts a new JSON-RPC server.
+func Serve(ctx context.Context, lis net.Listener, h Handler, opt ...ConnOpt) error {
+	for {
+		conn, err := lis.Accept()
+		if err != nil {
+			return err
+		}
+		NewConn(ctx, conn, h, opt...)
+	}
+}
+
+// mapRespsToReq returns a slice whose i'th element reports the index
+// in reqs of the i'th response in resps.
+//
+// It returns an error if a response's ID does not refer to that of a
+// request in reqs, or if two responses have the same ID, or if there
+// is a request (non-notification) that does not have a corresponding
+// response.
+func mapRespsToReq(reqs []*Request, resps []*Response) ([]int, error) {
+	reqIndexByID := make(map[uint64]int, len(reqs))
+	for i, req := range reqs {
+		if !req.Notif {
+			reqIndexByID[req.ID] = i
+		}
+	}
+
+	if len(resps) != len(reqIndexByID) {
+		return nil, fmt.Errorf("jsonrpc2: response batch too small: %d responses for %d non-notification requests", len(resps), len(reqIndexByID))
+	}
+
+	m := make([]int, len(resps))
+	seenIDs := make(map[uint64]struct{}, len(resps))
+	for i, resp := range resps {
+		reqIndex, present := reqIndexByID[resp.ID]
+		if !present {
+			return nil, fmt.Errorf("jsonrpc2: response batch contains response with ID %d that doesn't match any IDs in request batch", resp.ID)
+		}
+		m[i] = reqIndex
+
+		if _, seen := seenIDs[resp.ID]; seen {
+			return nil, fmt.Errorf("jsonrpc2: response batch contains multiple responses with same ID %d", resp.ID)
+		}
+		seenIDs[resp.ID] = struct{}{}
+	}
+
+	return m, nil
+}
+
+// call represents a JSON-RPC call over its entire lifecycle.
+type call struct {
+	request  *requestOrRequestBatch
+	response *responseOrResponseBatch
+	seq      uint64 // the seq of the request (or first request for a batch)
+	done     chan error
+}
+
+// anyMessage represents either a JSON Request or Response, or a batch
+// thereof.
+type anyMessage struct {
+	request  *requestOrRequestBatch
+	response *responseOrResponseBatch
+}
+
+func (m *anyMessage) requests() []*Request {
+	if m.request.single != nil {
+		return []*Request{m.request.single}
+	}
+	return m.request.batch
+}
+
+func (m *anyMessage) responses() []*Response {
+	if m.response.single != nil {
+		return []*Response{m.response.single}
+	}
+	return m.response.batch
+}
+
+func (m *anyMessage) MarshalJSON() ([]byte, error) {
+	var v interface{}
+	switch {
+	case m.request != nil && m.response == nil:
+		v = m.request
+	case m.request == nil && m.response != nil:
+		v = m.response
+	}
+	if v != nil {
+		return json.Marshal(v)
+	}
+	return nil, errors.New("jsonrpc2: message (or each message in batch) must have exactly one of the request or response fields set")
+}
+
+func (m *anyMessage) UnmarshalJSON(data []byte) error {
+	// The presence of these fields distinguishes between the 2
+	// message types.
+	type msg struct {
+		Method *string     `json:"method"`
+		Result interface{} `json:"result"`
+		Error  interface{} `json:"error"`
+	}
+
+	var isRequest, isResponse bool
+	checkType := func(m *msg) error {
+		mIsRequest := m.Method != nil
+		mIsResponse := m.Result != nil || m.Error != nil
+		if (!mIsRequest && !mIsResponse) || (mIsRequest && mIsResponse) {
+			return errors.New("jsonrpc2: unable to determine message type (request or response)")
+		}
+		if (mIsRequest && isResponse) || (mIsResponse && isRequest) {
+			return errors.New("jsonrpc2: batch message type mismatch (must be all requests or all responses)")
+		}
+		isRequest = mIsRequest
+		isResponse = mIsResponse
+		return nil
+	}
+
+	if isArray := len(data) > 0 && data[0] == '['; isArray {
+		var msgs []msg
+		if err := json.Unmarshal(data, &msgs); err != nil {
+			return err
+		}
+		if len(msgs) == 0 {
+			return errors.New("jsonrpc2: invalid empty batch")
+		}
+		for _, msg := range msgs {
+			if err := checkType(&msg); err != nil {
+				return err
+			}
+		}
+	} else {
+		var msg msg
+		if err := json.Unmarshal(data, &msg); err != nil {
+			return err
+		}
+		if err := checkType(&msg); err != nil {
+			return err
+		}
+	}
+
+	var v interface{}
+	switch {
+	case isRequest && !isResponse:
+		v = &m.request
+	case !isRequest && isResponse:
+		v = &m.response
+	}
+	return json.Unmarshal(data, v)
+}
+
+type requestOrRequestBatch struct {
+	batch  []*Request
+	single *Request
+}
+
+func (v *requestOrRequestBatch) MarshalJSON() ([]byte, error) {
+	if v.single != nil {
+		return json.Marshal(v.single)
+	}
+	return json.Marshal(v.batch)
+}
+
+func (v *requestOrRequestBatch) UnmarshalJSON(data []byte) error {
+	data = bytes.TrimLeft(data, " \t\n\r")
+	if len(data) == 0 {
+		return errInvalidRequestJSON
+	}
+	switch data[0] {
+	case '[':
+		*v = requestOrRequestBatch{}
+		if err := json.Unmarshal(data, &v.batch); err != nil {
+			return err
+		}
+		return nil
+
+	case '{':
+		*v = requestOrRequestBatch{}
+		if err := json.Unmarshal(data, &v.single); err != nil {
+			return err
+		}
+		return nil
+
+	default:
+		return errInvalidRequestJSON
+	}
+}
+
+type responseOrResponseBatch struct {
+	batch  []*Response
+	single *Response
+}
+
+func (v *responseOrResponseBatch) MarshalJSON() ([]byte, error) {
+	if v.single != nil {
+		return json.Marshal(v.single)
+	}
+	return json.Marshal(v.batch)
+}
+
+func (v *responseOrResponseBatch) UnmarshalJSON(data []byte) error {
+	data = bytes.TrimLeft(data, " \t\n\r")
+	if len(data) == 0 {
+		return errInvalidResponseJSON
+	}
+	switch data[0] {
+	case '[':
+		*v = responseOrResponseBatch{}
+		if err := json.Unmarshal(data, &v.batch); err != nil {
+			return err
+		}
+		return nil
+
+	case '{':
+		*v = responseOrResponseBatch{}
+		if err := json.Unmarshal(data, &v.single); err != nil {
+			return err
+		}
+		return nil
+
+	default:
+		return errInvalidResponseJSON
+	}
+}
+
+func readHeaderContentLength(r *bufio.Reader) (contentLength uint32, err error) {
+	for {
+		line, err := r.ReadString('\r')
+		if err != nil {
+			return 0, err
+		}
+		b, err := r.ReadByte()
+		if err != nil {
+			return 0, err
+		}
+		if b != '\n' {
+			return 0, 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)
+			n, err := strconv.ParseUint(line, 10, 32)
+			if err != nil {
+				return 0, err
+			}
+			contentLength = uint32(n)
+		}
+	}
+	if contentLength == 0 {
+		err = fmt.Errorf("jsonrpc2: no Content-Length header found")
+	}
+	return
+}
+
+func marshalHeadersAndBody(w io.Writer, v interface{}) error {
+	body, err := json.Marshal(v)
+	if err != nil {
+		return err
+	}
+	fmt.Fprintf(w, "Content-Length: %d\r\n", len(body))
+	fmt.Fprint(w, "Content-Type: application/vscode-jsonrpc; charset=utf8\r\n\r\n")
+	_, err = w.Write(body)
+	return err
+}
+
+var (
+	errInvalidRequestJSON  = errors.New("jsonrpc2: request must be either a JSON object or JSON array")
+	errInvalidResponseJSON = errors.New("jsonrpc2: response must be either a JSON object or JSON array")
+)

jsonrpc2_test.go

@@ -0,0 +1,356 @@
+package jsonrpc2
+
+import (
+	"bufio"
+	"bytes"
+	"context"
+	"encoding/json"
+	"fmt"
+	"io"
+	"net"
+	"reflect"
+	"strings"
+	"sync"
+	"testing"
+	"time"
+)
+
+func TestRequest_MarshalJSON_jsonrpc(t *testing.T) {
+	b, err := json.Marshal(&Request{})
+	if err != nil {
+		t.Fatal(err)
+	}
+	if want := `"jsonrpc":"2.0"`; !strings.Contains(string(b), want) {
+		t.Errorf("got %s, want it to include the string %s", b, want)
+	}
+}
+
+func TestResponse_MarshalJSON_jsonrpc(t *testing.T) {
+	b, err := json.Marshal(&Response{})
+	if err != nil {
+		t.Fatal(err)
+	}
+	if want := `"jsonrpc":"2.0"`; !strings.Contains(string(b), want) {
+		t.Errorf("got %s, want it to include the string %s", b, want)
+	}
+}
+
+func TestResponseMarshalJSON_Notif(t *testing.T) {
+	tests := map[*Request]bool{
+		&Request{ID: 0}:       true,
+		&Request{ID: 1}:       true,
+		&Request{Notif: true}: false,
+	}
+	for r, wantIDKey := range tests {
+		b, err := json.Marshal(r)
+		if err != nil {
+			t.Fatal(err)
+		}
+		hasIDKey := bytes.Contains(b, []byte(`"id"`))
+		if hasIDKey != wantIDKey {
+			t.Errorf("got %s, want contain id key: %v", b, wantIDKey)
+		}
+	}
+}
+
+func TestResponseUnmarshalJSON_Notif(t *testing.T) {
+	tests := map[string]bool{
+		`{"method":"f","id":0}`: false,
+		`{"method":"f","id":1}`: false,
+		`{"method":"f"}`:        true,
+	}
+	for s, want := range tests {
+		var r Request
+		if err := json.Unmarshal([]byte(s), &r); err != nil {
+			t.Fatal(err)
+		}
+		if r.Notif != want {
+			t.Errorf("%s: got %v, want %v", s, r.Notif, want)
+		}
+	}
+}
+
+// testHandlerA is the "server" handler.
+type testHandlerA struct{ t *testing.T }
+
+func (h *testHandlerA) Handle(ctx context.Context, conn *Conn, req *Request) {
+	if req.Notif {
+		return // notification
+	}
+	if err := conn.Reply(ctx, req.ID, fmt.Sprintf("hello, #%d: %s", req.ID, *req.Params)); err != nil {
+		h.t.Error(err)
+	}
+
+	if err := conn.Notify(ctx, "m", fmt.Sprintf("notif for #%d", req.ID)); err != nil {
+		h.t.Error(err)
+	}
+}
+
+// testHandlerB is the "client" handler.
+type testHandlerB struct {
+	t   *testing.T
+	mu  sync.Mutex
+	got []string
+}
+
+func (h *testHandlerB) Handle(ctx context.Context, conn *Conn, req *Request) {
+	if req.Notif {
+		h.mu.Lock()
+		defer h.mu.Unlock()
+		h.got = append(h.got, string(*req.Params))
+		return
+	}
+	h.t.Errorf("testHandlerB got unexpected request %+v", req)
+}
+
+func TestClientServer(t *testing.T) {
+	ctx := context.Background()
+	done := make(chan struct{})
+	lis, err := net.Listen("tcp", "127.0.0.1:0") // any available address
+	if err != nil {
+		t.Fatal("Listen:", err)
+	}
+	defer func() {
+		if lis == nil {
+			return // already closed
+		}
+		if err := lis.Close(); err != nil {
+			if !strings.HasSuffix(err.Error(), "use of closed network connection") {
+				t.Fatal(err)
+			}
+		}
+	}()
+
+	ha := testHandlerA{t: t}
+	go func() {
+		if err := Serve(ctx, lis, &ha); err != nil {
+			if !strings.HasSuffix(err.Error(), "use of closed network connection") {
+				t.Error(err)
+			}
+		}
+		close(done)
+	}()
+
+	conn, err := net.Dial("tcp", lis.Addr().String())
+	if err != nil {
+		t.Fatal("Dial:", err)
+	}
+	hb := testHandlerB{t: t}
+	cc := NewConn(ctx, conn, &hb)
+	defer func() {
+		if err := cc.Close(); err != nil {
+			t.Fatal(err)
+		}
+	}()
+
+	// Simple
+	const n = 100
+	for i := 0; i < n; i++ {
+		var got string
+		if err := cc.Call(ctx, "f", []int32{1, 2, 3}, &got); err != nil {
+			t.Fatal(err)
+		}
+		if want := fmt.Sprintf("hello, #%d: [1,2,3]", i); got != want {
+			t.Errorf("got result %q, want %q", got, want)
+		}
+	}
+	time.Sleep(100 * time.Millisecond)
+	hb.mu.Lock()
+	if len(hb.got) != n {
+		t.Errorf("testHandlerB got %d notifications, want %d", len(hb.got), n)
+	}
+	hb.mu.Unlock()
+
+	lis.Close()
+	<-done // ensure Serve's error return (if any) is caught by this test
+}
+
+type noopHandler struct{}
+
+func (noopHandler) Handle(ctx context.Context, conn *Conn, req *Request) {}
+
+type readWriteCloser struct {
+	read, write func(p []byte) (n int, err error)
+}
+
+func (x readWriteCloser) Read(p []byte) (n int, err error) {
+	return x.read(p)
+}
+
+func (x readWriteCloser) Write(p []byte) (n int, err error) {
+	return x.write(p)
+}
+
+func (readWriteCloser) Close() error { return nil }
+
+func eof(p []byte) (n int, err error) {
+	return 0, io.EOF
+}
+
+func TestConn_DisconnectNotify_EOF(t *testing.T) {
+	c := NewConn(context.Background(), &readWriteCloser{eof, eof}, nil)
+	select {
+	case <-c.DisconnectNotify():
+	case <-time.After(200 * time.Millisecond):
+		t.Fatal("no disconnect notification")
+	}
+}
+
+func TestConn_DisconnectNotify_Close(t *testing.T) {
+	c := NewConn(context.Background(), &readWriteCloser{eof, eof}, nil)
+	if err := c.Close(); err != nil {
+		t.Error(err)
+	}
+	select {
+	case <-c.DisconnectNotify():
+	case <-time.After(200 * time.Millisecond):
+		t.Fatal("no disconnect notification")
+	}
+}
+
+func TestConn_DisconnectNotify_Close_async(t *testing.T) {
+	done := make(chan struct{})
+	c := NewConn(context.Background(), &readWriteCloser{eof, eof}, nil)
+	go func() {
+		if err := c.Close(); err != nil && err != ErrClosed {
+			t.Error(err)
+		}
+		close(done)
+	}()
+	select {
+	case <-c.DisconnectNotify():
+	case <-time.After(200 * time.Millisecond):
+		t.Fatal("no disconnect notification")
+	}
+	<-done
+}
+
+func TestConn_Close_waitingForResponse(t *testing.T) {
+	c := NewConn(context.Background(), &readWriteCloser{eof, eof}, noopHandler{})
+	done := make(chan struct{})
+	go func() {
+		if err := c.Call(context.Background(), "m", nil, nil); err != ErrClosed {
+			t.Errorf("got error %v, want %v", err, ErrClosed)
+		}
+		close(done)
+	}()
+	if err := c.Close(); err != nil && err != ErrClosed {
+		t.Error(err)
+	}
+	select {
+	case <-c.DisconnectNotify():
+	case <-time.After(200 * time.Millisecond):
+		t.Fatal("no disconnect notification")
+	}
+	<-done
+}
+
+func TestAnyMessage(t *testing.T) {
+	tests := map[string]struct {
+		request, response bool
+	}{
+		// Single messages
+		`{}`:                                   {},
+		`{"foo":"bar"}`:                        {},
+		`{"method":"m"}`:                       {request: true},
+		`{"result":123}`:                       {response: true},
+		`{"error":{"code":456,"message":"m"}}`: {response: true},
+
+		// Batches
+		`[{"method":"m"}]`:                                      {request: true},
+		`[{"method":"m"},{"foo":"bar"}]`:                        {},
+		`[{"method":"m"},{"result":123}]`:                       {},
+		`[{"result":123},{"method":"foo"}]`:                     {},
+		`[{"result":123}]`:                                      {response: true},
+		`[{"error":{"code":456,"message":"m"}}]`:                {response: true},
+		`[{"result":123},{"error":{"code":456,"message":"m"}}]`: {response: true},
+	}
+	for s, want := range tests {
+		var m anyMessage
+		json.Unmarshal([]byte(s), &m)
+		if (m.request != nil) != want.request {
+			t.Errorf("%s: got request %v, want %v", s, m.request != nil, want.request)
+		}
+		if (m.response != nil) != want.response {
+			t.Errorf("%s: got response %v, want %v", s, m.response != nil, want.response)
+		}
+	}
+}
+
+func TestMessageCodec(t *testing.T) {
+	tests := []struct {
+		v, vempty interface{}
+	}{
+		{
+			v:      &requestOrRequestBatch{single: &Request{ID: 123}},
+			vempty: &requestOrRequestBatch{single: &Request{ID: 123}},
+		},
+		{
+			v:      &responseOrResponseBatch{single: &Response{ID: 123}},
+			vempty: &responseOrResponseBatch{},
+		},
+	}
+	for _, test := range tests {
+		b, err := json.Marshal(test.v)
+		if err != nil {
+			t.Fatal(err)
+		}
+
+		if err := json.Unmarshal(b, test.vempty); err != nil {
+			t.Fatal(err)
+		}
+
+		if !reflect.DeepEqual(test.vempty, test.v) {
+			t.Errorf("got %+v, want %+v", test.vempty, test.v)
+		}
+	}
+}
+
+func TestMapRespsToReq(t *testing.T) {
+	tests := []struct {
+		reqs      []*Request
+		resps     []*Response
+		want      []int
+		wantError bool
+	}{
+		{
+			reqs: nil, resps: nil, want: []int{}, wantError: false,
+		},
+		{
+			reqs: []*Request{{ID: 1}}, resps: []*Response{{ID: 1}}, want: []int{0},
+		},
+		{
+			reqs: []*Request{{ID: 2}}, resps: []*Response{}, wantError: true,
+		},
+		{
+			reqs: []*Request{}, resps: []*Response{{ID: 3}}, wantError: true,
+		},
+		{
+			reqs: []*Request{{ID: 4}}, resps: []*Response{{ID: 4}, {ID: 4}}, wantError: true,
+		},
+	}
+	for _, test := range tests {
+		m, err := mapRespsToReq(test.reqs, test.resps)
+		if (err != nil) != test.wantError {
+			t.Errorf("got error %v, wantError %v", err, test.wantError)
+			continue
+		}
+		if test.wantError {
+			continue
+		}
+		if !reflect.DeepEqual(m, test.want) {
+			t.Errorf("got %v, want %v", m, test.want)
+		}
+	}
+}
+
+func TestReadHeaderContentLength(t *testing.T) {
+	s := "Content-Type: foo\r\nContent-Length: 123\r\n\r\n{}"
+	n, err := readHeaderContentLength(bufio.NewReader(strings.NewReader(s)))
+	if err != nil {
+		t.Fatal(err)
+	}
+	if want := uint32(123); n != want {
+		t.Errorf("got %d, want %d", n, want)
+	}
+}