[relay] Replace net.Conn with context-aware Conn interface#5770
Merged
Conversation
…transports Introduce a listener.Conn interface with context-based Read/Write methods, replacing net.Conn throughout the relay server. This enables proper timeout propagation (e.g. handshake timeout) without goroutine-based workarounds and removes unused LocalAddr/SetDeadline methods from WS and QUIC conns.
Contributor
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Pro Run ID: 📒 Files selected for processing (1)
🚧 Files skipped from review as they are similar to previous changes (1)
📝 WalkthroughWalkthroughThis PR replaces net.Conn usage with a new context-aware listener.Conn across relay, listener implementations (QUIC, WebSocket), handshake and peer code, threads context/timeouts through handshake and peer I/O, and adjusts listener APIs and server wiring accordingly. (47 words) Changes
Sequence Diagram(s)sequenceDiagram
participant Client as External Client
participant Listener as Listener (WS / QUIC)
participant Relay as Relay
participant HS as Handshake
participant Peer as Peer
Client->>Listener: Establish connection
Listener->>Relay: Accept(conn listener.Conn)
Relay->>HS: Start handshake (with ctx timeout)
HS->>Listener: Read(ctx, buf)
Listener-->>HS: Handshake request
HS->>HS: Validate request
HS->>Listener: Write(ctx, response)
Listener-->>HS: Ack
Relay->>Peer: NewPeer(conn listener.Conn)
Peer->>Listener: Read(ctx, loop)
Listener-->>Peer: Message bytes
Peer->>Listener: Write(ctx, response)
Listener-->>Peer: Ack
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes Possibly related PRs
Suggested reviewers
Poem
🚥 Pre-merge checks | ✅ 2 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
pappz
changed the title
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 3
🧹 Nitpick comments (1)
relay/server/peer.go (1)
35-36:ctxandctxCancelare nil untilWork()is called.The
ctxandctxCancelfields are not initialized inNewPeerbut only later inWork()(line 62). If any code path accessesp.ctxbeforeWork()is called, it will be nil, potentially causing a panic.Looking at the call sites in
relay.go, the flow is:
NewPeer()creates the peerr.store.AddPeer(peer)adds it to the storer.notifier.PeerCameOnline(peer.ID())notifiespeer.Work()is started in a goroutineh.handshakeResponse(hsCtx)may callpeer.Close()on failureIf
Close()were to usep.ctx, it would panic. CurrentlyClose()doesn't usep.ctx, so this is safe. However, this is a subtle invariant that could break with future changes.Consider initializing context in NewPeer for defensive coding
func NewPeer(metrics *metrics.Metrics, id messages.PeerID, conn listener.Conn, store *store.Store, notifier *store.PeerNotifier) *Peer { + ctx, ctxCancel := context.WithCancel(context.Background()) p := &Peer{ metrics: metrics, log: log.WithField("peer_id", id.String()), id: id, conn: conn, store: store, notifier: notifier, + ctx: ctx, + ctxCancel: ctxCancel, } return p }Then in
Work(), you could either skip re-initialization or replace the context if needed.Also applies to: 45-56
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@relay/server/peer.go` around lines 35 - 36, NewPeer currently leaves the peer fields ctx and ctxCancel nil which is fragile; initialize them in NewPeer by creating a background context and cancel func (e.g., ctx, cancel := context.WithCancel(context.Background()) and assign to p.ctx and p.ctxCancel) so any early callers (including Close) can safely reference them, then update Work to not assume ctx is nil (either skip reinitialization or derive a child context from the existing p.ctx) and make Close defensively call p.ctxCancel only after a nil-check to avoid double cancels.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@relay/server/listener/quic/conn.go`:
- Around line 34-38: Conn.Write currently ignores the provided context (breaking
the listener.Conn contract) and may send QUIC datagrams after cancellation;
modify Conn.Write to honor ctx by checking ctx.Done()/ctx.Err() before
attempting SendDatagram and by applying the same timeout/cancellation semantics
as the WS implementation (e.g., respect write timeout or use select on
ctx.Done() with a send deadline), then call c.session.SendDatagram only if the
context is still active and translate errors via c.remoteCloseErrHandling as
before.
In `@relay/server/listener/ws/conn.go`:
- Around line 55-60: The current Conn.Write implementation returns len(b) even
when c.Conn.Write fails; modify Conn.Write (the method on Conn that uses
writeTimeout and calls c.Conn.Write with websocket.MessageBinary) to check the
error from c.Conn.Write and return 0, err on failure, otherwise return len(b),
nil, keeping the existing context timeout/cancel logic intact so behavior
matches relay/server/listener/quic/conn.go.
In `@relay/server/peer.go`:
- Around line 226-231: The crash occurs because dp.ctx can be nil when another
peer calls dp.Write (which calls context.WithTimeout) before peer.Work runs; fix
by ensuring the peer's context is created before the peer is added to the store
— initialize the context in the Peer constructor (so dp.ctx is non-nil) or,
alternatively, stop starting peer.Work in a goroutine and call peer.Work
synchronously immediately after creating the Peer so dp.ctx is set before
Add/Store; update the code paths that create Peers (where you currently add the
peer to the store) to either (a) set dp.ctx inside NewPeer or (b) call
peer.Work() inline instead of go peer.Work().
---
Nitpick comments:
In `@relay/server/peer.go`:
- Around line 35-36: NewPeer currently leaves the peer fields ctx and ctxCancel
nil which is fragile; initialize them in NewPeer by creating a background
context and cancel func (e.g., ctx, cancel :=
context.WithCancel(context.Background()) and assign to p.ctx and p.ctxCancel) so
any early callers (including Close) can safely reference them, then update Work
to not assume ctx is nil (either skip reinitialization or derive a child context
from the existing p.ctx) and make Close defensively call p.ctxCancel only after
a nil-check to avoid double cancels.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 6a719b5b-d5a3-4140-81b0-1df1df289b49
📒 Files selected for processing (11)
combined/cmd/root.gorelay/server/handshake.gorelay/server/listener/conn.gorelay/server/listener/listener.gorelay/server/listener/quic/conn.gorelay/server/listener/quic/listener.gorelay/server/listener/ws/conn.gorelay/server/listener/ws/listener.gorelay/server/peer.gorelay/server/relay.gorelay/server/server.go
💤 Files with no reviewable changes (1)
- relay/server/listener/listener.go
Integrate context creation (`context.WithCancel`) directly in `NewPeer` and remove redundant initialization in `Work`. Add `ctxCancel` calls to ensure context is properly canceled during `Close` operations.
|
lixmal
approved these changes
Apr 7, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Describe your changes
Replace net.Conn with context-aware Conn interface
Introduce a listener.Conn interface with context-based Read/Write methods, replacing net.Conn throughout the relay server. This enables proper timeout propagation (e.g. handshake timeout) without goroutine-based workarounds and removes unused LocalAddr/SetDeadline methods from WS and QUIC conns.
Issue ticket number and link
Stack
Checklist
Documentation
Select exactly one:
Docs PR URL (required if "docs added" is checked)
Paste the PR link from https://github.com/netbirdio/docs here:
https://github.com/netbirdio/docs/pull/__
Summary by CodeRabbit
Refactor
Bug Fixes