Skip to content

Commit 6599555

Browse files
committed
proxy: initial connection resilience
1 parent 063b60a commit 6599555

3 files changed

Lines changed: 197 additions & 140 deletions

File tree

http.go

Lines changed: 84 additions & 85 deletions
Original file line numberDiff line numberDiff line change
@@ -2,49 +2,49 @@ package proxy
22

33
import (
44
"context"
5-
"net"
6-
"net/http"
75
"crypto/tls"
86
"io"
9-
"time"
7+
"net"
8+
"net/http"
109
"strings"
10+
"time"
1111
// "errors"
1212
"fmt"
1313
// "sync"
14+
"bytes"
1415

1516
"github.com/urnetwork/connect"
1617
)
1718

18-
1919
// a simple http/https proxy focused on proper file descriptor management
2020

21-
2221
type HttpProxy struct {
23-
ProxyReadTimeout time.Duration
22+
ProxyReadTimeout time.Duration
2423
ProxyWriteTimeout time.Duration
25-
ProxyIdleTimeout time.Duration
24+
ProxyIdleTimeout time.Duration
2625
// only used for http proxy
2726
ProxyTlsHandshakeTimeout time.Duration
28-
ConnectDialWithRequest func(r *http.Request, network string, addr string) (net.Conn, error)
29-
GetTlsConfigForClient func(*tls.ClientHelloInfo) (*tls.Config, error)
27+
MaxHttpBodyBytes int64
28+
ConnectDialWithRequest func(r *http.Request, network string, addr string) (net.Conn, error)
29+
GetTlsConfigForClient func(*tls.ClientHelloInfo) (*tls.Config, error)
3030
}
3131

3232
func NewHttpProxy() *HttpProxy {
33-
return &HttpProxy{}
33+
return &HttpProxy{
34+
MaxHttpBodyBytes: 2 * 1024 * 1024,
35+
}
3436
}
3537

