Add ability to set custom logger

conn_opt.go

@@ -102,3 +102,10 @@ func LogMessages(logger Logger) ConnOpt {
 		})(c)
 	}
 }
+
+// SetLogger sets the logger for the connection.
+func SetLogger(logger Logger) ConnOpt {
+	return func(c *Conn) {
+		c.logger = logger
+	}
+}

conn_opt_test.go

@@ -0,0 +1,53 @@
+package jsonrpc2_test
+
+import (
+	"bufio"
+	"context"
+	"io"
+	"log"
+	"net"
+	"testing"
+
+	"github.com/sourcegraph/jsonrpc2"
+)
+
+func TestSetLogger(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	rd, wr := io.Pipe()
+	defer rd.Close()
+	defer wr.Close()
+
+	buf := bufio.NewReader(rd)
+	logger := log.New(wr, "", log.Lmsgprefix)
+
+	a, b := net.Pipe()
+	connA := jsonrpc2.NewConn(
+		ctx,
+		jsonrpc2.NewBufferedStream(a, jsonrpc2.VSCodeObjectCodec{}),
+		noopHandler{},
+		jsonrpc2.SetLogger(logger),
+	)
+	connB := jsonrpc2.NewConn(
+		ctx,
+		jsonrpc2.NewBufferedStream(b, jsonrpc2.VSCodeObjectCodec{}),
+		noopHandler{},
+	)
+	defer connA.Close()
+	defer connB.Close()
+
+	// Write a response with no corresponding request.
+	if err := connB.Reply(ctx, jsonrpc2.ID{Num: 0}, nil); err != nil {
+		t.Fatal(err)
+	}
+
+	want := "jsonrpc2: ignoring response #0 with no corresponding request\n"
+	got, err := buf.ReadString('\n')
+	if err != nil {
+		t.Fatal(err)
+	}
+	if got != want {
+		t.Fatalf("got %q, want %q", got, want)
+	}
+}