| 1 | package rpcwire |
| 2 | |
| 3 | import ( |
| 4 | "bufio" |
| 5 | "bytes" |
| 6 | "context" |
| 7 | "encoding/json" |
| 8 | "errors" |
| 9 | "fmt" |
| 10 | "io" |
| 11 | "strconv" |
| 12 | "sync" |
| 13 | "sync/atomic" |
| 14 | "time" |
| 15 | ) |
| 16 | |
| 17 | // RequestHandler answers an inbound JSON-RPC request. |
| 18 | type RequestHandler func(ctx context.Context, params json.RawMessage) (any, error) |
| 19 | |
| 20 | // HandlerResponse lets a protocol perform transport-local cleanup only after a |
| 21 | // successful response write. The callback runs exactly once with the result |
| 22 | // frame's write error (nil on success). It must be fast and must not write to |
| 23 | // the same Conn. This is intentionally transport-neutral: for example, a |
| 24 | // protocol can acknowledge detach before releasing its connection ownership. |
| 25 | type HandlerResponse struct { |
| 26 | Result any |
| 27 | AfterWrite func(error) |
| 28 | } |
| 29 | |
| 30 | // RespondThen wraps a handler result with an after-write callback. |
| 31 | func RespondThen(result any, afterWrite func(error)) HandlerResponse { |
| 32 | return HandlerResponse{Result: result, AfterWrite: afterWrite} |
| 33 | } |
| 34 | |
| 35 | // NotificationHandler handles an inbound JSON-RPC notification. |
| 36 | type NotificationHandler func(ctx context.Context, params json.RawMessage) |
| 37 | |
| 38 | // Options configures transport-only behavior. A non-positive frame limit means |
| 39 | // unlimited in that direction. Protocol adapters should always set an inbound |
| 40 | // limit for untrusted peers. |
| 41 | type Options struct { |
| 42 | MaxInboundBytes int |
| 43 | MaxOutboundBytes int |
| 44 | Name string |
| 45 | // StrictJSONRPC validates the jsonrpc member and mutually exclusive frame |
| 46 | // shapes. Extension Protocol peers always enable it. |
| 47 | StrictJSONRPC bool |
| 48 | // MaxConcurrentHandlers bounds inbound request and, unless a notification |
| 49 | // queue is configured, notification handlers without blocking response |
| 50 | // dispatch. Non-positive values use the safe default; overload requests |
| 51 | // receive ErrServerBusy. |
| 52 | MaxConcurrentHandlers int |
| 53 | // MaxQueuedNotifications enables ordered notification delivery through one |
| 54 | // bounded FIFO worker. A full queue fails the connection instead of silently |
| 55 | // losing a notification. Non-positive values preserve concurrent best-effort |
| 56 | // notification dispatch for protocols that do not require ordered delivery. |
| 57 | MaxQueuedNotifications int |
| 58 | // BeforeRequest runs synchronously on the read loop, after strict frame |
| 59 | // validation and before a handler goroutine is scheduled. It lets a protocol |
| 60 | // atomically record wire arrival order (for example, initialize-first) while |
| 61 | // preserving concurrent handler execution. Returning an error rejects only |
| 62 | // that request through the normal RPC error mapping. |
| 63 | BeforeRequest func(method string, params json.RawMessage) error |
| 64 | // BeforeNotification runs synchronously on the read loop before notification |
| 65 | // dispatch. Returning an error silently rejects the notification, as required |
| 66 | // by JSON-RPC, while allowing a protocol to poison transport-local state. |
| 67 | // The nil default preserves existing protocol behavior. |
| 68 | BeforeNotification func(method string, params json.RawMessage) error |
| 69 | // MaxWriteStall bounds how long a single outbound write may make no |
| 70 | // progress (the peer keeps the pipe open but has stopped reading) before |
| 71 | // the connection fails with WriteStallError. Non-positive disables the |
| 72 | // bound, preserving the historical block-forever behavior; stdio peers |
| 73 | // should always set it, since a wedged child otherwise hangs every caller. |
| 74 | MaxWriteStall time.Duration |
| 75 | } |
| 76 | |
| 77 | // Conn is one bidirectional JSON-RPC 2.0 connection framed as NDJSON. |
| 78 | type Conn struct { |
| 79 | r io.Reader |
| 80 | w io.Writer |
| 81 | opts Options |
| 82 | |
| 83 | // Exactly one writer goroutine owns w, fed by the bounded writeQ, so two |
| 84 | // frames can never interleave on the transport — even when a caller's |
| 85 | // context aborts mid-flight. writeActive/writeProgress back the optional |
| 86 | // stall watchdog (MaxWriteStall): a physical write making no progress for |
| 87 | // the bound fails the connection. |
| 88 | writeQ chan writeJob |
| 89 | writeSlots chan struct{} |
| 90 | writeGate sync.Mutex |
| 91 | writeClosed bool |
| 92 | writerDone chan struct{} |
| 93 | writeActive atomic.Bool |
| 94 | writeProgress atomic.Int64 |
| 95 | writerOnce sync.Once |
| 96 | |
| 97 | nextID atomic.Int64 |
| 98 | |
| 99 | pmu sync.Mutex |
| 100 | pending map[int64]chan rpcResult |
| 101 | |
| 102 | reqH map[string]RequestHandler |
| 103 | notH map[string]NotificationHandler |
| 104 | |
| 105 | wg sync.WaitGroup |
| 106 | closeOnce sync.Once |
| 107 | closed chan struct{} |
| 108 | closeMu sync.Mutex |
| 109 | closeErr error |
| 110 | handlerSlots chan struct{} |
| 111 | notifyQueue chan notificationCall |
| 112 | tryNotifySlots chan struct{} |
| 113 | } |
| 114 | |
| 115 | const DefaultMaxConcurrentHandlers = 64 |
| 116 | |
| 117 | type rpcResult struct { |
| 118 | result json.RawMessage |
| 119 | err error |
| 120 | } |
| 121 | |
| 122 | type notificationCall struct { |
| 123 | handler NotificationHandler |
| 124 | params json.RawMessage |
| 125 | } |
| 126 | |
| 127 | type outbound struct { |
| 128 | JSONRPC string `json:"jsonrpc"` |
| 129 | ID json.RawMessage `json:"id,omitempty"` |
| 130 | Method string `json:"method,omitempty"` |
| 131 | Params json.RawMessage `json:"params,omitempty"` |
| 132 | Result json.RawMessage `json:"result,omitempty"` |
| 133 | Error *ErrorObject `json:"error,omitempty"` |
| 134 | } |
| 135 | |
| 136 | type inbound struct { |
| 137 | JSONRPC string `json:"jsonrpc"` |
| 138 | ID json.RawMessage `json:"id"` |
| 139 | Method string `json:"method"` |
| 140 | Params json.RawMessage `json:"params"` |
| 141 | Result json.RawMessage `json:"result"` |
| 142 | Error *ErrorObject `json:"error"` |
| 143 | } |
| 144 | |
| 145 | // NewConn constructs a connection. Register handlers before calling Serve. |
| 146 | func NewConn(r io.Reader, w io.Writer, opts Options) *Conn { |
| 147 | if opts.Name == "" { |
| 148 | opts.Name = "rpcwire" |
| 149 | } |
| 150 | if opts.MaxConcurrentHandlers <= 0 { |
| 151 | opts.MaxConcurrentHandlers = DefaultMaxConcurrentHandlers |
| 152 | } |
| 153 | conn := &Conn{ |
| 154 | r: r, |
| 155 | w: w, |
| 156 | opts: opts, |
| 157 | pending: make(map[int64]chan rpcResult), |
| 158 | reqH: make(map[string]RequestHandler), |
| 159 | notH: make(map[string]NotificationHandler), |
| 160 | closed: make(chan struct{}), |
| 161 | handlerSlots: make(chan struct{}, opts.MaxConcurrentHandlers), |
| 162 | writeQ: make(chan writeJob, writeQueueLimit), |
| 163 | writeSlots: make(chan struct{}, writeQueueLimit), |
| 164 | writerDone: make(chan struct{}), |
| 165 | tryNotifySlots: make(chan struct{}, bestEffortNotifyQueueLimit), |
| 166 | } |
| 167 | if opts.MaxQueuedNotifications > 0 { |
| 168 | conn.notifyQueue = make(chan notificationCall, opts.MaxQueuedNotifications) |
| 169 | } |
| 170 | return conn |
| 171 | } |
| 172 | |
| 173 | // ensureWriter starts the single writer loop (and the stall watchdog when |
| 174 | // configured) exactly once, lazily on the first write or Serve. Lazy startup |
| 175 | // keeps a Conn that is constructed but never used — an attach rejected before |
| 176 | // Serve, for example — from leaking a permanent goroutine. |
| 177 | func (c *Conn) ensureWriter() { |
| 178 | c.writerOnce.Do(func() { |
| 179 | go c.writerLoop() |
| 180 | if c.opts.MaxWriteStall > 0 { |
| 181 | go c.stallWatchdog() |
| 182 | } |
| 183 | }) |
| 184 | } |
| 185 | |
| 186 | // Handle registers a request handler. It is not safe to mutate registrations |
| 187 | // concurrently with Serve. |
| 188 | func (c *Conn) Handle(method string, h RequestHandler) { c.reqH[method] = h } |
| 189 | |
| 190 | // HandleNotify registers a notification handler. |
| 191 | func (c *Conn) HandleNotify(method string, h NotificationHandler) { c.notH[method] = h } |
| 192 | |
| 193 | // Serve reads and dispatches frames until EOF, cancellation observed by the |
| 194 | // read loop, or a framing/read error. In-flight handler contexts are cancelled |
| 195 | // when the transport ends; a product that needs work to outlive the connection |
| 196 | // must derive that work from its own runtime context before returning. |
| 197 | func (c *Conn) Serve(ctx context.Context) error { |
| 198 | c.ensureWriter() |
| 199 | ctx, cancel := context.WithCancel(ctx) |
| 200 | defer cancel() |
| 201 | if c.notifyQueue != nil { |
| 202 | c.wg.Add(1) |
| 203 | go c.serveNotifications(ctx) |
| 204 | } |
| 205 | |
| 206 | br := bufio.NewReaderSize(c.r, 64<<10) |
| 207 | var loopErr error |
| 208 | for { |
| 209 | line, err := readLine(br, c.opts.MaxInboundBytes) |
| 210 | if len(line) > 0 { |
| 211 | c.dispatch(ctx, line) |
| 212 | } |
| 213 | if err != nil { |
| 214 | if !errors.Is(err, io.EOF) { |
| 215 | loopErr = c.decorateReadError(err) |
| 216 | } |
| 217 | break |
| 218 | } |
| 219 | if err := ctx.Err(); err != nil { |
| 220 | loopErr = err |
| 221 | break |
| 222 | } |
| 223 | } |
| 224 | |
| 225 | cancel() |
| 226 | if c.notifyQueue != nil { |
| 227 | close(c.notifyQueue) |
| 228 | } |
| 229 | c.wg.Wait() |
| 230 | if terminalErr := c.closeReason(); terminalErr != nil { |
| 231 | loopErr = terminalErr |
| 232 | } |
| 233 | c.shutdown(nil) |
| 234 | return loopErr |
| 235 | } |
| 236 | |
| 237 | func (c *Conn) serveNotifications(ctx context.Context) { |
| 238 | defer c.wg.Done() |
| 239 | for call := range c.notifyQueue { |
| 240 | call.handler(ctx, call.params) |
| 241 | } |
| 242 | } |
| 243 | |
| 244 | func (c *Conn) decorateReadError(err error) error { |
| 245 | var tooLarge *FrameTooLargeError |
| 246 | if errors.As(err, &tooLarge) { |
| 247 | return fmt.Errorf("%s: message exceeds size limit: %w", c.opts.Name, err) |
| 248 | } |
| 249 | return err |
| 250 | } |
| 251 | |
| 252 | func (c *Conn) dispatch(ctx context.Context, line []byte) { |
| 253 | var in inbound |
| 254 | if err := json.Unmarshal(line, &in); err != nil { |
| 255 | if c.opts.StrictJSONRPC && json.Valid(line) { |
| 256 | c.respondError(json.RawMessage("null"), ErrInvalidRequest, "invalid request", nil) |
| 257 | } else { |
| 258 | c.respondError(json.RawMessage("null"), ErrParse, "parse error", nil) |
| 259 | } |
| 260 | return |
| 261 | } |
| 262 | if c.opts.StrictJSONRPC { |
| 263 | if err := validateStrictFrame(line, in); err != nil { |
| 264 | c.respondError(ResponseIDForError(in.ID), ErrInvalidRequest, "invalid request", nil) |
| 265 | return |
| 266 | } |
| 267 | } |
| 268 | select { |
| 269 | case <-c.closed: |
| 270 | return |
| 271 | default: |
| 272 | } |
| 273 | hasID := len(in.ID) > 0 |
| 274 | switch { |
| 275 | case in.Method != "" && hasID: |
| 276 | if c.opts.BeforeRequest != nil { |
| 277 | if err := c.opts.BeforeRequest(in.Method, in.Params); err != nil { |
| 278 | c.respondHandlerError(in.ID, err) |
| 279 | return |
| 280 | } |
| 281 | } |
| 282 | if !c.tryStartHandler() { |
| 283 | c.respondError(in.ID, ErrServerBusy, "server busy", nil) |
| 284 | return |
| 285 | } |
| 286 | c.wg.Add(1) |
| 287 | go func() { |
| 288 | defer c.finishHandler() |
| 289 | defer c.wg.Done() |
| 290 | c.serveRequest(ctx, in.ID, in.Method, in.Params) |
| 291 | }() |
| 292 | case in.Method != "" && !hasID: |
| 293 | if c.opts.BeforeNotification != nil { |
| 294 | if err := c.opts.BeforeNotification(in.Method, in.Params); err != nil { |
| 295 | return |
| 296 | } |
| 297 | } |
| 298 | if h := c.notH[in.Method]; h != nil { |
| 299 | if c.notifyQueue != nil { |
| 300 | select { |
| 301 | case c.notifyQueue <- notificationCall{handler: h, params: in.Params}: |
| 302 | default: |
| 303 | c.fail(fmt.Errorf("%s: notification queue overflow", c.opts.Name)) |
| 304 | } |
| 305 | return |
| 306 | } |
| 307 | if !c.tryStartHandler() { |
| 308 | return |
| 309 | } |
| 310 | c.wg.Add(1) |
| 311 | go func() { |
| 312 | defer c.finishHandler() |
| 313 | defer c.wg.Done() |
| 314 | h(ctx, in.Params) |
| 315 | }() |
| 316 | } |
| 317 | case in.Method == "" && hasID: |
| 318 | c.resolve(in) |
| 319 | default: |
| 320 | c.respondError(json.RawMessage("null"), ErrInvalidRequest, "invalid request", nil) |
| 321 | } |
| 322 | } |
| 323 | |
| 324 | func (c *Conn) tryStartHandler() bool { |
| 325 | select { |
| 326 | case c.handlerSlots <- struct{}{}: |
| 327 | return true |
| 328 | default: |
| 329 | return false |
| 330 | } |
| 331 | } |
| 332 | |
| 333 | func (c *Conn) finishHandler() { <-c.handlerSlots } |
| 334 | |
| 335 | func validateStrictFrame(line []byte, in inbound) error { |
| 336 | var members map[string]json.RawMessage |
| 337 | if err := json.Unmarshal(line, &members); err != nil { |
| 338 | return err |
| 339 | } |
| 340 | if in.JSONRPC != "2.0" { |
| 341 | return errors.New("jsonrpc must be 2.0") |
| 342 | } |
| 343 | _, hasID := members["id"] |
| 344 | _, hasMethod := members["method"] |
| 345 | _, hasParams := members["params"] |
| 346 | _, hasResult := members["result"] |
| 347 | _, hasError := members["error"] |
| 348 | if hasID && !validRPCID(in.ID) { |
| 349 | return errors.New("id must be a string, integer, or null") |
| 350 | } |
| 351 | if hasMethod { |
| 352 | if in.Method == "" || hasResult || hasError { |
| 353 | return errors.New("invalid request shape") |
| 354 | } |
| 355 | if hasParams { |
| 356 | trimmed := bytes.TrimSpace(in.Params) |
| 357 | if len(trimmed) == 0 || (trimmed[0] != '{' && trimmed[0] != '[') { |
| 358 | return errors.New("params must be object or array") |
| 359 | } |
| 360 | } |
| 361 | return nil |
| 362 | } |
| 363 | if !hasID || hasParams || hasResult == hasError { |
| 364 | return errors.New("invalid response shape") |
| 365 | } |
| 366 | if hasError && in.Error == nil { |
| 367 | return errors.New("invalid error object") |
| 368 | } |
| 369 | if hasError { |
| 370 | var errorMembers map[string]json.RawMessage |
| 371 | if err := json.Unmarshal(members["error"], &errorMembers); err != nil { |
| 372 | return errors.New("invalid error object") |
| 373 | } |
| 374 | if _, ok := errorMembers["code"]; !ok { |
| 375 | return errors.New("error code is required") |
| 376 | } |
| 377 | if _, ok := errorMembers["message"]; !ok { |
| 378 | return errors.New("error message is required") |
| 379 | } |
| 380 | } |
| 381 | return nil |
| 382 | } |
| 383 | |
| 384 | func validRPCID(raw json.RawMessage) bool { |
| 385 | raw = bytes.TrimSpace(raw) |
| 386 | if bytes.Equal(raw, []byte("null")) { |
| 387 | return true |
| 388 | } |
| 389 | if len(raw) >= 2 && raw[0] == '"' && raw[len(raw)-1] == '"' { |
| 390 | return true |
| 391 | } |
| 392 | if len(raw) == 0 { |
| 393 | return false |
| 394 | } |
| 395 | i := 0 |
| 396 | if raw[0] == '-' { |
| 397 | i++ |
| 398 | if i == len(raw) { |
| 399 | return false |
| 400 | } |
| 401 | } |
| 402 | if raw[i] == '0' && i+1 != len(raw) { |
| 403 | return false |
| 404 | } |
| 405 | for ; i < len(raw); i++ { |
| 406 | if raw[i] < '0' || raw[i] > '9' { |
| 407 | return false |
| 408 | } |
| 409 | } |
| 410 | return true |
| 411 | } |
| 412 | |
| 413 | func (c *Conn) serveRequest(ctx context.Context, id json.RawMessage, method string, params json.RawMessage) { |
| 414 | h := c.reqH[method] |
| 415 | if h == nil { |
| 416 | c.respondError(id, ErrMethodNotFound, "method not found: "+method, nil) |
| 417 | return |
| 418 | } |
| 419 | result, err := h(ctx, params) |
| 420 | if err != nil { |
| 421 | c.respondHandlerError(id, err) |
| 422 | return |
| 423 | } |
| 424 | var afterWrite func(error) |
| 425 | if response, ok := result.(HandlerResponse); ok { |
| 426 | result = response.Result |
| 427 | afterWrite = response.AfterWrite |
| 428 | } |
| 429 | raw, err := json.Marshal(result) |
| 430 | if err != nil { |
| 431 | c.respondError(id, ErrInternal, "marshal result: "+err.Error(), nil) |
| 432 | c.runAfterWrite(afterWrite, err) |
| 433 | return |
| 434 | } |
| 435 | writeErr := c.write(context.Background(), outbound{JSONRPC: "2.0", ID: id, Result: raw}) |
| 436 | if writeErr != nil { |
| 437 | var tooLarge *FrameTooLargeError |
| 438 | if errors.As(writeErr, &tooLarge) { |
| 439 | c.respondError(id, ErrInternal, "response exceeds frame size limit", nil) |
| 440 | c.runAfterWrite(afterWrite, writeErr) |
| 441 | return |
| 442 | } |
| 443 | c.fail(writeErr) |
| 444 | } |
| 445 | c.runAfterWrite(afterWrite, writeErr) |
| 446 | } |
| 447 | |
| 448 | func (c *Conn) runAfterWrite(callback func(error), writeErr error) { |
| 449 | if callback == nil { |
| 450 | return |
| 451 | } |
| 452 | defer func() { |
| 453 | if recovered := recover(); recovered != nil { |
| 454 | c.fail(fmt.Errorf("%s: after-response callback panic: %v", c.opts.Name, recovered)) |
| 455 | } |
| 456 | }() |
| 457 | callback(writeErr) |
| 458 | } |
| 459 | |
| 460 | func (c *Conn) respondHandlerError(id json.RawMessage, err error) { |
| 461 | code := ErrInternal |
| 462 | message := err.Error() |
| 463 | var data any |
| 464 | var re *RPCError |
| 465 | if errors.As(err, &re) { |
| 466 | code = re.Code |
| 467 | message = re.Message |
| 468 | data = re.Data |
| 469 | } |
| 470 | c.respondError(id, code, message, data) |
| 471 | } |
| 472 | |
| 473 | func (c *Conn) resolve(in inbound) { |
| 474 | id, err := strconv.ParseInt(string(in.ID), 10, 64) |
| 475 | if err != nil { |
| 476 | return |
| 477 | } |
| 478 | c.pmu.Lock() |
| 479 | ch := c.pending[id] |
| 480 | delete(c.pending, id) |
| 481 | c.pmu.Unlock() |
| 482 | if ch == nil { |
| 483 | return |
| 484 | } |
| 485 | if in.Error != nil { |
| 486 | ch <- rpcResult{err: &ResponseError{Code: in.Error.Code, Message: in.Error.Message, Data: in.Error.Data}} |
| 487 | return |
| 488 | } |
| 489 | ch <- rpcResult{result: in.Result} |
| 490 | } |
| 491 | |
| 492 | // Notify sends a fire-and-forget notification. |
| 493 | func (c *Conn) Notify(method string, params any) error { |
| 494 | m, err := notification(method, params) |
| 495 | if err != nil { |
| 496 | return err |
| 497 | } |
| 498 | err = c.write(context.Background(), m) |
| 499 | if err != nil { |
| 500 | var tooLarge *FrameTooLargeError |
| 501 | if !errors.As(err, &tooLarge) { |
| 502 | c.fail(err) |
| 503 | } |
| 504 | } |
| 505 | return err |
| 506 | } |
| 507 | |
| 508 | // TryNotify enqueues a fire-and-forget notification without waiting for a |
| 509 | // physical write. A nil result means the bounded writer accepted the frame, |
| 510 | // not that the peer has processed it. When the queue is full it returns |
| 511 | // OutboundQueueFullError immediately, allowing observation-only callers to |
| 512 | // drop the event instead of adding sidecar backpressure to a host hot path. |
| 513 | func (c *Conn) TryNotify(method string, params any) error { |
| 514 | m, err := notification(method, params) |
| 515 | if err != nil { |
| 516 | return err |
| 517 | } |
| 518 | job, err := c.prepareWrite(m, context.Background()) |
| 519 | if err != nil { |
| 520 | return err |
| 521 | } |
| 522 | select { |
| 523 | case c.tryNotifySlots <- struct{}{}: |
| 524 | job.release = func() { <-c.tryNotifySlots } |
| 525 | default: |
| 526 | return &OutboundQueueFullError{Limit: cap(c.tryNotifySlots)} |
| 527 | } |
| 528 | if err := c.enqueueWrite(job, false); err != nil { |
| 529 | job.release() |
| 530 | return err |
| 531 | } |
| 532 | return nil |
| 533 | } |
| 534 | |
| 535 | func notification(method string, params any) (outbound, error) { |
| 536 | raw, err := json.Marshal(params) |
| 537 | if err != nil { |
| 538 | return outbound{}, err |
| 539 | } |
| 540 | return outbound{JSONRPC: "2.0", Method: method, Params: raw}, nil |
| 541 | } |
| 542 | |
| 543 | // Request sends a request and waits for its response, cancellation, or closure. |
| 544 | func (c *Conn) Request(ctx context.Context, method string, params any) (json.RawMessage, error) { |
| 545 | raw, err := json.Marshal(params) |
| 546 | if err != nil { |
| 547 | return nil, err |
| 548 | } |
| 549 | id := c.nextID.Add(1) |
| 550 | ch := make(chan rpcResult, 1) |
| 551 | c.pmu.Lock() |
| 552 | select { |
| 553 | case <-c.closed: |
| 554 | c.pmu.Unlock() |
| 555 | closedErr := c.terminalError() |
| 556 | if closedErr == nil { |
| 557 | closedErr = fmt.Errorf("%s: connection closed", c.opts.Name) |
| 558 | } |
| 559 | return nil, closedErr |
| 560 | default: |
| 561 | } |
| 562 | c.pending[id] = ch |
| 563 | c.pmu.Unlock() |
| 564 | defer func() { |
| 565 | c.pmu.Lock() |
| 566 | delete(c.pending, id) |
| 567 | c.pmu.Unlock() |
| 568 | }() |
| 569 | |
| 570 | idRaw, _ := json.Marshal(id) |
| 571 | if err := c.write(ctx, outbound{JSONRPC: "2.0", ID: idRaw, Method: method, Params: raw}); err != nil { |
| 572 | var tooLarge *FrameTooLargeError |
| 573 | // A caller-context abort (turn cancel, per-call timeout) fails only |
| 574 | // this request — the connection stays usable. Genuine transport |
| 575 | // failures, including a write that stalled past MaxWriteStall, fail |
| 576 | // the connection so a wedged peer is torn down. |
| 577 | if !errors.As(err, &tooLarge) && !errors.Is(err, context.Canceled) && !errors.Is(err, context.DeadlineExceeded) { |
| 578 | c.fail(err) |
| 579 | } |
| 580 | return nil, err |
| 581 | } |
| 582 | select { |
| 583 | case res := <-ch: |
| 584 | return res.result, res.err |
| 585 | case <-ctx.Done(): |
| 586 | return nil, ctx.Err() |
| 587 | } |
| 588 | } |
| 589 | |
| 590 | func (c *Conn) write(ctx context.Context, m outbound) error { |
| 591 | job, err := c.prepareWrite(m, ctx) |
| 592 | if err != nil { |
| 593 | return err |
| 594 | } |
| 595 | if err := c.enqueueWrite(job, true); err != nil { |
| 596 | return err |
| 597 | } |
| 598 | select { |
| 599 | case err := <-job.res: |
| 600 | return err |
| 601 | case <-job.ctx.Done(): |
| 602 | // The caller gave up: the writer loop will skip the frame if it has |
| 603 | // not physically started, or finish it serially if it has — the |
| 604 | // transport never sees a torn or interleaved frame. |
| 605 | return job.ctx.Err() |
| 606 | case <-c.closed: |
| 607 | // A completed write may race the connection's teardown; the buffered |
| 608 | // result is already there when that happened, so prefer it over the |
| 609 | // terminal error (the response-close regression class). |
| 610 | select { |
| 611 | case err := <-job.res: |
| 612 | return err |
| 613 | default: |
| 614 | } |
| 615 | return c.terminalError() |
| 616 | } |
| 617 | } |
| 618 | |
| 619 | // enqueueWrite reserves bounded queue capacity before entering writeGate. |
| 620 | // The reservation makes the send non-blocking while the gate is held, so |
| 621 | // shutdown can close writeQ without racing a producer or waiting behind a |
| 622 | // producer blocked on a full queue. blocking is false for TryNotify. |
| 623 | func (c *Conn) enqueueWrite(job writeJob, blocking bool) error { |
| 624 | c.ensureWriter() |
| 625 | if blocking { |
| 626 | select { |
| 627 | case c.writeSlots <- struct{}{}: |
| 628 | case <-job.ctx.Done(): |
| 629 | return job.ctx.Err() |
| 630 | case <-c.closed: |
| 631 | return c.terminalError() |
| 632 | } |
| 633 | } else { |
| 634 | select { |
| 635 | case c.writeSlots <- struct{}{}: |
| 636 | default: |
| 637 | return &OutboundQueueFullError{Limit: cap(c.tryNotifySlots)} |
| 638 | } |
| 639 | } |
| 640 | |
| 641 | c.writeGate.Lock() |
| 642 | defer c.writeGate.Unlock() |
| 643 | if c.writeClosed { |
| 644 | <-c.writeSlots |
| 645 | return c.terminalError() |
| 646 | } |
| 647 | // A reserved slot guarantees capacity; the send cannot block while the |
| 648 | // gate is held. Keeping the normal send makes accounting bugs fail loudly. |
| 649 | c.writeQ <- job |
| 650 | return nil |
| 651 | } |
| 652 | |
| 653 | func (c *Conn) prepareWrite(m outbound, ctx context.Context) (writeJob, error) { |
| 654 | var buf bytes.Buffer |
| 655 | enc := json.NewEncoder(&buf) |
| 656 | enc.SetEscapeHTML(false) |
| 657 | if err := enc.Encode(m); err != nil { |
| 658 | return writeJob{}, err |
| 659 | } |
| 660 | if limit := c.opts.MaxOutboundBytes; limit > 0 && buf.Len() > limit { |
| 661 | return writeJob{}, &FrameTooLargeError{Direction: "outbound", Size: buf.Len(), Limit: limit} |
| 662 | } |
| 663 | if ctx == nil { |
| 664 | ctx = context.Background() |
| 665 | } |
| 666 | return writeJob{frame: buf.Bytes(), ctx: ctx, res: make(chan error, 1)}, nil |
| 667 | } |
| 668 | |
| 669 | // writeQueueLimit bounds queued outbound frames per connection. A wedged |
| 670 | // peer fills the queue and then the stall watchdog fails the connection; |
| 671 | // senders never block unboundedly behind it. |
| 672 | const writeQueueLimit = 256 |
| 673 | |
| 674 | // bestEffortNotifyQueueLimit prevents observation events from filling the |
| 675 | // shared writer queue ahead of request/response traffic. Sixteen queued or |
| 676 | // in-flight events absorb healthy bursts while preserving capacity and |
| 677 | // latency for blocking intercept, provider, UI, and shutdown calls. |
| 678 | const bestEffortNotifyQueueLimit = 16 |
| 679 | |
| 680 | // writeJob is one outbound frame awaiting the single writer goroutine. |
| 681 | type writeJob struct { |
| 682 | frame []byte |
| 683 | ctx context.Context // pre-write cancellation only |
| 684 | res chan error // buffered 1 |
| 685 | release func() // releases optional best-effort notification capacity |
| 686 | } |
| 687 | |
| 688 | func completeWriteJob(job writeJob, err error) { |
| 689 | job.res <- err |
| 690 | if job.release != nil { |
| 691 | job.release() |
| 692 | } |
| 693 | } |
| 694 | |
| 695 | // writerLoop is the ONLY writer of c.w. It drains the queue in order, skips |
| 696 | // frames whose caller already gave up before the physical write began, and |
| 697 | // finishes any frame it started — frames are atomic and ordered by |
| 698 | // construction. On connection close it fails everything still queued. |
| 699 | func (c *Conn) writerLoop() { |
| 700 | defer close(c.writerDone) |
| 701 | for job := range c.writeQ { |
| 702 | <-c.writeSlots |
| 703 | c.writeGate.Lock() |
| 704 | closed := c.writeClosed |
| 705 | c.writeGate.Unlock() |
| 706 | if closed { |
| 707 | completeWriteJob(job, c.terminalError()) |
| 708 | continue |
| 709 | } |
| 710 | if job.ctx != nil { |
| 711 | if err := job.ctx.Err(); err != nil { |
| 712 | completeWriteJob(job, err) |
| 713 | continue |
| 714 | } |
| 715 | } |
| 716 | err := c.writeAll(job.frame) |
| 717 | if err != nil { |
| 718 | // When the connection is already terminal, report the root |
| 719 | // cause (e.g. the stall watchdog's WriteStallError) rather |
| 720 | // than its side effect — a transport closing underneath an |
| 721 | // in-flight write surfaces as a plain closed-pipe error. |
| 722 | select { |
| 723 | case <-c.closed: |
| 724 | if terminal := c.terminalError(); terminal != nil { |
| 725 | err = terminal |
| 726 | } |
| 727 | default: |
| 728 | } |
| 729 | } |
| 730 | completeWriteJob(job, err) |
| 731 | if err != nil { |
| 732 | // shutdown waits for writerDone. Run it outside this goroutine so |
| 733 | // writerLoop can return and satisfy that lifecycle handshake. |
| 734 | go c.fail(err) |
| 735 | return |
| 736 | } |
| 737 | } |
| 738 | } |
| 739 | |
| 740 | // writeAll serially writes one frame, marking activity for the stall |
| 741 | // watchdog before every blocking Write. It only returns on completion or a |
| 742 | // transport error — caller cancellation never tears a frame in half. |
| 743 | func (c *Conn) writeAll(b []byte) error { |
| 744 | for len(b) > 0 { |
| 745 | c.writeProgress.Store(time.Now().UnixNano()) |
| 746 | c.writeActive.Store(true) |
| 747 | n, err := c.w.Write(b) |
| 748 | c.writeActive.Store(false) |
| 749 | if err != nil { |
| 750 | return err |
| 751 | } |
| 752 | if n == 0 { |
| 753 | return io.ErrShortWrite |
| 754 | } |
| 755 | b = b[n:] |
| 756 | } |
| 757 | return nil |
| 758 | } |
| 759 | |
| 760 | // stallWatchdog fails the connection when a physical write makes no progress |
| 761 | // for MaxWriteStall: the peer is alive enough to hold the pipe open but has |
| 762 | // stopped reading, and without a bound every later frame would queue behind |
| 763 | // it forever. The watchdog is deliberately independent of any caller |
| 764 | // context, so a short per-call timeout cannot preempt it. |
| 765 | func (c *Conn) stallWatchdog() { |
| 766 | interval := c.opts.MaxWriteStall / 2 |
| 767 | if interval <= 0 { |
| 768 | interval = time.Millisecond |
| 769 | } |
| 770 | ticker := time.NewTicker(interval) |
| 771 | defer ticker.Stop() |
| 772 | for { |
| 773 | select { |
| 774 | case <-ticker.C: |
| 775 | if !c.writeActive.Load() { |
| 776 | continue |
| 777 | } |
| 778 | last := time.Unix(0, c.writeProgress.Load()) |
| 779 | if time.Since(last) > c.opts.MaxWriteStall { |
| 780 | c.fail(&WriteStallError{Direction: "outbound", Stall: c.opts.MaxWriteStall}) |
| 781 | return |
| 782 | } |
| 783 | case <-c.closed: |
| 784 | return |
| 785 | } |
| 786 | } |
| 787 | } |
| 788 | |
| 789 | func (c *Conn) writeError(id json.RawMessage, code int, message string, data any) error { |
| 790 | var raw json.RawMessage |
| 791 | if data != nil { |
| 792 | encoded, err := json.Marshal(data) |
| 793 | if err != nil { |
| 794 | code = ErrInternal |
| 795 | message = "marshal error data: " + err.Error() |
| 796 | } else if string(encoded) != "null" { |
| 797 | raw = encoded |
| 798 | } |
| 799 | } |
| 800 | return c.write(context.Background(), outbound{JSONRPC: "2.0", ID: id, Error: &ErrorObject{Code: code, Message: message, Data: raw}}) |
| 801 | } |
| 802 | |
| 803 | func (c *Conn) respondError(id json.RawMessage, code int, message string, data any) { |
| 804 | err := c.writeError(id, code, message, data) |
| 805 | var tooLarge *FrameTooLargeError |
| 806 | if errors.As(err, &tooLarge) && (data != nil || code != ErrInternal || message != "error response exceeds frame size limit") { |
| 807 | err = c.writeError(id, ErrInternal, "error response exceeds frame size limit", nil) |
| 808 | } |
| 809 | if err != nil { |
| 810 | c.fail(err) |
| 811 | } |
| 812 | } |
| 813 | |
| 814 | func (c *Conn) fail(err error) { |
| 815 | if err == nil { |
| 816 | return |
| 817 | } |
| 818 | c.shutdown(err) |
| 819 | if closer, ok := c.r.(io.Closer); ok { |
| 820 | _ = closer.Close() |
| 821 | } |
| 822 | } |
| 823 | |
| 824 | // closeReason returns the stored terminal error as-is (nil on a clean EOF); |
| 825 | // Serve uses it so a graceful end still reports nil. |
| 826 | func (c *Conn) closeReason() error { |
| 827 | c.closeMu.Lock() |
| 828 | defer c.closeMu.Unlock() |
| 829 | return c.closeErr |
| 830 | } |
| 831 | |
| 832 | // terminalError is the error every caller observes after the connection |
| 833 | // ends. It is never nil — a graceful EOF must not report silently dropped |
| 834 | // writes as successes. |
| 835 | func (c *Conn) terminalError() error { |
| 836 | c.closeMu.Lock() |
| 837 | defer c.closeMu.Unlock() |
| 838 | if c.closeErr != nil { |
| 839 | return c.closeErr |
| 840 | } |
| 841 | return fmt.Errorf("%s: connection closed", c.opts.Name) |
| 842 | } |
| 843 | |
| 844 | // writerExitWaitBound caps how long shutdown lets an in-flight physical write |
| 845 | // finish. A wedged writer is failed by the stall watchdog; teardown itself must |
| 846 | // still be bounded. |
| 847 | const writerExitWaitBound = 100 * time.Millisecond |
| 848 | |
| 849 | func (c *Conn) shutdown(err error) { |
| 850 | c.closeOnce.Do(func() { |
| 851 | c.closeMu.Lock() |
| 852 | c.closeErr = err |
| 853 | c.closeMu.Unlock() |
| 854 | |
| 855 | // Linearize closure against every producer, then close the queue. No |
| 856 | // producer can send after writeClosed becomes visible because enqueue |
| 857 | // performs its final check and send under the same gate. |
| 858 | c.ensureWriter() |
| 859 | c.writeGate.Lock() |
| 860 | c.writeClosed = true |
| 861 | close(c.writeQ) |
| 862 | c.writeGate.Unlock() |
| 863 | select { |
| 864 | case <-c.writerDone: |
| 865 | case <-time.After(writerExitWaitBound): |
| 866 | } |
| 867 | close(c.closed) |
| 868 | c.pmu.Lock() |
| 869 | for id, ch := range c.pending { |
| 870 | closedErr := err |
| 871 | if closedErr == nil { |
| 872 | closedErr = fmt.Errorf("%s: connection closed", c.opts.Name) |
| 873 | } |
| 874 | ch <- rpcResult{err: closedErr} |
| 875 | delete(c.pending, id) |
| 876 | } |
| 877 | c.pmu.Unlock() |
| 878 | }) |
| 879 | } |
| 880 |