36-
3738
func (self *HttpProxy) ListenAndServe(ctx context.Context, network string, addr string) error {
3839

3940
httpServer := &http.Server{
40-
Addr: addr,
41-
Handler: self,
42-
ReadTimeout: self.ProxyReadTimeout,
41+
Addr: addr,
42+
Handler: self,
43+
ReadTimeout: self.ProxyReadTimeout,
4344
WriteTimeout: self.ProxyWriteTimeout,
44-
IdleTimeout: self.ProxyIdleTimeout,
45+
IdleTimeout: self.ProxyIdleTimeout,
4546
}
4647

47-
4848
listenConfig := net.ListenConfig{}
4949

5050
l, err := listenConfig.Listen(
@@ -60,23 +60,21 @@ func (self *HttpProxy) ListenAndServe(ctx context.Context, network string, addr
6060
return httpServer.Serve(l)
6161
}
6262

63-
6463
func (self *HttpProxy) ListenAndServeTls(ctx context.Context, network string, addr string) error {
6564

6665
tlsConfig := &tls.Config{
6766
GetConfigForClient: self.GetTlsConfigForClient,
6867
}
6968

7069
httpServer := &http.Server{
71-
Addr: addr,
72-
Handler: self,
73-
TLSConfig: tlsConfig,
74-
ReadTimeout: self.ProxyReadTimeout,
70+
Addr: addr,
71+
Handler: self,
72+
TLSConfig: tlsConfig,
73+
ReadTimeout: self.ProxyReadTimeout,
7574
WriteTimeout: self.ProxyWriteTimeout,
76-
IdleTimeout: self.ProxyIdleTimeout,
75+
IdleTimeout: self.ProxyIdleTimeout,
7776
}
7877

79-
8078
listenConfig := net.ListenConfig{}
8179

8280
l, err := listenConfig.Listen(
@@ -92,7 +90,6 @@ func (self *HttpProxy) ListenAndServeTls(ctx context.Context, network string, ad
9290
return httpServer.ServeTLS(l, "", "")
9391
}
9492

95-
9693
func (self *HttpProxy) ServeHTTP(w http.ResponseWriter, r *http.Request) {
9794
connect.HandleError(func() {
9895
if r.Method == http.MethodConnect {
@@ -103,9 +100,6 @@ func (self *HttpProxy) ServeHTTP(w http.ResponseWriter, r *http.Request) {
103100
})
104101
}
105102

106-
107-
108-
109103
func (self *HttpProxy) handleHttps(w http.ResponseWriter, r *http.Request) {
110104
hij := w.(http.Hijacker)
111105

@@ -114,75 +108,94 @@ func (self *HttpProxy) handleHttps(w http.ResponseWriter, r *http.Request) {
114108
http.Error(w, http.StatusText(http.StatusInternalServerError), http.StatusInternalServerError)
115109
return
116110
}
117-
defer conn.Close()
118-
119-
// r.URL.Host contains both the host and port (if specified)
120-
proxyConn, err := self.ConnectDialWithRequest(r, "tcp", r.URL.Host)
121-
if err != nil {
122-
httpError(conn, http.StatusBadGateway, err)
123-
return
124-
}
125-
defer proxyConn.Close()
126-
127-
conn.Write([]byte("HTTP/1.0 200 Connection established\r\n\r\n"))
128-
129111
handleCtx, handleCancel := context.WithCancel(r.Context())
130112
defer handleCancel()
131-
132113
go connect.HandleError(func() {
133-
defer handleCancel()
134-
copyBufferWithTimeout(proxyConn, conn, nil, self.ProxyReadTimeout, self.ProxyWriteTimeout)
114+
defer conn.Close()
115+
select {
116+
case <-handleCtx.Done():
117+
}
135118
})
136119

137-
go connect.HandleError(func() {
138-
defer handleCancel()
139-
copyBufferWithTimeout(conn, proxyConn, nil, self.ProxyReadTimeout, self.ProxyWriteTimeout)
140-
})
120+
// r.URL.Host contains both the host and port (if specified)
121+
var proxyConn net.Conn
122+
for {
123+
select {
124+
case <-handleCtx.Done():
125+
httpError(conn, http.StatusBadGateway, err)
126+
return
127+
default:
128+
}
129+
proxyConn, err = self.ConnectDialWithRequest(r, "tcp", r.URL.Host)
130+
if err == nil {
131+
break
132+
}
133+
}
134+
defer proxyConn.Close()
141135

142-
select {
143-
case <- handleCtx.Done():
136+
_, err = conn.Write([]byte("HTTP/1.0 200 Connection established\r\n\r\n"))
137+
if err != nil {
138+
return
144139
}
145140

146-
return
141+
copyConn(handleCtx, handleCancel, conn, proxyConn, self.ProxyReadTimeout, self.ProxyWriteTimeout)
147142
}
148143

149144
func (self *HttpProxy) handleHttp(w http.ResponseWriter, r *http.Request) {
150145
// r.Header.Del("Accept-Encoding")
151146
// r.Header.Del("Proxy-Connection")
152147
// r.Header.Del("Proxy-Authenticate")
153148
// r.Header.Del("Proxy-Authorization")
154-
155149

156-
r2, err := http.NewRequestWithContext(
157-
r.Context(),
158-
r.Method,
159-
r.URL.String(),
160-
r.Body,
161-
)
150+
b := bytes.NewBuffer(nil)
151+
_, err := copyBufferWithTimeout(b, io.LimitReader(r.Body, self.MaxHttpBodyBytes), nil, self.ProxyReadTimeout, self.ProxyWriteTimeout)
152+
r.Body.Close()
162153
if err != nil {
163154
http.Error(w, http.StatusText(http.StatusInternalServerError), http.StatusInternalServerError)
164155
return
165156
}
157+
bodyBytes := b.Bytes()
158+
159+
handleCtx, handleCancel := context.WithCancel(r.Context())
160+
defer handleCancel()
166161

167162
tr := &http.Transport{
168163
Dial: func(network string, addr string) (net.Conn, error) {
169-
return connect.HandleError2(func()(net.Conn, error) {
164+
return connect.HandleError2(func() (net.Conn, error) {
170165
return self.ConnectDialWithRequest(r, network, addr)
171-
}, func()(net.Conn, error) {
166+
}, func() (net.Conn, error) {
172167
return nil, fmt.Errorf("Unexpected error")
173168
})
174169
},
175-
DisableKeepAlives: true,
176-
TLSHandshakeTimeout: self.ProxyTlsHandshakeTimeout,
170+
DisableKeepAlives: true,
171+
TLSHandshakeTimeout: self.ProxyTlsHandshakeTimeout,
177172
ResponseHeaderTimeout: self.ProxyReadTimeout,
178173
}
179174

180-
181-
response, err := tr.RoundTrip(r2)
182-
if err != nil {
183-
http.Error(w, http.StatusText(http.StatusBadGateway), http.StatusBadGateway)
184-
return
175+
var response *http.Response
176+
for {
177+
select {
178+
case <-handleCtx.Done():
179+
http.Error(w, http.StatusText(http.StatusBadGateway), http.StatusBadGateway)
180+
return
181+
default:
182+
}
183+
r2, err := http.NewRequestWithContext(
184+
r.Context(),
185+
r.Method,
186+
r.URL.String(),
187+
bytes.NewReader(bodyBytes),
188+
)
189+
if err != nil {
190+
http.Error(w, http.StatusText(http.StatusInternalServerError), http.StatusInternalServerError)
191+
return
192+
}
193+
response, err = tr.RoundTrip(r2)
194+
if err == nil {
195+
break
196+
}
185197
}
198+
defer response.Body.Close()
186199

187200
h := w.Header()
188201
for k := range w.Header() {
@@ -201,29 +214,17 @@ func (self *HttpProxy) handleHttp(w http.ResponseWriter, r *http.Request) {
201214
http.Error(w, http.StatusText(http.StatusInternalServerError), http.StatusInternalServerError)
202215
return
203216
}
204-
defer conn.Close()
205-
206-
proxyRw := response.Body.(io.ReadWriter)
207-
defer response.Body.Close()
208-
209-
handleCtx, handleCancel := context.WithCancel(r.Context())
210-
211217
go connect.HandleError(func() {
212-
defer handleCancel()
213-
copyBufferWithTimeout(conn, proxyRw, nil, self.ProxyReadTimeout, self.ProxyWriteTimeout)
218+
defer conn.Close()
219+
select {
220+
case <-handleCtx.Done():
221+
}
214222
})
215223

216-
go connect.HandleError(func() {
217-
defer handleCancel()
218-
copyBufferWithTimeout(proxyRw, conn, nil, self.ProxyReadTimeout, self.ProxyWriteTimeout)
219-
})
224+
proxyRw := response.Body.(io.ReadWriter)
220225

221-
select {
222-
case <- handleCtx.Done():
223-
}
226+
copyConn(handleCtx, handleCancel, conn, proxyRw, self.ProxyReadTimeout, self.ProxyWriteTimeout)
224227
} else {
225-
defer response.Body.Close()
226-
227228
var flush func()
228229

229230
chunked := false
@@ -256,7 +257,6 @@ func headerContains(h http.Header, name string, value string) bool {
256257
return false
257258
}
258259

259-
260260
// for a hijacked connection
261261
func httpError(w io.Writer, statusCode int, err error) error {
262262
errorMessage := err.Error()
@@ -270,4 +270,3 @@ func httpError(w io.Writer, statusCode int, err error) error {
270270
_, writeErr := io.WriteString(w, errStr)
271271
return writeErr
272272
}
273-

socks.go

Lines changed: 13 additions & 41 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@ import (
1313
"strings"
1414
// "syscall"
1515
"time"
16+
// "sync"
1617

1718
// "github.com/elazarl/goproxy"
1819
socks5 "github.com/things-go/go-socks5"
@@ -22,15 +23,13 @@ import (
2223
"github.com/urnetwork/glog"
2324
)
2425

25-
2626
type SocksRequest = *socks5.Request
2727

28-
2928
type SocksProxy struct {
30-
ProxyReadTimeout time.Duration
31-
ProxyWriteTimeout time.Duration
29+
ProxyReadTimeout time.Duration
30+
ProxyWriteTimeout time.Duration
3231
ConnectDialWithRequest func(ctx context.Context, r SocksRequest, network string, addr string) (net.Conn, error)
33-
ValidUser func(user string, password string, userAddr string) bool
32+
ValidUser func(user string, password string, userAddr string) bool
3433
}
3534

3635
func NewSocksProxy() *SocksProxy {
@@ -51,16 +50,15 @@ func (self *SocksProxy) ListenAndServe(ctx context.Context, network string, addr
5150
return nil, fmt.Errorf("Unexpected error")
5251
})
5352
}),
54-
socks5.WithConnectHandle(func(ctx context.Context, writer io.Writer, r SocksRequest)(error) {
55-
return connect.HandleError1(func()(error) {
53+
socks5.WithConnectHandle(func(ctx context.Context, writer io.Writer, r SocksRequest) error {
54+
return connect.HandleError1(func() error {
5655
return self.connectHandle(ctx, writer, r)
57-
}, func()(error) {
56+
}, func() error {
5857
return fmt.Errorf("Unexpected error")
5958
})
6059
}),
6160
)
6261

63-
6462
listenConfig := net.ListenConfig{}
6563

6664
l, err := listenConfig.Listen(
@@ -97,45 +95,21 @@ func (self *SocksProxy) connectHandle(ctx context.Context, writer io.Writer, r S
9795
socks5.SendReply(writer, resp, nil)
9896
return err
9997
}
100-
defer proxyConn.Close()
101-
102-
if err := socks5.SendReply(writer, statute.RepSuccess, proxyConn.LocalAddr()); err != nil {
103-
return err
104-
}
105-
10698
handleCtx, handleCancel := context.WithCancel(ctx)
10799
defer handleCancel()
108-
109-
errs := make(chan error)
110-
111100
go connect.HandleError(func() {
112-
defer handleCancel()
113-
_, err := copyBufferWithTimeout(proxyConn, r.Reader, nil, self.ProxyReadTimeout, self.ProxyWriteTimeout)
101+
defer proxyConn.Close()
114102
select {
115-
case <- handleCtx.Done():
116-
case errs <- err:
103+
case <-handleCtx.Done():
117104
}
118105
})
119106

120-
go connect.HandleError(func() {
121-
defer handleCancel()
122-
_, err := copyBufferWithTimeout(writer, proxyConn, nil, self.ProxyReadTimeout, self.ProxyWriteTimeout)
123-
select {
124-
case <- handleCtx.Done():
125-
case errs <- err:
126-
}
127-
})
128-
129-
for {
130-
select {
131-
case <- handleCtx.Done():
132-
return nil
133-
case err := <- errs:
134-
return err
135-
}
107+
if err := socks5.SendReply(writer, statute.RepSuccess, proxyConn.LocalAddr()); err != nil {
108+
return err
136109
}
137-
}
138110

111+
return copyRw(handleCtx, handleCancel, r.Reader, writer, proxyConn, proxyConn, self.ProxyReadTimeout, self.ProxyWriteTimeout)
112+
}
139113

140114
// socks.Logger
141115
func (self *SocksProxy) Errorf(format string, args ...any) {
@@ -156,5 +130,3 @@ func (self *SocksProxy) Resolve(ctx context.Context, name string) (context.Conte
156130
// names are not resolved locally
157131
return ctx, net.ParseIP("0.0.0.0").To4(), nil
158132
}
159-
160-

0 commit comments

Comments
 (0)