From 9e49e6863737005a38f3b98b6589a20fb736bfdb Mon Sep 17 00:00:00 2001
From: lcd1232 <8745863+lcd1232@users.noreply.github.com>
Date: Thu, 1 Apr 2021 14:03:58 +0300
Subject: [PATCH 1/2] Add response and request to context

---
 context.go      | 33 +++++++++++++++++++++++++++++++++
 context_test.go | 25 +++++++++++++++++++++++++
 handler.go      |  6 ++++--
 3 files changed, 62 insertions(+), 2 deletions(-)

diff --git a/context.go b/context.go
index fc09944..9366b7b 100644
--- a/context.go
+++ b/context.go
@@ -2,6 +2,7 @@ package jsonrpc
 
 import (
 	"context"
+	"net/http"
 
 	"github.com/goccy/go-json"
 )
@@ -10,6 +11,8 @@ type (
 	requestIDKey  struct{}
 	metadataIDKey struct{}
 	methodNameKey struct{}
+	requestKey    struct{}
+	responseKey   struct{}
 )
 
 // RequestID takes request id from context.
@@ -41,3 +44,33 @@ func MethodName(c context.Context) string {
 func WithMethodName(c context.Context, name string) context.Context {
 	return context.WithValue(c, methodNameKey{}, name)
 }
+
+// WithRequest adds request to context.
+func WithRequest(c context.Context, r *http.Request) context.Context {
+	return context.WithValue(c, requestKey{}, r)
+}
+
+// GetRequest takes request from context.
+func GetRequest(c context.Context) *http.Request {
+	v := c.Value(requestKey{})
+	if r, ok := v.(*http.Request); ok {
+		return r
+	}
+
+	return nil
+}
+
+// WithRequest adds response to context.
+func WithResponse(c context.Context, r http.ResponseWriter) context.Context {
+	return context.WithValue(c, responseKey{}, r)
+}
+
+// GetResponse takes response from context.
+func GetResponse(c context.Context) http.ResponseWriter {
+	v := c.Value(responseKey{})
+	if r, ok := v.(http.ResponseWriter); ok {
+		return r
+	}
+
+	return nil
+}
diff --git a/context_test.go b/context_test.go
index cb7c887..67a357b 100644
--- a/context_test.go
+++ b/context_test.go
@@ -2,9 +2,12 @@ package jsonrpc
 
 import (
 	"context"
+	"net/http"
+	"net/http/httptest"
 	"testing"
 
 	"github.com/goccy/go-json"
+	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
 
@@ -39,3 +42,25 @@ func TestMethodName(t *testing.T) {
 	})
 	require.Equal(t, t.Name(), pick)
 }
+
+func TestRequest(t *testing.T) {
+	assert.NotPanics(t, func() {
+		r := GetRequest(context.Background())
+		assert.Nil(t, r)
+	})
+	c := context.Background()
+	r := httptest.NewRequest(http.MethodPost, "/", nil)
+	c = WithRequest(c, r)
+	assert.Equal(t, r, GetRequest(c))
+}
+
+func TestResponse(t *testing.T) {
+	assert.NotPanics(t, func() {
+		r := GetResponse(context.Background())
+		assert.Nil(t, r)
+	})
+	c := context.Background()
+	r := httptest.NewRecorder()
+	c = WithResponse(c, r)
+	assert.Equal(t, r, GetResponse(c))
+}
diff --git a/handler.go b/handler.go
index c748e74..c24a077 100644
--- a/handler.go
+++ b/handler.go
@@ -43,7 +43,7 @@ func (mr *MethodRepository) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 
 	resp := make([]*Response, len(rs))
 	for i := range rs {
-		resp[i] = mr.InvokeMethod(r.Context(), rs[i])
+		resp[i] = mr.InvokeMethod(r.Context(), rs[i], r, w)
 	}
 
 	if err := SendResponse(w, resp, batch); err != nil {
@@ -53,7 +53,7 @@ func (mr *MethodRepository) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 }
 
 // InvokeMethod invokes JSON-RPC method.
-func (mr *MethodRepository) InvokeMethod(c context.Context, r *Request) *Response {
+func (mr *MethodRepository) InvokeMethod(c context.Context, r *Request, req *http.Request, w http.ResponseWriter) *Response {
 	var md Metadata
 	res := NewResponse(r)
 	md, res.Error = mr.TakeMethodMetadata(r)
@@ -64,6 +64,8 @@ func (mr *MethodRepository) InvokeMethod(c context.Context, r *Request) *Respons
 	wrappedContext := WithRequestID(c, r.ID)
 	wrappedContext = WithMethodName(wrappedContext, r.Method)
 	wrappedContext = WithMetadata(wrappedContext, md)
+	wrappedContext = WithRequest(wrappedContext, req)
+	wrappedContext = WithResponse(wrappedContext, w)
 	res.Result, res.Error = md.Handler.ServeJSONRPC(wrappedContext, r.Params)
 	if res.Error != nil {
 		res.Result = nil

From cbc8b7349fdebe58d76cfc864cb65fda175efa85 Mon Sep 17 00:00:00 2001
From: lcd1232 <8745863+lcd1232@users.noreply.github.com>
Date: Thu, 1 Apr 2021 15:37:12 +0300
Subject: [PATCH 2/2] Fix comment

---
 context.go | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/context.go b/context.go
index 9366b7b..2f7fee5 100644
--- a/context.go
+++ b/context.go
@@ -60,7 +60,7 @@ func GetRequest(c context.Context) *http.Request {
 	return nil
 }
 
-// WithRequest adds response to context.
+// WithResponse adds response to context.
 func WithResponse(c context.Context, r http.ResponseWriter) context.Context {
 	return context.WithValue(c, responseKey{}, r)
 }