Skip to content

Commit f852689

Browse files
authored
proxy: add panic recovery to Process.ProxyRequest (#363)
Switching to use httputil.ReverseProxy in #342 introduced a possible panic if a client disconnects while streaming the body. Since llama-swap does not use http.Server the recover() is not automatically there. - introduce a recover() in Process.ProxyRequest to recover and log the event - add TestProcess_ReverseProxyPanicIsHandled to reproduce and test the fix fixes: #362
1 parent e250e71 commit f852689

File tree

2 files changed

+83
-0
lines changed

2 files changed

+83
-0
lines changed

proxy/process.go

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -499,6 +499,18 @@ func (p *Process) ProxyRequest(w http.ResponseWriter, r *http.Request) {
499499
startDuration = time.Since(beginStartTime)
500500
}
501501

502+
// recover from http.ErrAbortHandler panics that can occur when the client
503+
// disconnects before the response is sent
504+
defer func() {
505+
if r := recover(); r != nil {
506+
if r == http.ErrAbortHandler {
507+
p.proxyLogger.Infof("<%s> recovered from client disconnection during streaming", p.ID)
508+
} else {
509+
p.proxyLogger.Infof("<%s> recovered from panic: %v", p.ID, r)
510+
}
511+
}
512+
}()
513+
502514
if p.reverseProxy != nil {
503515
p.reverseProxy.ServeHTTP(w, r)
504516
} else {

proxy/process_test.go

Lines changed: 71 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -494,3 +494,74 @@ func TestProcess_EnvironmentSetCorrectly(t *testing.T) {
494494
assert.Equal(t, len(process1.cmd.Environ())+2, len(process2.cmd.Environ()), "process2 should have 2 more environment variables than process1")
495495

496496
}
497+
498+
// TestProcess_ReverseProxyPanicIsHandled tests that panics from
499+
// httputil.ReverseProxy in Process.ProxyRequest(w, r) do not bubble up and are
500+
// handled appropriately.
501+
//
502+
// httputil.ReverseProxy will panic with http.ErrAbortHandler when it has sent headers
503+
// can't copy the body. This can be caused by a client disconnecting before the full
504+
// response is sent from some reason.
505+
//
506+
// bug: https://github.com/mostlygeek/llama-swap/issues/362
507+
// see: https://github.com/golang/go/issues/23643 (where panic was added to httputil.ReverseProxy)
508+
func TestProcess_ReverseProxyPanicIsHandled(t *testing.T) {
509+
// Add defer/recover to catch any panics that aren't handled by ProxyRequest
510+
// If this recover() is hit, it means ProxyRequest didn't handle the panic properly
511+
defer func() {
512+
if r := recover(); r != nil {
513+
t.Fatalf("ProxyRequest should handle panics from reverseProxy.ServeHTTP, but panic was not caught: %v", r)
514+
}
515+
}()
516+
517+
expectedMessage := "panic_test"
518+
config := getTestSimpleResponderConfig(expectedMessage)
519+
520+
process := NewProcess("panic-test", 5, config, debugLogger, debugLogger)
521+
defer process.Stop()
522+
523+
// Start the process
524+
err := process.start()
525+
assert.Nil(t, err)
526+
assert.Equal(t, StateReady, process.CurrentState())
527+
528+
// Create a custom ResponseWriter that simulates a client disconnect
529+
// by panicking when Write is called after headers are sent
530+
panicWriter := &panicOnWriteResponseWriter{
531+
ResponseRecorder: httptest.NewRecorder(),
532+
shouldPanic: true,
533+
}
534+
535+
// Make a request that will trigger the panic
536+
req := httptest.NewRequest("GET", "/slow-respond?echo=test&delay=100ms", nil)
537+
538+
// This should panic inside reverseProxy.ServeHTTP when the panicWriter.Write() is called.
539+
// ProxyRequest should catch and handle this panic gracefully.
540+
process.ProxyRequest(panicWriter, req)
541+
542+
// If we get here, the panic was properly recovered in ProxyRequest
543+
// The process should still be in a ready state
544+
assert.Equal(t, StateReady, process.CurrentState())
545+
}
546+
547+
// panicOnWriteResponseWriter is a ResponseWriter that panics on Write
548+
// to simulate a client disconnect after headers are sent
549+
// used by: TestProcess_ReverseProxyPanicIsHandled
550+
type panicOnWriteResponseWriter struct {
551+
*httptest.ResponseRecorder
552+
shouldPanic bool
553+
headerWritten bool
554+
}
555+
556+
func (w *panicOnWriteResponseWriter) WriteHeader(statusCode int) {
557+
w.headerWritten = true
558+
w.ResponseRecorder.WriteHeader(statusCode)
559+
}
560+
561+
func (w *panicOnWriteResponseWriter) Write(b []byte) (int, error) {
562+
if w.shouldPanic && w.headerWritten {
563+
// Simulate the panic that httputil.ReverseProxy throws
564+
panic(http.ErrAbortHandler)
565+
}
566+
return w.ResponseRecorder.Write(b)
567+
}

0 commit comments

Comments
 (0)