120 lines
2.2 KiB
Go
120 lines
2.2 KiB
Go
package proxy
|
|
|
|
import (
|
|
"io"
|
|
"net"
|
|
"net/http"
|
|
"sync"
|
|
"sync/atomic"
|
|
"time"
|
|
)
|
|
|
|
type writeTimeoutConn struct {
|
|
net.Conn
|
|
timeout time.Duration
|
|
}
|
|
|
|
func (c *writeTimeoutConn) Write(p []byte) (int, error) {
|
|
if err := c.Conn.SetWriteDeadline(time.Now().Add(c.timeout)); err != nil {
|
|
return 0, err
|
|
}
|
|
n, err := c.Conn.Write(p)
|
|
clearErr := c.Conn.SetWriteDeadline(time.Time{})
|
|
if err == nil {
|
|
err = clearErr
|
|
}
|
|
return n, err
|
|
}
|
|
|
|
type trackingBody struct {
|
|
io.ReadCloser
|
|
failed *atomic.Bool
|
|
}
|
|
|
|
func (b *trackingBody) Read(p []byte) (int, error) {
|
|
n, err := b.ReadCloser.Read(p)
|
|
if err != nil && err != io.EOF {
|
|
b.failed.Store(true)
|
|
}
|
|
return n, err
|
|
}
|
|
|
|
func responseBodyTimeout(resp *http.Response, opts Options) time.Duration {
|
|
if resp.StatusCode == http.StatusSwitchingProtocols {
|
|
return 0
|
|
}
|
|
if isStreamResponse(resp.Header) {
|
|
return opts.SSEIdleTimeout
|
|
}
|
|
return opts.ResponseTimeout
|
|
}
|
|
|
|
type timeoutBody struct {
|
|
body io.ReadCloser
|
|
idle bool
|
|
timeout time.Duration
|
|
timer *time.Timer
|
|
mu sync.Mutex
|
|
done bool
|
|
sequence uint64
|
|
closeOnce sync.Once
|
|
closeErr error
|
|
}
|
|
|
|
func newTimeoutBody(body io.ReadCloser, timeout time.Duration, idle bool) *timeoutBody {
|
|
t := &timeoutBody{body: body, idle: idle, timeout: timeout}
|
|
t.resetLocked()
|
|
return t
|
|
}
|
|
|
|
func (b *timeoutBody) Read(p []byte) (int, error) {
|
|
n, err := b.body.Read(p)
|
|
b.mu.Lock()
|
|
if !b.done {
|
|
if err != nil {
|
|
b.done = true
|
|
b.timer.Stop()
|
|
} else if n > 0 && b.idle {
|
|
b.resetLocked()
|
|
}
|
|
}
|
|
b.mu.Unlock()
|
|
return n, err
|
|
}
|
|
|
|
func (b *timeoutBody) Close() error {
|
|
b.mu.Lock()
|
|
if !b.done {
|
|
b.done = true
|
|
b.sequence++
|
|
b.timer.Stop()
|
|
}
|
|
b.mu.Unlock()
|
|
return b.closeUnderlying()
|
|
}
|
|
|
|
func (b *timeoutBody) resetLocked() {
|
|
if b.timer != nil {
|
|
b.timer.Stop()
|
|
}
|
|
b.sequence++
|
|
sequence := b.sequence
|
|
b.timer = time.AfterFunc(b.timeout, func() { b.expire(sequence) })
|
|
}
|
|
|
|
func (b *timeoutBody) expire(sequence uint64) {
|
|
b.mu.Lock()
|
|
if b.done || sequence != b.sequence {
|
|
b.mu.Unlock()
|
|
return
|
|
}
|
|
b.done = true
|
|
b.mu.Unlock()
|
|
_ = b.closeUnderlying()
|
|
}
|
|
|
|
func (b *timeoutBody) closeUnderlying() error {
|
|
b.closeOnce.Do(func() { b.closeErr = b.body.Close() })
|
|
return b.closeErr
|
|
}
|