Fix sse tunneling #190
Fix sse tunneling #190
Conversation
WalkthroughThe changes add specialized handling for WebSocket protocol upgrades (HTTP 101 Switching Protocols) and Server-Sent Events (SSE) streaming within the HTTP tunnel path. WebSocket connections are detected and immediately switched to TCP tunnel mode, while SSE streams are forwarded live without buffering. Regular HTTP responses maintain existing behavior. Changes
Sequence Diagram(s)sequenceDiagram
participant Client
participant Tunnel
participant Remote
rect rgb(240, 248, 255)
Note over Tunnel: HTTP Response Received
end
alt WebSocket Upgrade (101)
Tunnel->>Tunnel: Check StatusCode == 101
Tunnel->>Client: Write upgrade response
Tunnel->>Client: Flush
Tunnel->>Tunnel: Switch to TCP tunnel mode
Tunnel->>Remote: Establish TCP tunnel
else SSE Stream (Content-Type: text/event-stream)
Tunnel->>Tunnel: Detect SSE content type
Tunnel->>Client: Write status line & headers
Tunnel->>Client: Write header terminator
loop Stream Data
Remote->>Tunnel: Read chunk
Tunnel->>Client: Write chunk
Tunnel->>Client: Flush immediately
end
else Regular HTTP Response
Tunnel->>Tunnel: Read full response body
Tunnel->>Client: Write status, headers, body
end
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes Multiple distinct code paths (WebSocket upgrade, SSE streaming, regular HTTP) with conditional logic, streaming behavior with per-chunk flushing, header manipulation, and error handling require careful verification of protocol-specific edge cases and state transitions. Poem
Pre-merge checks and finishing touches❌ Failed checks (1 inconclusive)
✅ Passed checks (4 passed)
✨ Finishing touches
🧪 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 |
Deploying portr-docs with
|
| Latest commit: |
ff2f946
|
| Status: | ✅ Deploy successful! |
| Preview URL: | https://4be2d603.portr-docs.pages.dev |
| Branch Preview URL: | https://sse-ws-tunnel.portr-docs.pages.dev |
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (2)
tunnel/internal/client/ssh/ssh.go (2)
13-13: Import looks fine; consider addingmimefor robust SSE detection.If you adopt the SSE media-type parsing suggestion below, also import
mime. Otherwise, keepingstringsis fine.
314-360: SSE streaming path LGTM; tighten detection and close body on exit.This likely resolves #188 by streaming without buffering. Make detection resilient and ensure cleanup.
Minimal patch:
@@ - contentType := response.Header.Get("Content-Type") - if strings.Contains(contentType, "text/event-stream") { + // Robust SSE detection (case-insensitive, ignores parameters) + // Requires: import "mime" + ct, _, _ := mime.ParseMediaType(response.Header.Get("Content-Type")) + if strings.EqualFold(ct, "text/event-stream") { @@ - // Stream the body with immediate flushing for real-time delivery + // Stream the body with immediate flushing for real-time delivery + // Ensure the body is closed when we exit this branch + defer response.Body.Close() buf := make([]byte, 32*1024) // 32KB buffer for { n, err := response.Body.Read(buf) if n > 0 { _, writeErr := srcWriter.Write(buf[:n]) if writeErr != nil { if s.config.Debug { s.logDebug("Failed to write SSE data", writeErr) } return } // Flush immediately to ensure real-time streaming srcWriter.Flush() } if err != nil { if err != io.EOF && s.config.Debug { s.logDebug("SSE stream ended", err) } break } } return }And add the import:
import ( @@ - "strings" + "strings" + "mime" )Optional: confirm if skipping
s.logHttpRequestfor SSE is intentional to avoid DB bloat; otherwise, consider metadata-only logging (status, headers, no body).
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (1)
tunnel/internal/client/ssh/ssh.go(2 hunks)
🔇 Additional comments (1)
tunnel/internal/client/ssh/ssh.go (1)
361-361: Regular HTTP path unchanged — OK.No regressions spotted in the non-streaming branch.
| // Handle WebSocket upgrades and SSE streams with TCP tunneling | ||
| if response.StatusCode == http.StatusSwitchingProtocols { | ||
| // WebSocket upgrade - write response headers and switch to TCP tunneling | ||
| err = response.Write(srcWriter) | ||
| if err != nil { | ||
| if s.config.Debug { | ||
| s.logDebug("Failed to write WebSocket upgrade response", err) | ||
| } | ||
| return | ||
| } | ||
| srcWriter.Flush() | ||
| s.tcpTunnel(src, dst) | ||
| return | ||
| } | ||
|
|
There was a problem hiding this comment.
WebSocket upgrade: possible data loss due to buffered bytes not forwarded.
After http.ReadResponse, dstReader may already hold server WS frames; similarly, srcReader may hold client bytes beyond the HTTP request. Switching straight to tcpTunnel(src, dst) bypasses these buffers (io.Copy reads from net.Conn, not the bufio.Reader), dropping bytes and corrupting the stream.
Apply this patch to drain any buffered bytes on both sides before tunneling:
if response.StatusCode == http.StatusSwitchingProtocols {
// WebSocket upgrade - write response headers and switch to TCP tunneling
err = response.Write(srcWriter)
if err != nil {
if s.config.Debug {
s.logDebug("Failed to write WebSocket upgrade response", err)
}
return
}
srcWriter.Flush()
+ // Drain any bytes already buffered post-handshake to avoid loss when switching to raw TCP
+ if n := dstReader.Buffered(); n > 0 {
+ buf := make([]byte, n)
+ if _, err := io.ReadFull(dstReader, buf); err == nil {
+ if _, err := srcWriter.Write(buf); err != nil {
+ if s.config.Debug {
+ s.logDebug("Failed to flush buffered server bytes on WS upgrade", err)
+ }
+ return
+ }
+ srcWriter.Flush()
+ }
+ }
+ if n := srcReader.Buffered(); n > 0 {
+ buf := make([]byte, n)
+ if _, err := io.ReadFull(srcReader, buf); err == nil {
+ if _, err := dstWriter.Write(buf); err != nil {
+ if s.config.Debug {
+ s.logDebug("Failed to flush buffered client bytes on WS upgrade", err)
+ }
+ return
+ }
+ dstWriter.Flush()
+ }
+ }
s.tcpTunnel(src, dst)
return
}📝 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.
| // Handle WebSocket upgrades and SSE streams with TCP tunneling | |
| if response.StatusCode == http.StatusSwitchingProtocols { | |
| // WebSocket upgrade - write response headers and switch to TCP tunneling | |
| err = response.Write(srcWriter) | |
| if err != nil { | |
| if s.config.Debug { | |
| s.logDebug("Failed to write WebSocket upgrade response", err) | |
| } | |
| return | |
| } | |
| srcWriter.Flush() | |
| s.tcpTunnel(src, dst) | |
| return | |
| } | |
| // Handle WebSocket upgrades and SSE streams with TCP tunneling | |
| if response.StatusCode == http.StatusSwitchingProtocols { | |
| // WebSocket upgrade - write response headers and switch to TCP tunneling | |
| err = response.Write(srcWriter) | |
| if err != nil { | |
| if s.config.Debug { | |
| s.logDebug("Failed to write WebSocket upgrade response", err) | |
| } | |
| return | |
| } | |
| srcWriter.Flush() | |
| // Drain any bytes already buffered post-handshake to avoid loss when switching to raw TCP | |
| if n := dstReader.Buffered(); n > 0 { | |
| buf := make([]byte, n) | |
| if _, err := io.ReadFull(dstReader, buf); err == nil { | |
| if _, err := srcWriter.Write(buf); err != nil { | |
| if s.config.Debug { | |
| s.logDebug("Failed to flush buffered server bytes on WS upgrade", err) | |
| } | |
| return | |
| } | |
| srcWriter.Flush() | |
| } | |
| } | |
| if n := srcReader.Buffered(); n > 0 { | |
| buf := make([]byte, n) | |
| if _, err := io.ReadFull(srcReader, buf); err == nil { | |
| if _, err := dstWriter.Write(buf); err != nil { | |
| if s.config.Debug { | |
| s.logDebug("Failed to flush buffered client bytes on WS upgrade", err) | |
| } | |
| return | |
| } | |
| dstWriter.Flush() | |
| } | |
| } | |
| s.tcpTunnel(src, dst) | |
| return | |
| } |
🤖 Prompt for AI Agents
In tunnel/internal/client/ssh/ssh.go around lines 299 to 313, the WebSocket
upgrade path switches directly to tcpTunnel(src, dst) which can drop data
already buffered in the bufio readers (dstReader/srcReader) after
http.ReadResponse; drain any buffered bytes from both readers into their
counterpart connections before starting the tunnel: for each side obtain the
buffered bytes (e.g. via Reader.Buffered/Peek or by using io.MultiReader to
prepend the buffer), write those bytes to the opposite writer/connection (handle
write errors and flush the srcWriter), then only call tcpTunnel(src, dst) once
buffered data has been forwarded and flushed. Ensure you log errors when
draining fails and return appropriately.
Fixes #188
Summary by CodeRabbit
Release Notes