From c9c77b6bb9be655af1a5870e54930aa2fab012c2 Mon Sep 17 00:00:00 2001
From: Sam Herrmann <samherrmann.github@gmail.com>
Date: Fri, 19 Nov 2021 03:30:03 -0500
Subject: [PATCH] Add ability to set custom logger (#48)

Before this commit, a custom logger could be set with the LogMessages
function. However, by using LogMessages not only is a custom logger set
but also all received and sent messages are logged. Use cases exist
where a custom logger is desired to log errors but not all messages.
---
 conn_opt.go      |  7 +++++++
 conn_opt_test.go | 53 ++++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 60 insertions(+)
 create mode 100644 conn_opt_test.go

diff --git a/conn_opt.go b/conn_opt.go
index 3779d7f..423cf80 100644
--- a/conn_opt.go
+++ b/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
+	}
+}
diff --git a/conn_opt_test.go b/conn_opt_test.go
new file mode 100644
index 0000000..df53a1a
--- /dev/null
+++ b/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)
+	}
+}