Add a way to specify more non-standard-compliant fields to Request (#50)

This change introduces `ExtraField`, a `CallOption` that can add
arbitrary fields to the top-level JSON-RPC Request message.

call_opt.go

@@ -19,6 +19,15 @@ func Meta(meta interface{}) CallOption {
 	})
 }
 
+// ExtraField returns a call option which attaches the given name/value pair to
+// the JSON-RPC 2.0 request. This can be used to add arbitrary extensions to
+// JSON RPC 2.0.
+func ExtraField(name string, value interface{}) CallOption {
+	return callOptionFunc(func(r *Request) error {
+		return r.SetExtraField(name, value)
+	})
+}
+
 // PickID returns a call option which sets the ID on a request. Care must be
 // taken to ensure there are no conflicts with any previously picked ID, nor
 // with the default sequence ID.

call_opt_test.go

@@ -90,3 +90,53 @@ func TestStringID(t *testing.T) {
 		t.Fatal(err)
 	}
 }
+
+func TestExtraField(t *testing.T) {
+
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	a, b := inMemoryPeerConns()
+	defer a.Close()
+	defer b.Close()
+
+	handler := handlerFunc(func(ctx context.Context, conn *jsonrpc2.Conn, req *jsonrpc2.Request) {
+		replyWithError := func(msg string) {
+			respErr := &jsonrpc2.Error{Code: jsonrpc2.CodeInvalidRequest, Message: msg}
+			if err := conn.ReplyWithError(ctx, req.ID, respErr); err != nil {
+				t.Error(err)
+			}
+		}
+		var sessionID string
+		for _, field := range req.ExtraFields {
+			if field.Name != "sessionId" {
+				continue
+			}
+			var ok bool
+			sessionID, ok = field.Value.(string)
+			if !ok {
+				t.Errorf("\"sessionId\" is not a string: %v", field.Value)
+			}
+		}
+		if sessionID == "" {
+			replyWithError("sessionId must be set")
+			return
+		}
+		if sessionID != "session" {
+			replyWithError("sessionId has the wrong value")
+			return
+		}
+		if err := conn.Reply(ctx, req.ID, "ok"); err != nil {
+			t.Error(err)
+		}
+	})
+	connA := jsonrpc2.NewConn(ctx, jsonrpc2.NewBufferedStream(a, jsonrpc2.VSCodeObjectCodec{}), handler)
+	connB := jsonrpc2.NewConn(ctx, jsonrpc2.NewBufferedStream(b, jsonrpc2.VSCodeObjectCodec{}), noopHandler{})
+	defer connA.Close()
+	defer connB.Close()
+
+	var res string
+	if err := connB.Call(ctx, "f", nil, &res, jsonrpc2.ExtraField("sessionId", "session")); err != nil {
+		t.Fatal(err)
+	}
+}

jsonrpc2.go

@@ -30,6 +30,12 @@ type JSONRPC2 interface {
 	Close() error
 }
 
+// RequestField is a top-level field that can be added to the JSON-RPC request.
+type RequestField struct {
+	Name  string
+	Value interface{}
+}
+
 // Request represents a JSON-RPC request or
 // notification. See
 // http://www.jsonrpc.org/specification#request_object and
