| 1 | package extension |
| 2 | |
| 3 | import ( |
| 4 | "bufio" |
| 5 | "bytes" |
| 6 | "context" |
| 7 | "encoding/json" |
| 8 | "errors" |
| 9 | "fmt" |
| 10 | "io" |
| 11 | "log" |
| 12 | "strconv" |
| 13 | "sync" |
| 14 | "sync/atomic" |
| 15 | ) |
| 16 | |
| 17 | // Standard JSON-RPC 2.0 error codes, plus the extension domain code and the |
| 18 | // transport-local overload code. |
| 19 | const ( |
| 20 | CodeParseError = -32700 |
| 21 | CodeInvalidRequest = -32600 |
| 22 | CodeMethodNotFound = -32601 |
| 23 | CodeInvalidParams = -32602 |
| 24 | CodeInternal = -32603 |
| 25 | // CodeServerBusy is a stable transport-local overload response. It is |
| 26 | // outside the JSON-RPC reserved range and intentionally carries no peer |
| 27 | // data. |
| 28 | CodeServerBusy = -32099 |
| 29 | ) |
| 30 | |
| 31 | // Transport bounds. |
| 32 | const ( |
| 33 | // maxConcurrentHandlers bounds inbound request and notification handlers. |
| 34 | maxConcurrentHandlers = 32 |
| 35 | // maxQueuedNotifications bounds the outbound notification queue. A full |
| 36 | // queue fails the connection rather than silently dropping a provider |
| 37 | // stream chunk (a dropped chunk would surface as a stream_gap host-side); |
| 38 | // this mirrors the host side's policy. |
| 39 | maxQueuedNotifications = 256 |
| 40 | ) |
| 41 | |
| 42 | // ResponseError is returned by outbound calls when the peer answers with a |
| 43 | // JSON-RPC error. Data remains raw so callers can decode ProtocolErrorData. |
| 44 | type ResponseError struct { |
| 45 | Code int |
| 46 | Message string |
| 47 | Data json.RawMessage |
| 48 | } |
| 49 | |
| 50 | func (e *ResponseError) Error() string { |
| 51 | if e == nil { |
| 52 | return "" |
| 53 | } |
| 54 | return e.Message |
| 55 | } |
| 56 | |
| 57 | // FrameTooLargeError reports a frame that violates the frozen NDJSON budget. |
| 58 | // Size and Limit include the trailing newline, matching the bytes sent over |
| 59 | // the transport. |
| 60 | type FrameTooLargeError struct { |
| 61 | Direction string |
| 62 | Size int |
| 63 | Limit int |
| 64 | } |
| 65 | |
| 66 | func (e *FrameTooLargeError) Error() string { |
| 67 | return fmt.Sprintf("extension: %s frame is %d bytes; limit is %d", e.Direction, e.Size, e.Limit) |
| 68 | } |
| 69 | |
| 70 | // rpcErrorObject is the JSON-RPC error object carried on the wire. |
| 71 | type rpcErrorObject struct { |
| 72 | Code int `json:"code"` |
| 73 | Message string `json:"message"` |
| 74 | Data json.RawMessage `json:"data,omitempty"` |
| 75 | } |
| 76 | |
| 77 | // requestHandler answers an inbound JSON-RPC request. |
| 78 | type requestHandler func(ctx context.Context, params json.RawMessage) (any, error) |
| 79 | |
| 80 | // notificationHandler handles an inbound JSON-RPC notification. |
| 81 | type notificationHandler func(ctx context.Context, params json.RawMessage) |
| 82 | |
| 83 | // deferredResult lets a request handler run cleanup only after a successful |
| 84 | // response write (for example, starting a provider stream pump once the |
| 85 | // stream/open acknowledgment is on the wire). |
| 86 | type deferredResult struct { |
| 87 | result any |
| 88 | after func() |
| 89 | } |
| 90 | |
| 91 | type rpcResult struct { |
| 92 | result json.RawMessage |
| 93 | err error |
| 94 | } |
| 95 | |
| 96 | type outbound struct { |
| 97 | JSONRPC string `json:"jsonrpc"` |
| 98 | ID json.RawMessage `json:"id,omitempty"` |
| 99 | Method string `json:"method,omitempty"` |
| 100 | Params json.RawMessage `json:"params,omitempty"` |
| 101 | Result json.RawMessage `json:"result,omitempty"` |
| 102 | Error *rpcErrorObject `json:"error,omitempty"` |
| 103 | } |
| 104 | |
| 105 | type inbound struct { |
| 106 | JSONRPC string `json:"jsonrpc"` |
| 107 | ID json.RawMessage `json:"id"` |
| 108 | Method string `json:"method"` |
| 109 | Params json.RawMessage `json:"params"` |
| 110 | Result json.RawMessage `json:"result"` |
| 111 | Error *rpcErrorObject `json:"error"` |
| 112 | } |
| 113 | |
| 114 | // conn is one bidirectional strict JSON-RPC 2.0 connection framed as NDJSON. |
| 115 | // The extension dialect narrows generic JSON-RPC: ids are integers only and |
| 116 | // params must be JSON objects. |
| 117 | type conn struct { |
| 118 | r io.Reader |
| 119 | w io.Writer |
| 120 | log *log.Logger |
| 121 | |
| 122 | wmu sync.Mutex |
| 123 | |
| 124 | nextID atomic.Int64 |
| 125 | |
| 126 | pmu sync.Mutex |
| 127 | pending map[int64]chan rpcResult |
| 128 | |
| 129 | reqH map[string]requestHandler |
| 130 | notH map[string]notificationHandler |
| 131 | |
| 132 | // beforeRequest and beforeNotification run synchronously on the read loop |
| 133 | // after strict frame validation and before dispatch, letting the |
| 134 | // handshake barrier observe wire arrival order. |
| 135 | beforeRequest func(method string) error |
| 136 | beforeNotification func(method string) error |
| 137 | |
| 138 | wg sync.WaitGroup |
| 139 | closeOnce sync.Once |
| 140 | closed chan struct{} |
| 141 | closeMu sync.Mutex |
| 142 | closeErr error |
| 143 | handlerSlots chan struct{} |
| 144 | notifyQueue chan []byte |
| 145 | } |
| 146 | |
| 147 | func newConn(r io.Reader, w io.Writer, logger *log.Logger) *conn { |
| 148 | return &conn{ |
| 149 | r: r, |
| 150 | w: w, |
| 151 | log: logger, |
| 152 | pending: make(map[int64]chan rpcResult), |
| 153 | reqH: make(map[string]requestHandler), |
| 154 | notH: make(map[string]notificationHandler), |
| 155 | closed: make(chan struct{}), |
| 156 | handlerSlots: make(chan struct{}, maxConcurrentHandlers), |
| 157 | notifyQueue: make(chan []byte, maxQueuedNotifications), |
| 158 | } |
| 159 | } |
| 160 | |
| 161 | // serve reads and dispatches frames until EOF, cancellation, or a |
| 162 | // framing/read error. In-flight handler contexts are cancelled when the |
| 163 | // transport ends. |
| 164 | func (c *conn) serve(ctx context.Context) error { |
| 165 | serveCtx, cancel := context.WithCancel(ctx) |
| 166 | defer cancel() |
| 167 | |
| 168 | c.wg.Add(1) |
| 169 | go c.serveOutboundNotifications() |
| 170 | |
| 171 | // Unblock a read parked on ctx cancellation: closing the reader is the |
| 172 | // only reliable way to interrupt it. |
| 173 | if closer, ok := c.r.(io.Closer); ok { |
| 174 | c.wg.Add(1) |
| 175 | go func() { |
| 176 | defer c.wg.Done() |
| 177 | select { |
| 178 | case <-serveCtx.Done(): |
| 179 | _ = closer.Close() |
| 180 | case <-c.closed: |
| 181 | } |
| 182 | }() |
| 183 | } |
| 184 | |
| 185 | br := bufio.NewReaderSize(c.r, 64<<10) |
| 186 | var loopErr error |
| 187 | for { |
| 188 | line, err := readLine(br, FrameBytes) |
| 189 | if len(line) > 0 { |
| 190 | c.dispatch(serveCtx, line) |
| 191 | } |
| 192 | if err != nil { |
| 193 | if !errors.Is(err, io.EOF) { |
| 194 | loopErr = err |
| 195 | } |
| 196 | break |
| 197 | } |
| 198 | if err := serveCtx.Err(); err != nil { |
| 199 | loopErr = err |
| 200 | break |
| 201 | } |
| 202 | } |
| 203 | |
| 204 | cancel() |
| 205 | close(c.notifyQueue) |
| 206 | c.wg.Wait() |
| 207 | // A connection that was failed or shut down deliberately makes the |
| 208 | // resulting read error a consequence, not the cause: report the recorded |
| 209 | // terminal error (nil for an orderly shutdown). Otherwise a parent ctx |
| 210 | // cancellation explains the forced reader close. |
| 211 | select { |
| 212 | case <-c.closed: |
| 213 | loopErr = c.recordedCloseError() |
| 214 | default: |
| 215 | if err := ctx.Err(); err != nil { |
| 216 | loopErr = err |
| 217 | } |
| 218 | } |
| 219 | c.shutdown(loopErr) |
| 220 | return loopErr |
| 221 | } |
| 222 | |
| 223 | // serveOutboundNotifications is the single ordered writer for fire-and-forget |
| 224 | // notifications (provider stream chunks). The queue is bounded; a full queue |
| 225 | // fails the connection instead of dropping a frame. |
| 226 | func (c *conn) serveOutboundNotifications() { |
| 227 | defer c.wg.Done() |
| 228 | for frame := range c.notifyQueue { |
| 229 | if err := c.writeFrame(frame); err != nil { |
| 230 | c.fail(err) |
| 231 | return |
| 232 | } |
| 233 | } |
| 234 | } |
| 235 | |
| 236 | func (c *conn) dispatch(ctx context.Context, line []byte) { |
| 237 | var in inbound |
| 238 | if err := json.Unmarshal(line, &in); err != nil { |
| 239 | if json.Valid(line) { |
| 240 | c.respondError(json.RawMessage("null"), CodeInvalidRequest, "invalid request", nil) |
| 241 | } else { |
| 242 | c.respondError(json.RawMessage("null"), CodeParseError, "parse error", nil) |
| 243 | } |
| 244 | return |
| 245 | } |
| 246 | if err := validateStrictFrame(line, &in); err != nil { |
| 247 | if c.log != nil { |
| 248 | c.log.Printf("extension: rejecting frame: %v", err) |
| 249 | } |
| 250 | c.respondError(responseIDForError(in.ID), CodeInvalidRequest, "invalid request", nil) |
| 251 | return |
| 252 | } |
| 253 | select { |
| 254 | case <-c.closed: |
| 255 | return |
| 256 | default: |
| 257 | } |
| 258 | hasID := len(in.ID) > 0 |
| 259 | switch { |
| 260 | case in.Method != "" && hasID: |
| 261 | if c.beforeRequest != nil { |
| 262 | if err := c.beforeRequest(in.Method); err != nil { |
| 263 | c.respondHandlerError(in.ID, err) |
| 264 | return |
| 265 | } |
| 266 | } |
| 267 | if !c.tryStartHandler() { |
| 268 | c.respondError(in.ID, CodeServerBusy, "server busy", nil) |
| 269 | return |
| 270 | } |
| 271 | c.wg.Add(1) |
| 272 | go func() { |
| 273 | defer c.finishHandler() |
| 274 | defer c.wg.Done() |
| 275 | c.serveRequest(ctx, in.ID, in.Method, in.Params) |
| 276 | }() |
| 277 | case in.Method != "" && !hasID: |
| 278 | if c.beforeNotification != nil { |
| 279 | if err := c.beforeNotification(in.Method); err != nil { |
| 280 | return |
| 281 | } |
| 282 | } |
| 283 | h := c.notH[in.Method] |
| 284 | if h == nil { |
| 285 | if c.log != nil { |
| 286 | c.log.Printf("extension: dropping notification for unhandled method %q", in.Method) |
| 287 | } |
| 288 | return |
| 289 | } |
| 290 | // No response is possible for a notification, so a saturated handler |
| 291 | // pool drops with a diagnostic rather than failing the connection. |
| 292 | if !c.tryStartHandler() { |
| 293 | if c.log != nil { |
| 294 | c.log.Printf("extension: dropping %q notification: handler pool saturated", in.Method) |
| 295 | } |
| 296 | return |
| 297 | } |
| 298 | c.wg.Add(1) |
| 299 | go func() { |
| 300 | defer c.finishHandler() |
| 301 | defer c.wg.Done() |
| 302 | c.runNotification(ctx, h, in.Params) |
| 303 | }() |
| 304 | case in.Method == "" && hasID: |
| 305 | c.resolve(&in) |
| 306 | default: |
| 307 | c.respondError(json.RawMessage("null"), CodeInvalidRequest, "invalid request", nil) |
| 308 | } |
| 309 | } |
| 310 | |
| 311 | func (c *conn) tryStartHandler() bool { |
| 312 | select { |
| 313 | case c.handlerSlots <- struct{}{}: |
| 314 | return true |
| 315 | default: |
| 316 | return false |
| 317 | } |
| 318 | } |
| 319 | |
| 320 | func (c *conn) finishHandler() { <-c.handlerSlots } |
| 321 | |
| 322 | // validateStrictFrame enforces the extension dialect of JSON-RPC 2.0: |
| 323 | // jsonrpc=="2.0", request/response shapes are mutually exclusive, ids are |
| 324 | // integers (or null), and params, when present, is a JSON object. |
| 325 | func validateStrictFrame(line []byte, in *inbound) error { |
| 326 | var members map[string]json.RawMessage |
| 327 | if err := json.Unmarshal(line, &members); err != nil { |
| 328 | return err |
| 329 | } |
| 330 | if in.JSONRPC != "2.0" { |
| 331 | return errors.New("jsonrpc must be 2.0") |
| 332 | } |
| 333 | _, hasID := members["id"] |
| 334 | _, hasMethod := members["method"] |
| 335 | _, hasParams := members["params"] |
| 336 | _, hasResult := members["result"] |
| 337 | _, hasError := members["error"] |
| 338 | if hasID && !validRPCID(in.ID) { |
| 339 | return errors.New("id must be an integer or null") |
| 340 | } |
| 341 | if hasMethod { |
| 342 | if in.Method == "" || hasResult || hasError { |
| 343 | return errors.New("invalid request shape") |
| 344 | } |
| 345 | if hasParams { |
| 346 | trimmed := bytes.TrimSpace(in.Params) |
| 347 | if len(trimmed) == 0 || trimmed[0] != '{' { |
| 348 | return errors.New("params must be a JSON object") |
| 349 | } |
| 350 | } |
| 351 | return nil |
| 352 | } |
| 353 | if !hasID || hasParams || hasResult == hasError { |
| 354 | return errors.New("invalid response shape") |
| 355 | } |
| 356 | if hasError { |
| 357 | if in.Error == nil { |
| 358 | return errors.New("invalid error object") |
| 359 | } |
| 360 | var errorMembers map[string]json.RawMessage |
| 361 | if err := json.Unmarshal(members["error"], &errorMembers); err != nil { |
| 362 | return errors.New("invalid error object") |
| 363 | } |
| 364 | if _, ok := errorMembers["code"]; !ok { |
| 365 | return errors.New("error code is required") |
| 366 | } |
| 367 | if _, ok := errorMembers["message"]; !ok { |
| 368 | return errors.New("error message is required") |
| 369 | } |
| 370 | } |
| 371 | return nil |
| 372 | } |
| 373 | |
| 374 | // validRPCID reports whether raw is an integer or null id. Unlike generic |
| 375 | // JSON-RPC, the extension protocol does not use string ids. |
| 376 | func validRPCID(raw json.RawMessage) bool { |
| 377 | raw = bytes.TrimSpace(raw) |
| 378 | if bytes.Equal(raw, []byte("null")) { |
| 379 | return true |
| 380 | } |
| 381 | if len(raw) == 0 { |
| 382 | return false |
| 383 | } |
| 384 | i := 0 |
| 385 | if raw[0] == '-' { |
| 386 | i++ |
| 387 | if i == len(raw) { |
| 388 | return false |
| 389 | } |
| 390 | } |
| 391 | if raw[i] == '0' && i+1 != len(raw) { |
| 392 | return false |
| 393 | } |
| 394 | for ; i < len(raw); i++ { |
| 395 | if raw[i] < '0' || raw[i] > '9' { |
| 396 | return false |
| 397 | } |
| 398 | } |
| 399 | return true |
| 400 | } |
| 401 | |
| 402 | // responseIDForError extracts the id member for an error response to a |
| 403 | // rejected frame, falling back to null when the id is absent or invalid. |
| 404 | func responseIDForError(raw json.RawMessage) json.RawMessage { |
| 405 | if len(bytes.TrimSpace(raw)) == 0 || !validRPCID(raw) { |
| 406 | return json.RawMessage("null") |
| 407 | } |
| 408 | return raw |
| 409 | } |
| 410 | |
| 411 | func (c *conn) serveRequest(ctx context.Context, id json.RawMessage, method string, params json.RawMessage) { |
| 412 | h := c.reqH[method] |
| 413 | if h == nil { |
| 414 | notFound := MustProtocolError(ErrUnknownMethod) |
| 415 | spec := frozenErrorSpecs[ErrUnknownMethod] |
| 416 | c.respondError(id, spec.Code, "method not found: "+method, ProtocolErrorData{Reason: notFound.Reason, Retryable: spec.Retryable}) |
| 417 | return |
| 418 | } |
| 419 | result, err := c.runHandler(ctx, h, params) |
| 420 | if err != nil { |
| 421 | c.respondHandlerError(id, err) |
| 422 | return |
| 423 | } |
| 424 | var after func() |
| 425 | if deferred, ok := result.(deferredResult); ok { |
| 426 | result = deferred.result |
| 427 | after = deferred.after |
| 428 | } |
| 429 | raw, err := json.Marshal(result) |
| 430 | if err != nil { |
| 431 | c.respondError(id, CodeInternal, "marshal result: "+err.Error(), nil) |
| 432 | return |
| 433 | } |
| 434 | writeErr := c.write(outbound{JSONRPC: "2.0", ID: id, Result: raw}) |
| 435 | if writeErr != nil { |
| 436 | var tooLarge *FrameTooLargeError |
| 437 | if errors.As(writeErr, &tooLarge) { |
| 438 | c.respondError(id, CodeInternal, "response exceeds frame size limit", nil) |
| 439 | return |
| 440 | } |
| 441 | c.fail(writeErr) |
| 442 | return |
| 443 | } |
| 444 | if after != nil { |
| 445 | c.runAfterWrite(after) |
| 446 | } |
| 447 | } |
| 448 | |
| 449 | // runNotification executes one notification handler, converting a panic into |
| 450 | // a diagnostic so the read loop and the connection survive. |
| 451 | func (c *conn) runNotification(ctx context.Context, h notificationHandler, params json.RawMessage) { |
| 452 | defer func() { |
| 453 | if recovered := recover(); recovered != nil && c.log != nil { |
| 454 | c.log.Printf("extension: notification handler panic: %v", recovered) |
| 455 | } |
| 456 | }() |
| 457 | h(ctx, params) |
| 458 | } |
| 459 | |
| 460 | // runHandler executes one request handler, converting a panic into the frozen |
| 461 | // internal error so the read loop and the connection survive. |
| 462 | func (c *conn) runHandler(ctx context.Context, h requestHandler, params json.RawMessage) (result any, err error) { |
| 463 | defer func() { |
| 464 | if recovered := recover(); recovered != nil { |
| 465 | if c.log != nil { |
| 466 | c.log.Printf("extension: handler panic: %v", recovered) |
| 467 | } |
| 468 | result = nil |
| 469 | err = MustProtocolError(ErrInternal) |
| 470 | } |
| 471 | }() |
| 472 | return h(ctx, params) |
| 473 | } |
| 474 | |
| 475 | func (c *conn) runAfterWrite(after func()) { |
| 476 | defer func() { |
| 477 | if recovered := recover(); recovered != nil { |
| 478 | c.fail(fmt.Errorf("extension: after-response callback panic: %v", recovered)) |
| 479 | } |
| 480 | }() |
| 481 | after() |
| 482 | } |
| 483 | |
| 484 | func (c *conn) respondHandlerError(id json.RawMessage, err error) { |
| 485 | // A fatal error (a failed handshake) is answered first and only then ends |
| 486 | // the connection, so the peer sees the reason. |
| 487 | var fatal *fatalError |
| 488 | isFatal := errors.As(err, &fatal) |
| 489 | respond := err |
| 490 | if isFatal { |
| 491 | respond = fatal.err |
| 492 | } |
| 493 | var protocolErr *ProtocolError |
| 494 | if errors.As(respond, &protocolErr) { |
| 495 | spec := frozenErrorSpecs[protocolErr.Reason] |
| 496 | message := protocolErr.Message |
| 497 | if message == "" { |
| 498 | message = spec.Message |
| 499 | } |
| 500 | c.respondError(id, spec.Code, message, ProtocolErrorData{Reason: protocolErr.Reason, Retryable: spec.Retryable}) |
| 501 | } else { |
| 502 | if c.log != nil { |
| 503 | c.log.Printf("extension: handler error: %v", respond) |
| 504 | } |
| 505 | // Unknown handler errors never leak internals onto the wire: the peer |
| 506 | // sees the frozen internal error, the diagnostic goes to the logger. |
| 507 | spec := frozenErrorSpecs[ErrInternal] |
| 508 | c.respondError(id, spec.Code, spec.Message, ProtocolErrorData{Reason: ErrInternal, Retryable: spec.Retryable}) |
| 509 | } |
| 510 | if isFatal { |
| 511 | c.fail(fatal.err) |
| 512 | } |
| 513 | } |
| 514 | |
| 515 | func (c *conn) resolve(in *inbound) { |
| 516 | id, err := strconv.ParseInt(string(in.ID), 10, 64) |
| 517 | if err != nil { |
| 518 | return |
| 519 | } |
| 520 | c.pmu.Lock() |
| 521 | ch := c.pending[id] |
| 522 | delete(c.pending, id) |
| 523 | c.pmu.Unlock() |
| 524 | if ch == nil { |
| 525 | return |
| 526 | } |
| 527 | if in.Error != nil { |
| 528 | ch <- rpcResult{err: &ResponseError{Code: in.Error.Code, Message: in.Error.Message, Data: in.Error.Data}} |
| 529 | return |
| 530 | } |
| 531 | ch <- rpcResult{result: in.Result} |
| 532 | } |
| 533 | |
| 534 | // notify queues a fire-and-forget notification. Notifications travel through |
| 535 | // one bounded FIFO queue so provider stream chunks stay ordered; a full queue |
| 536 | // fails the connection rather than silently dropping a frame (mirroring the |
| 537 | // host side). A marshaled frame beyond FrameBytes fails only that call. |
| 538 | func (c *conn) notify(method string, params any) error { |
| 539 | raw, err := json.Marshal(params) |
| 540 | if err != nil { |
| 541 | return err |
| 542 | } |
| 543 | var buf bytes.Buffer |
| 544 | enc := json.NewEncoder(&buf) |
| 545 | enc.SetEscapeHTML(false) |
| 546 | if err := enc.Encode(outbound{JSONRPC: "2.0", Method: method, Params: raw}); err != nil { |
| 547 | return err |
| 548 | } |
| 549 | if buf.Len() > FrameBytes { |
| 550 | return &FrameTooLargeError{Direction: "outbound", Size: buf.Len(), Limit: FrameBytes} |
| 551 | } |
| 552 | select { |
| 553 | case <-c.closed: |
| 554 | return c.closedError() |
| 555 | default: |
| 556 | } |
| 557 | select { |
| 558 | case c.notifyQueue <- buf.Bytes(): |
| 559 | return nil |
| 560 | default: |
| 561 | err := fmt.Errorf("extension: outbound notification queue overflow (%d)", maxQueuedNotifications) |
| 562 | c.fail(err) |
| 563 | return err |
| 564 | } |
| 565 | } |
| 566 | |
| 567 | // call sends a request and waits for its response, cancellation, or closure. |
| 568 | func (c *conn) call(ctx context.Context, method string, params any) (json.RawMessage, error) { |
| 569 | raw, err := json.Marshal(params) |
| 570 | if err != nil { |
| 571 | return nil, err |
| 572 | } |
| 573 | id := c.nextID.Add(1) |
| 574 | ch := make(chan rpcResult, 1) |
| 575 | c.pmu.Lock() |
| 576 | select { |
| 577 | case <-c.closed: |
| 578 | c.pmu.Unlock() |
| 579 | return nil, c.closedError() |
| 580 | default: |
| 581 | } |
| 582 | c.pending[id] = ch |
| 583 | c.pmu.Unlock() |
| 584 | defer func() { |
| 585 | c.pmu.Lock() |
| 586 | delete(c.pending, id) |
| 587 | c.pmu.Unlock() |
| 588 | }() |
| 589 | |
| 590 | idRaw, _ := json.Marshal(id) |
| 591 | if err := c.write(outbound{JSONRPC: "2.0", ID: idRaw, Method: method, Params: raw}); err != nil { |
| 592 | var tooLarge *FrameTooLargeError |
| 593 | if !errors.As(err, &tooLarge) { |
| 594 | c.fail(err) |
| 595 | } |
| 596 | return nil, err |
| 597 | } |
| 598 | select { |
| 599 | case res := <-ch: |
| 600 | return res.result, res.err |
| 601 | case <-ctx.Done(): |
| 602 | return nil, ctx.Err() |
| 603 | case <-c.closed: |
| 604 | return nil, c.closedError() |
| 605 | } |
| 606 | } |
| 607 | |
| 608 | func (c *conn) write(m outbound) error { |
| 609 | var buf bytes.Buffer |
| 610 | enc := json.NewEncoder(&buf) |
| 611 | enc.SetEscapeHTML(false) |
| 612 | if err := enc.Encode(m); err != nil { |
| 613 | return err |
| 614 | } |
| 615 | if buf.Len() > FrameBytes { |
| 616 | return &FrameTooLargeError{Direction: "outbound", Size: buf.Len(), Limit: FrameBytes} |
| 617 | } |
| 618 | return c.writeFrame(buf.Bytes()) |
| 619 | } |
| 620 | |
| 621 | func (c *conn) writeFrame(frame []byte) error { |
| 622 | c.wmu.Lock() |
| 623 | defer c.wmu.Unlock() |
| 624 | for len(frame) > 0 { |
| 625 | n, err := c.w.Write(frame) |
| 626 | if err != nil { |
| 627 | return err |
| 628 | } |
| 629 | if n == 0 { |
| 630 | return io.ErrShortWrite |
| 631 | } |
| 632 | frame = frame[n:] |
| 633 | } |
| 634 | return nil |
| 635 | } |
| 636 | |
| 637 | func (c *conn) respondError(id json.RawMessage, code int, message string, data any) { |
| 638 | var raw json.RawMessage |
| 639 | if data != nil { |
| 640 | encoded, err := json.Marshal(data) |
| 641 | if err != nil { |
| 642 | code = CodeInternal |
| 643 | message = "marshal error data: " + err.Error() |
| 644 | } else if string(encoded) != "null" { |
| 645 | raw = encoded |
| 646 | } |
| 647 | } |
| 648 | if err := c.write(outbound{JSONRPC: "2.0", ID: id, Error: &rpcErrorObject{Code: code, Message: message, Data: raw}}); err != nil { |
| 649 | var tooLarge *FrameTooLargeError |
| 650 | if !errors.As(err, &tooLarge) { |
| 651 | c.fail(err) |
| 652 | } |
| 653 | } |
| 654 | } |
| 655 | |
| 656 | func (c *conn) fail(err error) { |
| 657 | if err == nil { |
| 658 | return |
| 659 | } |
| 660 | c.shutdown(err) |
| 661 | if closer, ok := c.r.(io.Closer); ok { |
| 662 | _ = closer.Close() |
| 663 | } |
| 664 | } |
| 665 | |
| 666 | func (c *conn) closedError() error { |
| 667 | c.closeMu.Lock() |
| 668 | defer c.closeMu.Unlock() |
| 669 | if c.closeErr != nil { |
| 670 | return c.closeErr |
| 671 | } |
| 672 | return errors.New("extension: connection closed") |
| 673 | } |
| 674 | |
| 675 | // recordedCloseError returns the terminal error recorded at shutdown, which |
| 676 | // may be nil for an orderly close. |
| 677 | func (c *conn) recordedCloseError() error { |
| 678 | c.closeMu.Lock() |
| 679 | defer c.closeMu.Unlock() |
| 680 | return c.closeErr |
| 681 | } |
| 682 | |
| 683 | func (c *conn) shutdown(err error) { |
| 684 | c.closeOnce.Do(func() { |
| 685 | c.closeMu.Lock() |
| 686 | c.closeErr = err |
| 687 | c.closeMu.Unlock() |
| 688 | close(c.closed) |
| 689 | c.pmu.Lock() |
| 690 | for id, ch := range c.pending { |
| 691 | pendingErr := err |
| 692 | if pendingErr == nil { |
| 693 | pendingErr = errors.New("extension: connection closed") |
| 694 | } |
| 695 | ch <- rpcResult{err: pendingErr} |
| 696 | delete(c.pending, id) |
| 697 | } |
| 698 | c.pmu.Unlock() |
| 699 | }) |
| 700 | } |
| 701 | |
| 702 | // readLine reads one NDJSON frame, enforcing the byte budget across bufio |
| 703 | // refills and trimming the trailing line ending. |
| 704 | func readLine(br *bufio.Reader, maxBytes int) ([]byte, error) { |
| 705 | var buf []byte |
| 706 | for { |
| 707 | chunk, err := br.ReadSlice('\n') |
| 708 | buf = append(buf, chunk...) |
| 709 | if maxBytes > 0 && len(buf) > maxBytes { |
| 710 | return nil, &FrameTooLargeError{Direction: "inbound", Size: len(buf), Limit: maxBytes} |
| 711 | } |
| 712 | if errors.Is(err, bufio.ErrBufferFull) { |
| 713 | continue |
| 714 | } |
| 715 | n := len(buf) |
| 716 | for n > 0 && (buf[n-1] == '\n' || buf[n-1] == '\r') { |
| 717 | n-- |
| 718 | } |
| 719 | return trimSpaceBytes(buf[:n]), err |
| 720 | } |
| 721 | } |
| 722 | |
| 723 | func trimSpaceBytes(b []byte) []byte { |
| 724 | i, j := 0, len(b) |
| 725 | for i < j && isSpaceByte(b[i]) { |
| 726 | i++ |
| 727 | } |
| 728 | for j > i && isSpaceByte(b[j-1]) { |
| 729 | j-- |
| 730 | } |
| 731 | return b[i:j] |
| 732 | } |
| 733 | |
| 734 | func isSpaceByte(c byte) bool { return c == ' ' || c == '\t' || c == '\n' || c == '\r' } |
| 735 |