package jsonrpc2
import (
"context"
"encoding/json"
"fmt"
"sync"
"sync/atomic"
)
type Conn struct {
seq int64 handlers []Handler
stream Stream
err error
pendingMu sync.Mutex pending map[ID]chan *WireResponse
handlingMu sync.Mutex handling map[ID]*Request
}
type requestState int
const (
requestWaiting = requestState(iota)
requestSerial
requestParallel
requestReplied
requestDone
)
type Request struct {
conn *Conn
cancel context.CancelFunc
state requestState
nextRequest chan struct{}
WireRequest
}
func NewErrorf(code int64, format string, args ...interface{}) *Error {
return &Error{
Code: code,
Message: fmt.Sprintf(format, args...),
}
}
func NewConn(s Stream) *Conn {
conn := &Conn{
handlers: []Handler{defaultHandler{}},
stream: s,
pending: make(map[ID]chan *WireResponse),
handling: make(map[ID]*Request),
}
return conn
}
func (c *Conn) AddHandler(handler Handler) {
c.handlers = append([]Handler{handler}, c.handlers...)
}
func (c *Conn) Cancel(id ID) {
c.handlingMu.Lock()
handling, found := c.handling[id]
c.handlingMu.Unlock()
if found {
handling.cancel()
}
}
func (c *Conn) Notify(ctx context.Context, method string, params interface{}) (err error) {
jsonParams, err := marshalToRaw(params)
if err != nil {
return fmt.Errorf("marshalling notify parameters: %v", err)
}
request := &WireRequest{
Method: method,
Params: jsonParams,
}
data, err := json.Marshal(request)
if err != nil {
return fmt.Errorf("marshalling notify request: %v", err)
}
for _, h := range c.handlers {
ctx = h.Request(ctx, c, Send, request)
}
defer func() {
for _, h := range c.handlers {
h.Done(ctx, err)
}
}()
n, err := c.stream.Write(ctx, data)
for _, h := range c.handlers {
ctx = h.Wrote(ctx, n)
}
return err
}
func (c *Conn) Call(ctx context.Context, method string, params, result interface{}) (err error) {
id := ID{Number: atomic.AddInt64(&c.seq, 1)}
jsonParams, err := marshalToRaw(params)
if err != nil {
return fmt.Errorf("marshalling call parameters: %v", err)
}
request := &WireRequest{
ID: &id,
Method: method,
Params: jsonParams,
}
data, err := json.Marshal(request)
if err != nil {
return fmt.Errorf("marshalling call request: %v", err)
}
for _, h := range c.handlers {
ctx = h.Request(ctx, c, Send, request)
}
rchan := make(chan *WireResponse)
c.pendingMu.Lock()
c.pending[id] = rchan
c.pendingMu.Unlock()
defer func() {
c.pendingMu.Lock()
delete(c.pending, id)
c.pendingMu.Unlock()
for _, h := range c.handlers {
h.Done(ctx, err)
}
}()
n, err := c.stream.Write(ctx, data)
for _, h := range c.handlers {
ctx = h.Wrote(ctx, n)
}
if err != nil {
return err
}
select {
case response := <-rchan:
for _, h := range c.handlers {
ctx = h.Response(ctx, c, Receive, response)
}
if response.Error != nil {
return response.Error
}
if result == nil || response.Result == nil {
return nil
}
if err := json.Unmarshal(*response.Result, result); err != nil {
return fmt.Errorf("unmarshalling result: %v", err)
}
return nil
case <-ctx.Done():
cancelled := false
for _, h := range c.handlers {
if h.Cancel(ctx, c, id, cancelled) {
cancelled = true
}
}
return ctx.Err()
}
}
func (r *Request) Conn() *Conn { return r.conn }
func (r *Request) IsNotify() bool {
return r.ID == nil
}
func (r *Request) Parallel() {
if r.state >= requestParallel {
return
}
r.state = requestParallel
close(r.nextRequest)
}
func (r *Request) Reply(ctx context.Context, result interface{}, err error) error {
if r.state >= requestReplied {
return fmt.Errorf("reply invoked more than once")
}
if r.IsNotify() {
return fmt.Errorf("reply not invoked with a valid call")
}
r.Parallel()
r.state = requestReplied
var raw *json.RawMessage
if err == nil {
raw, err = marshalToRaw(result)
}
response := &WireResponse{
Result: raw,
ID: r.ID,
}
if err != nil {
if callErr, ok := err.(*Error); ok {
response.Error = callErr
} else {
response.Error = NewErrorf(0, "%s", err)
}
}
data, err := json.Marshal(response)
if err != nil {
return err
}
for _, h := range r.conn.handlers {
ctx = h.Response(ctx, r.conn, Send, response)
}
n, err := r.conn.stream.Write(ctx, data)
for _, h := range r.conn.handlers {
ctx = h.Wrote(ctx, n)
}
if err != nil {
return err
}
return nil
}
func (c *Conn) setHandling(r *Request, active bool) {
if r.ID == nil {
return
}
r.conn.handlingMu.Lock()
defer r.conn.handlingMu.Unlock()
if active {
r.conn.handling[*r.ID] = r
} else {
delete(r.conn.handling, *r.ID)
}
}
type combined struct {
VersionTag VersionTag `json:"jsonrpc"`
ID *ID `json:"id,omitempty"`
Method string `json:"method"`
Params *json.RawMessage `json:"params,omitempty"`
Result *json.RawMessage `json:"result,omitempty"`
Error *Error `json:"error,omitempty"`
}
func (c *Conn) Run(runCtx context.Context) error {
nextRequest := make(chan struct{})
close(nextRequest)
for {
data, n, err := c.stream.Read(runCtx)
if err != nil {
return err
}
msg := &combined{}
if err := json.Unmarshal(data, msg); err != nil {
for _, h := range c.handlers {
h.Error(runCtx, fmt.Errorf("unmarshal failed: %v", err))
}
continue
}
switch {
case msg.Method != "":
reqCtx, cancelReq := context.WithCancel(runCtx)
thisRequest := nextRequest
nextRequest = make(chan struct{})
req := &Request{
conn: c,
cancel: cancelReq,
nextRequest: nextRequest,
WireRequest: WireRequest{
VersionTag: msg.VersionTag,
Method: msg.Method,
Params: msg.Params,
ID: msg.ID,
},
}
for _, h := range c.handlers {
reqCtx = h.Request(reqCtx, c, Receive, &req.WireRequest)
reqCtx = h.Read(reqCtx, n)
}
c.setHandling(req, true)
go func() {
<-thisRequest
req.state = requestSerial
defer func() {
c.setHandling(req, false)
if !req.IsNotify() && req.state < requestReplied {
req.Reply(reqCtx, nil, NewErrorf(CodeInternalError, "method %q did not reply", req.Method))
}
req.Parallel()
for _, h := range c.handlers {
h.Done(reqCtx, err)
}
cancelReq()
}()
delivered := false
for _, h := range c.handlers {
if h.Deliver(reqCtx, req, delivered) {
delivered = true
}
}
}()
case msg.ID != nil:
c.pendingMu.Lock()
rchan := c.pending[*msg.ID]
if rchan != nil {
delete(c.pending, *msg.ID)
}
c.pendingMu.Unlock()
response := &WireResponse{
Result: msg.Result,
Error: msg.Error,
ID: msg.ID,
}
rchan <- response
close(rchan)
default:
for _, h := range c.handlers {
h.Error(runCtx, fmt.Errorf("message not a call, notify or response, ignoring"))
}
}
}
}
func marshalToRaw(obj interface{}) (*json.RawMessage, error) {
data, err := json.Marshal(obj)
if err != nil {
return nil, err
}
raw := json.RawMessage(data)
return &raw, nil
}