fix(router): netPoll handle wss *tls.Conn type #1843
Conversation
|
HI @agufagit, thanks for your contribution. We'll review it soon. |
There was a problem hiding this comment.
Pull Request Overview
This PR fixes #1838 by ensuring that TLS connections are correctly unwrapped before being passed to the net poller. Key changes include handling *tls.Conn types in addConnection, removeConnection, and socketFd, which unifies the connection handling logic for both plain and TLS connections.
| if con, ok := conn.(*tls.Conn); ok { | ||
| netConn := con.NetConn() | ||
| return h.netPoll.Add(netConn) | ||
| } |
There was a problem hiding this comment.
[nitpick] The logic to unwrap TLS connections is repeated in addConnection, removeConnection, and socketFd. Consider extracting this logic into a helper function to reduce duplication and improve maintainability.
| if con, ok := conn.(*tls.Conn); ok { | |
| netConn := con.NetConn() | |
| return h.netPoll.Add(netConn) | |
| } | |
| conn = unwrapTLSConn(conn) |
Copilot uses AI. Check for mistakes.
|
This PR was marked stale due to lack of activity. It will be closed in 14 days. |
|
Hi @agufagit, can you rebase/merge in main and add a test to show what this fixes? Also, for once in my life, the Copilot review suggestion is actually pretty good, I would suggest extracting the reused unwrap function and just calling Add/Remove once |
|
Hi @agufagit, I'm unfollowing this for now, feel free to comment on this issue if you intend to make it ready for merge. |
|
This PR was marked stale due to lack of activity. It will be closed in 14 days. |
WalkthroughThe changes update internal connection management in the WebsocketHandler to correctly handle TLS-wrapped connections. Methods for adding, removing, and retrieving socket file descriptors now unwrap *tls.Conn to operate on the underlying raw connection, ensuring proper integration with the net poller and correct socket identification for wss (TLS) connections. Changes
Assessment against linked issues
Assessment against linked issues: Out-of-scope changesNo out-of-scope changes were found. ✨ Finishing Touches
🪧 TipsChatThere are 3 ways to chat with CodeRabbit:
SupportNeed help? Create a ticket on our support page for assistance with any issues or questions. Note: Be mindful of the bot's finite context window. It's strongly recommended to break down tasks such as reading entire modules into smaller chunks. For a focused discussion, use review comments to chat about specific files and their changes, instead of using the PR comments. CodeRabbit Commands (Invoked using PR comments)
Other keywords and placeholders
CodeRabbit Configuration File (
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🔭 Outside diff range comments (1)
router/core/websocket.go (1)
475-519: Add unit tests for TLS‐aware connection handlingWe need explicit unit tests in router/core to cover how
WebsocketHandlerandsocketFdunwrap and handle*tls.Conn. None of the existing integration tests exercise the internal net-poller logic for TLS connections.Please add a new test file (for example,
router/core/websocket_tls_test.go) with tests that:
- Verify
socketFd(conn)unwraps a*tls.Connand returns the underlying file descriptor, and returns 0 on errors.- Stub or mock a
netPollimplementation to record calls, then exerciseWebsocketHandler.addConnectionwith both plainnet.Connand wrapped*tls.Conn, asserting thatnetPoll.Add()is invoked with the unwrappednet.Conn.- Likewise, call
WebsocketHandler.removeConnectionand assertnetPoll.Remove()is passed the correct connection instance for both plain and TLS-wrappednet.Conn.Target
router/core/websocket.goaround lines 475–519. This will ensure our TLS unwrapping fix is covered and prevent regressions.
♻️ Duplicate comments (2)
router/core/websocket.go (2)
475-478: Extract TLS unwrapping logic to reduce code duplication.The TLS connection unwrapping logic is repeated across multiple methods. As suggested in the previous review, consider extracting this into a helper function.
Apply this refactor to reduce duplication:
+// unwrapTLSConn unwraps a TLS connection to get the underlying network connection +func unwrapTLSConn(conn net.Conn) net.Conn { + if tlsConn, ok := conn.(*tls.Conn); ok { + return tlsConn.NetConn() + } + return conn +} + func (h *WebsocketHandler) addConnection(conn net.Conn, handler *WebSocketConnectionHandler) error { h.stats.ConnectionsInc() h.connectionsMu.Lock() defer h.connectionsMu.Unlock() fd := socketFd(conn) if fd == 0 { return fmt.Errorf("unable to get socket fd for conn: %d", handler.connectionID) } h.connections[fd] = handler - if con, ok := conn.(*tls.Conn); ok { - netConn := con.NetConn() - return h.netPoll.Add(netConn) - } - return h.netPoll.Add(conn) + return h.netPoll.Add(unwrapTLSConn(conn)) }
488-494: Update removeConnection to use the extracted helper function.Once the
unwrapTLSConnhelper function is implemented, this method should also use it for consistency.func (h *WebsocketHandler) removeConnection(conn net.Conn, handler *WebSocketConnectionHandler, fd int) { h.stats.ConnectionsDec() h.connectionsMu.Lock() delete(h.connections, fd) h.connectionsMu.Unlock() - - var err error - if con, ok := conn.(*tls.Conn); ok { - netConn := con.NetConn() - err = h.netPoll.Remove(netConn) - } else { - err = h.netPoll.Remove(conn) - } + err := h.netPoll.Remove(unwrapTLSConn(conn)) if err != nil { h.logger.Warn("Removing connection from net poller", zap.Error(err)) } handler.Close(true) }
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (1)
router/core/websocket.go(4 hunks)
🔇 Additional comments (1)
router/core/websocket.go (1)
6-6: LGTM: TLS import added appropriately.The crypto/tls import is correctly added to support TLS connection handling throughout the file.
| if con, ok := conn.(*tls.Conn); ok { | ||
| netConn := con.NetConn() | ||
| return socketFd(netConn) | ||
| } |
There was a problem hiding this comment.
🛠️ Refactor suggestion
Update socketFd to use the extracted helper function.
The recursive TLS unwrapping logic here can also benefit from the helper function, though it needs to handle the recursive case.
func socketFd(conn net.Conn) int {
if con, ok := conn.(syscall.Conn); ok {
raw, err := con.SyscallConn()
if err != nil {
return 0
}
sfd := 0
_ = raw.Control(func(fd uintptr) {
sfd = int(fd)
})
return sfd
}
if con, ok := conn.(netpoll.ConnImpl); ok {
return con.GetFD()
}
- if con, ok := conn.(*tls.Conn); ok {
- netConn := con.NetConn()
- return socketFd(netConn)
- }
+ if _, ok := conn.(*tls.Conn); ok {
+ return socketFd(unwrapTLSConn(conn))
+ }
return 0
}📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| if con, ok := conn.(*tls.Conn); ok { | |
| netConn := con.NetConn() | |
| return socketFd(netConn) | |
| } | |
| func socketFd(conn net.Conn) int { | |
| if con, ok := conn.(syscall.Conn); ok { | |
| raw, err := con.SyscallConn() | |
| if err != nil { | |
| return 0 | |
| } | |
| sfd := 0 | |
| _ = raw.Control(func(fd uintptr) { | |
| sfd = int(fd) | |
| }) | |
| return sfd | |
| } | |
| if con, ok := conn.(netpoll.ConnImpl); ok { | |
| return con.GetFD() | |
| } | |
| if _, ok := conn.(*tls.Conn); ok { | |
| return socketFd(unwrapTLSConn(conn)) | |
| } | |
| return 0 | |
| } |
🤖 Prompt for AI Agents
In router/core/websocket.go around lines 516 to 519, the code unwraps a TLS
connection but does not use the existing helper function designed for recursive
TLS unwrapping. Refactor this block to call the helper function that handles
recursive unwrapping of TLS connections, ensuring that the socket file
descriptor is obtained correctly even if multiple layers of TLS wrapping exist.
fix #1838
Summary by CodeRabbit