| 1 | package rpcwire |
| 2 | |
| 3 | import ( |
| 4 | "bytes" |
| 5 | "context" |
| 6 | "encoding/json" |
| 7 | "errors" |
| 8 | "fmt" |
| 9 | "io" |
| 10 | "runtime" |
| 11 | "strings" |
| 12 | "sync" |
| 13 | "sync/atomic" |
| 14 | "testing" |
| 15 | "time" |
| 16 | ) |
| 17 | |
| 18 | type strictTestResponse struct { |
| 19 | JSONRPC string `json:"jsonrpc"` |
| 20 | ID json.RawMessage `json:"id"` |
| 21 | Result json.RawMessage `json:"result"` |
| 22 | Error *ErrorObject `json:"error"` |
| 23 | } |
| 24 | |
| 25 | func decodeStrictTestResponses(t *testing.T, raw []byte) []strictTestResponse { |
| 26 | t.Helper() |
| 27 | dec := json.NewDecoder(bytes.NewReader(raw)) |
| 28 | var frames []strictTestResponse |
| 29 | for { |
| 30 | var frame strictTestResponse |
| 31 | if err := dec.Decode(&frame); errors.Is(err, io.EOF) { |
| 32 | return frames |
| 33 | } else if err != nil { |
| 34 | t.Fatalf("decode response %d from %q: %v", len(frames), raw, err) |
| 35 | } |
| 36 | frames = append(frames, frame) |
| 37 | } |
| 38 | } |
| 39 | |
| 40 | func TestEOFDelimitedFinalFrameIsDispatched(t *testing.T) { |
| 41 | t.Run("complete request without newline", func(t *testing.T) { |
| 42 | var out bytes.Buffer |
| 43 | conn := NewConn(strings.NewReader(`{"jsonrpc":"2.0","id":1,"method":"echo","params":{"v":1}}`), &out, Options{StrictJSONRPC: true}) |
| 44 | conn.Handle("echo", func(_ context.Context, params json.RawMessage) (any, error) { |
| 45 | return json.RawMessage(params), nil |
| 46 | }) |
| 47 | if err := conn.Serve(context.Background()); err != nil { |
| 48 | t.Fatalf("Serve: %v", err) |
| 49 | } |
| 50 | frames := decodeStrictTestResponses(t, out.Bytes()) |
| 51 | if len(frames) != 1 || frames[0].Error != nil || string(frames[0].Result) != `{"v":1}` { |
| 52 | t.Fatalf("frames = %+v, raw = %q", frames, out.String()) |
| 53 | } |
| 54 | }) |
| 55 | |
| 56 | t.Run("truncated JSON without newline", func(t *testing.T) { |
| 57 | var out bytes.Buffer |
| 58 | conn := NewConn(strings.NewReader(`{"jsonrpc":"2.0","id":1,"method":"echo"`), &out, Options{StrictJSONRPC: true}) |
| 59 | if err := conn.Serve(context.Background()); err != nil { |
| 60 | t.Fatalf("Serve: %v", err) |
| 61 | } |
| 62 | frames := decodeStrictTestResponses(t, out.Bytes()) |
| 63 | if len(frames) != 1 || frames[0].Error == nil || frames[0].Error.Code != ErrParse || string(frames[0].ID) != "null" { |
| 64 | t.Fatalf("frames = %+v, raw = %q", frames, out.String()) |
| 65 | } |
| 66 | }) |
| 67 | } |
| 68 | |
| 69 | func TestInboundLimitAcceptsFrameAtExactBoundary(t *testing.T) { |
| 70 | frame := "{\"jsonrpc\":\"2.0\",\"method\":\"note\",\"params\":{}}\n" |
| 71 | var called atomic.Int32 |
| 72 | conn := NewConn(strings.NewReader(frame), io.Discard, Options{MaxInboundBytes: len(frame), StrictJSONRPC: true}) |
| 73 | conn.HandleNotify("note", func(context.Context, json.RawMessage) { called.Add(1) }) |
| 74 | if err := conn.Serve(context.Background()); err != nil { |
| 75 | t.Fatalf("Serve at exact boundary: %v", err) |
| 76 | } |
| 77 | if got := called.Load(); got != 1 { |
| 78 | t.Fatalf("notification calls = %d, want 1", got) |
| 79 | } |
| 80 | } |
| 81 | |
| 82 | type busyResponseWriter struct { |
| 83 | mu sync.Mutex |
| 84 | buf bytes.Buffer |
| 85 | busy chan struct{} |
| 86 | } |
| 87 | |
| 88 | func (w *busyResponseWriter) Write(p []byte) (int, error) { |
| 89 | w.mu.Lock() |
| 90 | defer w.mu.Unlock() |
| 91 | n, err := w.buf.Write(p) |
| 92 | if bytes.Contains(p, []byte(`"message":"server busy"`)) { |
| 93 | select { |
| 94 | case w.busy <- struct{}{}: |
| 95 | default: |
| 96 | } |
| 97 | } |
| 98 | return n, err |
| 99 | } |
| 100 | |
| 101 | func (w *busyResponseWriter) Bytes() []byte { |
| 102 | w.mu.Lock() |
| 103 | defer w.mu.Unlock() |
| 104 | return append([]byte(nil), w.buf.Bytes()...) |
| 105 | } |
| 106 | |
| 107 | func TestInboundHandlerConcurrencyIsBoundedWithoutBlockingResponses(t *testing.T) { |
| 108 | var input strings.Builder |
| 109 | for id := 1; id <= 5; id++ { |
| 110 | fmt.Fprintf(&input, "{\"jsonrpc\":\"2.0\",\"id\":%d,\"method\":\"block\",\"params\":{}}\n", id) |
| 111 | } |
| 112 | out := &busyResponseWriter{busy: make(chan struct{}, 3)} |
| 113 | release := make(chan struct{}) |
| 114 | started := make(chan struct{}, 2) |
| 115 | var running atomic.Int32 |
| 116 | var maximum atomic.Int32 |
| 117 | conn := NewConn(strings.NewReader(input.String()), out, Options{ |
| 118 | StrictJSONRPC: true, MaxConcurrentHandlers: 2, |
| 119 | }) |
| 120 | conn.Handle("block", func(context.Context, json.RawMessage) (any, error) { |
| 121 | current := running.Add(1) |
| 122 | for { |
| 123 | observed := maximum.Load() |
| 124 | if current <= observed || maximum.CompareAndSwap(observed, current) { |
| 125 | break |
| 126 | } |
| 127 | } |
| 128 | started <- struct{}{} |
| 129 | <-release |
| 130 | running.Add(-1) |
| 131 | return struct{}{}, nil |
| 132 | }) |
| 133 | done := make(chan error, 1) |
| 134 | go func() { done <- conn.Serve(context.Background()) }() |
| 135 | for i := 0; i < 2; i++ { |
| 136 | select { |
| 137 | case <-started: |
| 138 | case <-time.After(time.Second): |
| 139 | t.Fatal("bounded handlers did not start") |
| 140 | } |
| 141 | } |
| 142 | for i := 0; i < 3; i++ { |
| 143 | select { |
| 144 | case <-out.busy: |
| 145 | case <-time.After(time.Second): |
| 146 | t.Fatal("overload responses were blocked behind active handlers") |
| 147 | } |
| 148 | } |
| 149 | close(release) |
| 150 | if err := <-done; err != nil { |
| 151 | t.Fatal(err) |
| 152 | } |
| 153 | if got := maximum.Load(); got != 2 { |
| 154 | t.Fatalf("maximum concurrent handlers = %d, want 2", got) |
| 155 | } |
| 156 | raw := out.Bytes() |
| 157 | frames := decodeStrictTestResponses(t, raw) |
| 158 | busy := 0 |
| 159 | for _, frame := range frames { |
| 160 | if frame.Error != nil && frame.Error.Code == ErrServerBusy && frame.Error.Message == "server busy" { |
| 161 | busy++ |
| 162 | } |
| 163 | } |
| 164 | if len(frames) != 5 || busy != 3 { |
| 165 | t.Fatalf("responses=%d busy=%d, want 5/3; raw=%q", len(frames), busy, raw) |
| 166 | } |
| 167 | } |
| 168 | |
| 169 | func TestQueuedNotificationsPreserveBurstOrder(t *testing.T) { |
| 170 | const count = 500 |
| 171 | var input strings.Builder |
| 172 | for i := 0; i < count; i++ { |
| 173 | fmt.Fprintf(&input, "{\"jsonrpc\":\"2.0\",\"method\":\"note\",\"params\":{\"index\":%d}}\n", i) |
| 174 | } |
| 175 | conn := NewConn(strings.NewReader(input.String()), io.Discard, Options{ |
| 176 | Name: "ordered-notifications", StrictJSONRPC: true, MaxQueuedNotifications: count, |
| 177 | }) |
| 178 | got := make([]int, 0, count) |
| 179 | conn.HandleNotify("note", func(_ context.Context, params json.RawMessage) { |
| 180 | var value struct { |
| 181 | Index int `json:"index"` |
| 182 | } |
| 183 | if err := json.Unmarshal(params, &value); err != nil { |
| 184 | t.Errorf("decode notification: %v", err) |
| 185 | return |
| 186 | } |
| 187 | got = append(got, value.Index) |
| 188 | }) |
| 189 | if err := conn.Serve(context.Background()); err != nil { |
| 190 | t.Fatalf("Serve: %v", err) |
| 191 | } |
| 192 | if len(got) != count { |
| 193 | t.Fatalf("notification calls = %d, want %d", len(got), count) |
| 194 | } |
| 195 | for i, value := range got { |
| 196 | if value != i { |
| 197 | t.Fatalf("notification[%d] = %d, want %d", i, value, i) |
| 198 | } |
| 199 | } |
| 200 | } |
| 201 | |
| 202 | func TestQueuedNotificationOverflowFailsConnection(t *testing.T) { |
| 203 | var input strings.Builder |
| 204 | for i := 0; i < 100; i++ { |
| 205 | fmt.Fprintf(&input, "{\"jsonrpc\":\"2.0\",\"method\":\"note\",\"params\":{\"index\":%d}}\n", i) |
| 206 | } |
| 207 | conn := NewConn(strings.NewReader(input.String()), io.Discard, Options{ |
| 208 | Name: "notification-overflow", StrictJSONRPC: true, MaxQueuedNotifications: 1, |
| 209 | }) |
| 210 | conn.HandleNotify("note", func(ctx context.Context, _ json.RawMessage) { |
| 211 | <-ctx.Done() |
| 212 | }) |
| 213 | err := conn.Serve(context.Background()) |
| 214 | if err == nil || !strings.Contains(err.Error(), "notification-overflow: notification queue overflow") { |
| 215 | t.Fatalf("Serve error = %v, want notification queue overflow", err) |
| 216 | } |
| 217 | } |
| 218 | |
| 219 | func TestStrictJSONRPCAcceptsLegalRequestNotificationAndResponses(t *testing.T) { |
| 220 | t.Run("request and notification", func(t *testing.T) { |
| 221 | input := strings.Join([]string{ |
| 222 | `{"jsonrpc":"2.0","id":1,"method":"sum","params":[2,3]}`, |
| 223 | `{"jsonrpc":"2.0","method":"note","params":{"ok":true}}`, |
| 224 | }, "\n") + "\n" |
| 225 | var out bytes.Buffer |
| 226 | conn := NewConn(strings.NewReader(input), &out, Options{StrictJSONRPC: true}) |
| 227 | var requestCalls atomic.Int32 |
| 228 | var notificationCalls atomic.Int32 |
| 229 | conn.Handle("sum", func(_ context.Context, params json.RawMessage) (any, error) { |
| 230 | requestCalls.Add(1) |
| 231 | if string(params) != `[2,3]` { |
| 232 | t.Errorf("request params = %s", params) |
| 233 | } |
| 234 | return map[string]int{"sum": 5}, nil |
| 235 | }) |
| 236 | conn.HandleNotify("note", func(_ context.Context, params json.RawMessage) { |
| 237 | notificationCalls.Add(1) |
| 238 | if string(params) != `{"ok":true}` { |
| 239 | t.Errorf("notification params = %s", params) |
| 240 | } |
| 241 | }) |
| 242 | if err := conn.Serve(context.Background()); err != nil { |
| 243 | t.Fatalf("Serve: %v", err) |
| 244 | } |
| 245 | if requestCalls.Load() != 1 || notificationCalls.Load() != 1 { |
| 246 | t.Fatalf("request calls = %d, notification calls = %d", requestCalls.Load(), notificationCalls.Load()) |
| 247 | } |
| 248 | frames := decodeStrictTestResponses(t, out.Bytes()) |
| 249 | if len(frames) != 1 || frames[0].Error != nil || string(frames[0].Result) != `{"sum":5}` { |
| 250 | t.Fatalf("frames = %+v, raw = %q", frames, out.String()) |
| 251 | } |
| 252 | }) |
| 253 | |
| 254 | for _, tt := range []struct { |
| 255 | name string |
| 256 | frame string |
| 257 | want string |
| 258 | wantError int |
| 259 | }{ |
| 260 | {name: "result response", frame: `{"jsonrpc":"2.0","id":9,"result":{"ok":true}}`, want: `{"ok":true}`}, |
| 261 | {name: "error response", frame: `{"jsonrpc":"2.0","id":9,"error":{"code":-32007,"message":"busy","data":{"retry":true}}}`, wantError: -32007}, |
| 262 | } { |
| 263 | t.Run(tt.name, func(t *testing.T) { |
| 264 | conn := NewConn(strings.NewReader(tt.frame), io.Discard, Options{StrictJSONRPC: true}) |
| 265 | ch := make(chan rpcResult, 1) |
| 266 | conn.pending[9] = ch |
| 267 | if err := conn.Serve(context.Background()); err != nil { |
| 268 | t.Fatalf("Serve: %v", err) |
| 269 | } |
| 270 | result := <-ch |
| 271 | if tt.wantError != 0 { |
| 272 | var responseErr *ResponseError |
| 273 | if !errors.As(result.err, &responseErr) || responseErr.Code != tt.wantError || !bytes.Contains(responseErr.Data, []byte(`"retry":true`)) { |
| 274 | t.Fatalf("response error = %#v", result.err) |
| 275 | } |
| 276 | return |
| 277 | } |
| 278 | if result.err != nil || string(result.result) != tt.want { |
| 279 | t.Fatalf("result = %s, error = %v", result.result, result.err) |
| 280 | } |
| 281 | }) |
| 282 | } |
| 283 | } |
| 284 | |
| 285 | func TestStrictJSONRPCRejectsResultAndErrorBadVersionAndScalarParams(t *testing.T) { |
| 286 | tests := []struct { |
| 287 | name string |
| 288 | frame string |
| 289 | }{ |
| 290 | { |
| 291 | name: "response has result and error", |
| 292 | frame: `{"jsonrpc":"2.0","id":1,"result":{},"error":{"code":-1,"message":"bad"}}`, |
| 293 | }, |
| 294 | { |
| 295 | name: "wrong jsonrpc version", |
| 296 | frame: `{"jsonrpc":"1.0","id":1,"method":"run","params":{}}`, |
| 297 | }, |
| 298 | { |
| 299 | name: "string params", |
| 300 | frame: `{"jsonrpc":"2.0","id":1,"method":"run","params":"bad"}`, |
| 301 | }, |
| 302 | { |
| 303 | name: "number params", |
| 304 | frame: `{"jsonrpc":"2.0","id":1,"method":"run","params":7}`, |
| 305 | }, |
| 306 | { |
| 307 | name: "null params", |
| 308 | frame: `{"jsonrpc":"2.0","id":1,"method":"run","params":null}`, |
| 309 | }, |
| 310 | } |
| 311 | |
| 312 | for _, tt := range tests { |
| 313 | t.Run(tt.name, func(t *testing.T) { |
| 314 | var out bytes.Buffer |
| 315 | conn := NewConn(strings.NewReader(tt.frame+"\n"), &out, Options{StrictJSONRPC: true}) |
| 316 | var called atomic.Int32 |
| 317 | conn.Handle("run", func(context.Context, json.RawMessage) (any, error) { |
| 318 | called.Add(1) |
| 319 | return nil, nil |
| 320 | }) |
| 321 | if err := conn.Serve(context.Background()); err != nil { |
| 322 | t.Fatalf("Serve: %v", err) |
| 323 | } |
| 324 | if called.Load() != 0 { |
| 325 | t.Fatalf("handler called %d times", called.Load()) |
| 326 | } |
| 327 | frames := decodeStrictTestResponses(t, out.Bytes()) |
| 328 | if len(frames) != 1 || frames[0].Error == nil || frames[0].Error.Code != ErrInvalidRequest || string(frames[0].ID) != "1" { |
| 329 | t.Fatalf("frames = %+v, raw = %q", frames, out.String()) |
| 330 | } |
| 331 | }) |
| 332 | } |
| 333 | } |
| 334 | |
| 335 | type oneByteWriter struct { |
| 336 | mu sync.Mutex |
| 337 | buf bytes.Buffer |
| 338 | } |
| 339 | |
| 340 | func (w *oneByteWriter) Write(p []byte) (int, error) { |
| 341 | w.mu.Lock() |
| 342 | defer w.mu.Unlock() |
| 343 | if len(p) == 0 { |
| 344 | return 0, nil |
| 345 | } |
| 346 | _ = w.buf.WriteByte(p[0]) |
| 347 | runtime.Gosched() |
| 348 | return 1, nil |
| 349 | } |
| 350 | |
| 351 | func (w *oneByteWriter) Bytes() []byte { |
| 352 | w.mu.Lock() |
| 353 | defer w.mu.Unlock() |
| 354 | return append([]byte(nil), w.buf.Bytes()...) |
| 355 | } |
| 356 | |
| 357 | func TestConcurrentWritesDoNotInterleaveFrames(t *testing.T) { |
| 358 | const count = 64 |
| 359 | w := &oneByteWriter{} |
| 360 | conn := NewConn(strings.NewReader(""), w, Options{}) |
| 361 | var wg sync.WaitGroup |
| 362 | errs := make(chan error, count) |
| 363 | for i := 0; i < count; i++ { |
| 364 | wg.Add(1) |
| 365 | go func(i int) { |
| 366 | defer wg.Done() |
| 367 | errs <- conn.Notify("event", map[string]int{"index": i}) |
| 368 | }(i) |
| 369 | } |
| 370 | wg.Wait() |
| 371 | close(errs) |
| 372 | for err := range errs { |
| 373 | if err != nil { |
| 374 | t.Fatalf("Notify: %v", err) |
| 375 | } |
| 376 | } |
| 377 | |
| 378 | lines := bytes.Split(bytes.TrimSpace(w.Bytes()), []byte{'\n'}) |
| 379 | if len(lines) != count { |
| 380 | t.Fatalf("frame count = %d, want %d", len(lines), count) |
| 381 | } |
| 382 | seen := make(map[int]bool, count) |
| 383 | for i, line := range lines { |
| 384 | var frame struct { |
| 385 | JSONRPC string `json:"jsonrpc"` |
| 386 | Method string `json:"method"` |
| 387 | Params struct { |
| 388 | Index int `json:"index"` |
| 389 | } `json:"params"` |
| 390 | } |
| 391 | if err := json.Unmarshal(line, &frame); err != nil { |
| 392 | t.Fatalf("frame %d is interleaved or invalid (%q): %v", i, line, err) |
| 393 | } |
| 394 | if frame.JSONRPC != "2.0" || frame.Method != "event" { |
| 395 | t.Fatalf("frame %d = %+v", i, frame) |
| 396 | } |
| 397 | if seen[frame.Params.Index] { |
| 398 | t.Fatalf("duplicate index %d", frame.Params.Index) |
| 399 | } |
| 400 | seen[frame.Params.Index] = true |
| 401 | } |
| 402 | if len(seen) != count { |
| 403 | t.Fatalf("unique payloads = %d, want %d", len(seen), count) |
| 404 | } |
| 405 | } |
| 406 | |
| 407 | type failWriter struct{ err error } |
| 408 | |
| 409 | func (w failWriter) Write([]byte) (int, error) { return 0, w.err } |
| 410 | |
| 411 | func TestHandlerResponseWriteFailureTerminatesConnection(t *testing.T) { |
| 412 | wantErr := errors.New("write failed") |
| 413 | request := "{\"jsonrpc\":\"2.0\",\"id\":1,\"method\":\"ok\",\"params\":{}}\n" |
| 414 | conn := NewConn(strings.NewReader(request), failWriter{err: wantErr}, Options{}) |
| 415 | conn.Handle("ok", func(context.Context, json.RawMessage) (any, error) { |
| 416 | return map[string]bool{"ok": true}, nil |
| 417 | }) |
| 418 | if err := conn.Serve(context.Background()); !errors.Is(err, wantErr) { |
| 419 | t.Fatalf("Serve error = %v, want %v", err, wantErr) |
| 420 | } |
| 421 | } |
| 422 | |
| 423 | func TestOversizedHandlerResultFailsConnectionWhenErrorCannotFit(t *testing.T) { |
| 424 | request := "{\"jsonrpc\":\"2.0\",\"id\":1,\"method\":\"large\",\"params\":{}}\n" |
| 425 | var out bytes.Buffer |
| 426 | conn := NewConn(strings.NewReader(request), &out, Options{MaxOutboundBytes: 24}) |
| 427 | conn.Handle("large", func(context.Context, json.RawMessage) (any, error) { |
| 428 | return map[string]string{"body": strings.Repeat("x", 1024)}, nil |
| 429 | }) |
| 430 | err := conn.Serve(context.Background()) |
| 431 | var tooLarge *FrameTooLargeError |
| 432 | if !errors.As(err, &tooLarge) || tooLarge.Direction != "outbound" || tooLarge.Limit != 24 { |
| 433 | t.Fatalf("Serve error = %v", err) |
| 434 | } |
| 435 | if out.Len() != 0 { |
| 436 | t.Fatalf("wrote oversized fallback response %q", out.String()) |
| 437 | } |
| 438 | } |
| 439 | |
| 440 | func TestRequestWriteFailureReturnsOriginalError(t *testing.T) { |
| 441 | wantErr := errors.New("request write failed") |
| 442 | conn := NewConn(strings.NewReader(""), failWriter{err: wantErr}, Options{Name: "test"}) |
| 443 | _, err := conn.Request(context.Background(), "call", map[string]string{"v": fmt.Sprint(1)}) |
| 444 | if !errors.Is(err, wantErr) { |
| 445 | t.Fatalf("Request error = %v, want %v", err, wantErr) |
| 446 | } |
| 447 | } |
| 448 | |
| 449 | func TestRequestStartedAfterServeEOFReturnsClosed(t *testing.T) { |
| 450 | conn := NewConn(strings.NewReader(""), io.Discard, Options{Name: "closed-race"}) |
| 451 | done := make(chan error, 1) |
| 452 | go func() { done <- conn.Serve(context.Background()) }() |
| 453 | if err := <-done; err != nil { |
| 454 | t.Fatalf("Serve: %v", err) |
| 455 | } |
| 456 | ctx, cancel := context.WithTimeout(context.Background(), time.Second) |
| 457 | defer cancel() |
| 458 | _, err := conn.Request(ctx, "late", struct{}{}) |
| 459 | if err == nil || errors.Is(err, context.DeadlineExceeded) { |
| 460 | t.Fatalf("late Request error = %v, want immediate connection closed", err) |
| 461 | } |
| 462 | } |
| 463 |