diff --git a/proxy.go b/proxy.go
index 0058333..7221334 100644
--- a/proxy.go
+++ b/proxy.go
@@ -512,7 +512,7 @@ func withCloseNotifier(handler http.HandlerFunc) http.HandlerFunc {
 			cancel()
 		}()
 		select {
-		case <-w.(http.CloseNotifier).CloseNotify():
+		case <-r.Context().Done():
 			cancel()
 		case <-ctx.Done():
 		}