Add subscription support to jsonrpc package

jsonrpc/http.go

@@ -100,7 +100,7 @@ func (h *HTTP) ServeHTTP(writer http.ResponseWriter, req *http.Request) {
 	if conn == nil {
 		return
 	}
-	if err := h.rpc.Handle(conn); err != nil {
+	if err := h.rpc.Handle(req.Context(), conn); err != nil {
 		writer.WriteHeader(http.StatusInternalServerError)
 	}
 }

jsonrpc/memconn.go

@@ -1,29 +1,49 @@
 package jsonrpc
 
-import "io"
+import (
+	"context"
+	"errors"
+	"io"
+)
 
 // MemConn is a simple in-memory JSON RPC connection.
 type MemConn struct {
+	ctx  context.Context
 	send chan<- []byte
 	recv <-chan []byte
 }
 
-func NewMemConn(send chan<- []byte, recv <-chan []byte) *MemConn {
+func NewMemConn(ctx context.Context, send chan<- []byte, recv <-chan []byte) *MemConn {
 	return &MemConn{
+		ctx:  ctx,
 		send: send,
 		recv: recv,
 	}
 }
 
 func (m *MemConn) Read(p []byte) (int, error) {
-	msg, ok := <-m.recv
-	if !ok {
-		return 0, io.EOF
+	select {
+	case <-m.ctx.Done():
+		return 0, ErrRead{
+			Inner: m.ctx.Err(),
+		}
+	case msg, ok := <-m.recv:
+		if !ok {
+			return 0, io.EOF
+		}
+		return copy(p, msg), nil
 	}
-	return copy(p, msg), nil
 }
 
 func (m *MemConn) Write(p []byte) (l int, err error) {
+	defer func() {
+		if panicked := recover(); panicked != nil {
+			l = 0
+			err = errors.New("memconn: write on closed connection")
+		}
+	}()
 	m.send <- p
-	return len(p), nil
+	l = len(p)
+	err = nil
+	return
 }