From a3d86c792f0f5a0c0c2c4ed9157125e914cb5534 Mon Sep 17 00:00:00 2001
From: Keegan Carruthers-Smith <keegan.csmith@gmail.com>
Date: Tue, 1 May 2018 19:02:17 +0100
Subject: [PATCH] allow multiple OnRecv and OnSend

---
 conn_opt.go |  4 ++--
 jsonrpc2.go | 28 ++++++++++++++++++----------
 2 files changed, 20 insertions(+), 12 deletions(-)

diff --git a/conn_opt.go b/conn_opt.go
index e2f8ab5..0a42904 100644
--- a/conn_opt.go
+++ b/conn_opt.go
@@ -13,13 +13,13 @@ 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 }
+	return func(c *Conn) { c.onRecv = append(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 }
+	return func(c *Conn) { c.onSend = append(c.onSend, f) }
 }
 
 // LogMessages causes all messages sent and received on conn to be
diff --git a/jsonrpc2.go b/jsonrpc2.go
index 54f6833..3e0763d 100644
--- a/jsonrpc2.go
+++ b/jsonrpc2.go
@@ -298,8 +298,8 @@ type Conn struct {
 	disconnect chan struct{}
 
 	// Set by ConnOpt funcs.
-	onRecv func(*Request, *Response)
-	onSend func(*Request, *Response)
+	onRecv []func(*Request, *Response)
+	onSend []func(*Request, *Response)
 }
 
 var _ JSONRPC2 = (*Conn)(nil)
@@ -370,12 +370,19 @@ func (c *Conn) send(ctx context.Context, m *anyMessage, wait bool) (cc *call, er
 	}
 	c.mu.Unlock()
 
-	if c.onSend != nil {
+	if len(c.onSend) > 0 {
+		var (
+			req  *Request
+			resp *Response
+		)
 		switch {
 		case m.request != nil:
-			c.onSend(m.request, nil)
+			req = m.request
 		case m.response != nil:
-			c.onSend(nil, m.response)
+			resp = m.response
+		}
+		for _, onSend := range c.onSend {
+			onSend(req, resp)
 		}
 	}
 
@@ -498,8 +505,8 @@ func (c *Conn) readMessages(ctx context.Context) {
 
 		switch {
 		case m.request != nil:
-			if c.onRecv != nil {
-				c.onRecv(m.request, nil)
+			for _, onRecv := range c.onRecv {
+				onRecv(m.request, nil)
 			}
 			c.h.Handle(ctx, c, m.request)
 
@@ -516,13 +523,14 @@ func (c *Conn) readMessages(ctx context.Context) {
 					call.response = resp
 				}
 
-				if c.onRecv != nil {
+				if len(c.onRecv) > 0 {
 					var req *Request
-
 					if call != nil {
 						req = call.request
 					}
-					c.onRecv(req, resp)
+					for _, onRecv := range c.onRecv {
+						onRecv(req, resp)
+					}
 				}
 
 				switch {