@@ -45,61 +51,104 @@ type Request struct {
 	// NOTE: It is not part of spec. However, it is useful for propogating
 	// tracing context, etc.
 	Meta *json.RawMessage `json:"meta,omitempty"`
+
+	// ExtraFields optionally adds fields to the root of the JSON-RPC request.
+	//
+	// NOTE: It is not part of the spec, but there are other protocols based on
+	// JSON-RPC 2 that require it.
+	ExtraFields []RequestField `json:"-"`
 }
 
 // MarshalJSON implements json.Marshaler and adds the "jsonrpc":"2.0"
 // property.
 func (r Request) MarshalJSON() ([]byte, error) {
-	r2 := struct {
-		Method  string           `json:"method"`
-		Params  *json.RawMessage `json:"params,omitempty"`
-		ID      *ID              `json:"id,omitempty"`
-		Meta    *json.RawMessage `json:"meta,omitempty"`
-		JSONRPC string           `json:"jsonrpc"`
-	}{
-		Method:  r.Method,
-		Params:  r.Params,
-		Meta:    r.Meta,
-		JSONRPC: "2.0",
+	r2 := map[string]interface{}{
+		"jsonrpc": "2.0",
+		"method":  r.Method,
+	}
+	for _, field := range r.ExtraFields {
+		r2[field.Name] = field.Value
 	}
 	if !r.Notif {
-		r2.ID = &r.ID
+		r2["id"] = &r.ID
+	}
+	if r.Params != nil {
+		r2["params"] = r.Params
+	}
+	if r.Meta != nil {
+		r2["meta"] = r.Meta
 	}
 	return json.Marshal(r2)
 }
 
 // UnmarshalJSON implements json.Unmarshaler.
 func (r *Request) UnmarshalJSON(data []byte) error {
-	var r2 struct {
-		Method string           `json:"method"`
-		Params *json.RawMessage `json:"params,omitempty"`
-		Meta   *json.RawMessage `json:"meta,omitempty"`
-		ID     *ID              `json:"id"`
-	}
+	r2 := make(map[string]interface{})
 
 	// Detect if the "params" field is JSON "null" or just not present
 	// by seeing if the field gets overwritten to nil.
-	r2.Params = &json.RawMessage{}
+	emptyParams := &json.RawMessage{}
+	r2["params"] = emptyParams
 
-	if err := json.Unmarshal(data, &r2); err != nil {
+	decoder := json.NewDecoder(bytes.NewReader(data))
+	decoder.UseNumber()
+	if err := decoder.Decode(&r2); err != nil {
 		return err
 	}
-	r.Method = r2.Method
+	var ok bool
+	r.Method, ok = r2["method"].(string)
+	if !ok {
+		return errors.New("missing method field")
+	}
 	switch {
-	case r2.Params == nil:
+	case r2["params"] == nil:
 		r.Params = &jsonNull
-	case len(*r2.Params) == 0:
+	case r2["params"] == emptyParams:
 		r.Params = nil
 	default:
-		r.Params = r2.Params
+		b, err := json.Marshal(r2["params"])
+		if err != nil {
+			return fmt.Errorf("failed to marshal params: %w", err)
+		}
+		r.Params = (*json.RawMessage)(&b)
+	}
+	meta, ok := r2["meta"]
+	if ok {
+		b, err := json.Marshal(meta)
+		if err != nil {
+			return fmt.Errorf("failed to marshal Meta: %w", err)
+		}
+		r.Meta = (*json.RawMessage)(&b)
 	}
-	r.Meta = r2.Meta
-	if r2.ID == nil {
+	switch rawID := r2["id"].(type) {
+	case nil:
 		r.ID = ID{}
 		r.Notif = true
-	} else {
-		r.ID = *r2.ID
+	case string:
+		r.ID = ID{Str: rawID, IsString: true}
+		r.Notif = false
+	case json.Number:
+		id, err := rawID.Int64()
+		if err != nil {
+			return fmt.Errorf("failed to unmarshal ID: %w", err)
+		}
+		r.ID = ID{Num: uint64(id)}
 		r.Notif = false
+	default:
+		return fmt.Errorf("unexpected ID type: %T", rawID)
+	}
+
+	// Clear the extra fields before populating them again.
+	r.ExtraFields = nil
+	for name, value := range r2 {
+		switch name {
+		case "id", "jsonrpc", "meta", "method", "params":
+			continue
+		}
+		r.ExtraFields = append(r.ExtraFields, RequestField{
+			Name:  name,
+			Value: value,
+		})
 	}
 	return nil
 }
@@ -126,6 +175,21 @@ func (r *Request) SetMeta(v interface{}) error {
 	return nil
 }
 
+// SetExtraField adds an entry to r.ExtraFields, so that it is added to the
+// JSON representation of the request, as a way to add arbitrary extensions to
+// JSON RPC 2.0. If JSON marshaling fails, it returns an error.
+func (r *Request) SetExtraField(name string, v interface{}) error {
+	switch name {
+	case "id", "jsonrpc", "meta", "method", "params":
+		return fmt.Errorf("invalid extra field %q", name)
+	}
+	r.ExtraFields = append(r.ExtraFields, RequestField{
+		Name:  name,
+		Value: v,
+	})
+	return nil
+}
+
 // Response represents a JSON-RPC response. See
 // http://www.jsonrpc.org/specification#response_object.
 type Response struct {

jsonrpc2_test.go

@@ -24,7 +24,7 @@ func TestRequest_MarshalJSON_jsonrpc(t *testing.T) {
 	if err != nil {
 		t.Fatal(err)
 	}
-	if want := `{"method":"","id":0,"jsonrpc":"2.0"}`; string(b) != want {
+	if want := `{"id":0,"jsonrpc":"2.0","method":""}`; string(b) != want {
 		t.Errorf("got %q, want %q", b, want)
 	}
 }

object_test.go

@@ -44,17 +44,21 @@ func TestRequest_MarshalUnmarshalJSON(t *testing.T) {
 		want Request
 	}{
 		{
-			data: []byte(`{"method":"m","params":{"foo":"bar"},"id":123,"jsonrpc":"2.0"}`),
+			data: []byte(`{"id":123,"jsonrpc":"2.0","method":"m","params":{"foo":"bar"}}`),
 			want: Request{ID: ID{Num: 123}, Method: "m", Params: &obj},
 		},
 		{
-			data: []byte(`{"method":"m","params":null,"id":123,"jsonrpc":"2.0"}`),
+			data: []byte(`{"id":123,"jsonrpc":"2.0","method":"m","params":null}`),
 			want: Request{ID: ID{Num: 123}, Method: "m", Params: &null},
 		},
 		{
-			data: []byte(`{"method":"m","id":123,"jsonrpc":"2.0"}`),
+			data: []byte(`{"id":123,"jsonrpc":"2.0","method":"m"}`),
 			want: Request{ID: ID{Num: 123}, Method: "m", Params: nil},
 		},
+		{
+			data: []byte(`{"id":123,"jsonrpc":"2.0","method":"m","sessionId":"session"}`),
+			want: Request{ID: ID{Num: 123}, Method: "m", Params: nil, ExtraFields: []RequestField{{Name: "sessionId", Value: "session"}}},
+		},
 	}
 	for _, test := range tests {
 		var got Request