fix connection state in tunnel entrypoint

This commit is contained in:
ginuerzh 2024-08-01 20:52:08 +08:00
parent 12ef82e41f
commit 1a776dc759
9 changed files with 121 additions and 16 deletions

View File

@ -106,7 +106,7 @@ func (h *forwardHandler) Handle(ctx context.Context, conn net.Conn, opts ...hand
}
if protocol == forward.ProtoHTTP {
h.handleHTTP(ctx, rw, conn.RemoteAddr(), log)
h.handleHTTP(ctx, xio.NewReadWriteCloser(rw, rw, conn), conn.RemoteAddr(), log)
return nil
}
@ -177,7 +177,7 @@ func (h *forwardHandler) Handle(ctx context.Context, conn net.Conn, opts ...hand
return nil
}
func (h *forwardHandler) handleHTTP(ctx context.Context, rw io.ReadWriter, remoteAddr net.Addr, log logger.Logger) (err error) {
func (h *forwardHandler) handleHTTP(ctx context.Context, rw io.ReadWriteCloser, remoteAddr net.Addr, log logger.Logger) (err error) {
br := bufio.NewReader(rw)
for {
@ -334,12 +334,18 @@ func (h *forwardHandler) handleHTTP(ctx context.Context, rw io.ReadWriter, remot
log.Trace(string(dump))
}
if res.Close {
defer rw.Close()
}
if err := h.rewriteBody(res, bodyRewrites...); err != nil {
rw.Close()
log.Errorf("rewrite body: %v", err)
return
}
if err = res.Write(rw); err != nil {
rw.Close()
log.Errorf("write response from node %s(%s): %v", target.Name, target.Addr, err)
}
}()

View File

@ -108,7 +108,7 @@ func (h *forwardHandler) Handle(ctx context.Context, conn net.Conn, opts ...hand
}
}
if protocol == forward.ProtoHTTP {
h.handleHTTP(ctx, rw, conn.RemoteAddr(), localAddr, log)
h.handleHTTP(ctx, xio.NewReadWriteCloser(rw, rw, conn), conn.RemoteAddr(), localAddr, log)
return nil
}
@ -178,7 +178,7 @@ func (h *forwardHandler) Handle(ctx context.Context, conn net.Conn, opts ...hand
return nil
}
func (h *forwardHandler) handleHTTP(ctx context.Context, rw io.ReadWriter, remoteAddr net.Addr, localAddr net.Addr, log logger.Logger) (err error) {
func (h *forwardHandler) handleHTTP(ctx context.Context, rw io.ReadWriteCloser, remoteAddr net.Addr, localAddr net.Addr, log logger.Logger) (err error) {
br := bufio.NewReader(rw)
for {
@ -335,12 +335,18 @@ func (h *forwardHandler) handleHTTP(ctx context.Context, rw io.ReadWriter, remot
log.Trace(string(dump))
}
if res.Close {
defer rw.Close()
}
if err := h.rewriteBody(res, bodyRewrites...); err != nil {
rw.Close()
log.Errorf("rewrite body: %v", err)
return
}
if err = res.Write(rw); err != nil {
rw.Close()
log.Errorf("write response from node %s(%s): %v", target.Name, target.Addr, err)
}
}()

View File

@ -279,6 +279,9 @@ func (h *httpHandler) handleProxy(rw, cc io.ReadWriter, req *http.Request, log l
err := func() error {
req, err := http.ReadRequest(bufio.NewReader(rw))
if err != nil {
if err == io.EOF {
return nil
}
return err
}

View File

@ -101,6 +101,7 @@ func (h *socks5Handler) Handle(ctx context.Context, conn net.Conn, opts ...handl
if clientID := sc.ID(); clientID != "" {
ctx = ctxvalue.ContextWithClientID(ctx, ctxvalue.ClientID(clientID))
log = log.WithFields(map[string]any{"user": clientID})
}
conn = sc

View File

@ -123,13 +123,15 @@ func (ep *entrypoint) handle(ctx context.Context, conn net.Conn) error {
timeout: 15 * time.Second,
log: log,
}
cc, node, cid, err := d.Dial(ctx, "tcp", tunnelID.String())
c, node, cid, err := d.Dial(ctx, "tcp", tunnelID.String())
if err != nil {
log.Error(err)
return resp.Write(conn)
}
log.Debugf("new connection to tunnel: %s, connector: %s", tunnelID, cid)
cc = c
host := req.Host
if h, _, _ := net.SplitHostPort(host); h == "" {
host = net.JoinHostPort(strings.Trim(host, "[]"), "80")
@ -149,17 +151,26 @@ func (ep *entrypoint) handle(ctx context.Context, conn net.Conn) error {
Version: relay.Version1,
Status: relay.StatusOK,
Features: features,
}).WriteTo(cc)
}).WriteTo(c)
}
if err := req.Write(cc); err != nil {
cc.Close()
// HTTP/1.0
if req.ProtoMajor == 1 && req.ProtoMinor == 0 {
if strings.ToLower(req.Header.Get("Connection")) == "keep-alive" {
req.Header.Del("Connection")
} else {
req.Header.Set("Connection", "close")
}
}
if err := req.Write(c); err != nil {
c.Close()
log.Errorf("send request: %v", err)
return resp.Write(conn)
}
if req.Header.Get("Upgrade") == "websocket" {
err := xnet.Transport(cc, xio.NewReadWriter(br, conn))
err := xnet.Transport(c, xio.NewReadWriter(br, conn))
if err == nil {
err = io.EOF
}
@ -167,7 +178,7 @@ func (ep *entrypoint) handle(ctx context.Context, conn net.Conn) error {
}
go func() {
defer cc.Close()
defer c.Close()
t := time.Now()
log.Debugf("%s <-> %s", remoteAddr, host)
@ -178,7 +189,7 @@ func (ep *entrypoint) handle(ctx context.Context, conn net.Conn) error {
}).Debugf("%s >-< %s", remoteAddr, host)
}()
res, err := http.ReadResponse(bufio.NewReader(cc), req)
res, err := http.ReadResponse(bufio.NewReader(c), req)
if err != nil {
log.Errorf("read response: %v", err)
resp.Write(conn)
@ -190,7 +201,21 @@ func (ep *entrypoint) handle(ctx context.Context, conn net.Conn) error {
log.Trace(string(dump))
}
if res.Close {
defer conn.Close()
}
// HTTP/1.0
if req.ProtoMajor == 1 && req.ProtoMinor == 0 {
if !res.Close {
res.Header.Set("Connection", "keep-alive")
}
res.ProtoMajor = req.ProtoMajor
res.ProtoMinor = req.ProtoMinor
}
if err = res.Write(conn); err != nil {
conn.Close()
log.Errorf("write response: %v", err)
}
}()

View File

@ -264,6 +264,10 @@ func (h *tunnelHandler) Handle(ctx context.Context, conn net.Conn, opts ...handl
// Close implements io.Closer interface.
func (h *tunnelHandler) Close() error {
if h.epSvc != nil {
h.epSvc.Close()
}
h.pool.Close()
return nil
}

View File

@ -66,6 +66,14 @@ func (c *Connector) Session() *mux.Session {
return c.s
}
func (c *Connector) Close() error {
if c == nil || c.s == nil {
return nil
}
return c.s.Close()
}
type Tunnel struct {
node string
id relay.TunnelID
@ -75,7 +83,7 @@ type Tunnel struct {
mu sync.RWMutex
sd sd.SD
ttl time.Duration
rw *selector.RandomWeighted[*Connector]
// rw *selector.RandomWeighted[*Connector]
}
func NewTunnel(node string, tid relay.TunnelID, ttl time.Duration) *Tunnel {
@ -85,7 +93,7 @@ func NewTunnel(node string, tid relay.TunnelID, ttl time.Duration) *Tunnel {
t: time.Now(),
close: make(chan struct{}),
ttl: ttl,
rw: selector.NewRandomWeighted[*Connector](),
// rw: selector.NewRandomWeighted[*Connector](),
}
if t.ttl <= 0 {
t.ttl = defaultTTL
@ -117,8 +125,14 @@ func (t *Tunnel) GetConnector(network string) *Connector {
t.mu.RLock()
defer t.mu.RUnlock()
rw := t.rw
rw.Reset()
// rw := t.rw
// rw.Reset()
if len(t.connectors) == 1 {
return t.connectors[0]
}
rw := selector.NewRandomWeighted[*Connector]()
found := false
for _, c := range t.connectors {
@ -147,6 +161,22 @@ func (t *Tunnel) GetConnector(network string) *Connector {
return rw.Next()
}
func (t *Tunnel) Close() error {
t.mu.Lock()
defer t.mu.Unlock()
select {
case <-t.close:
default:
for _, c := range t.connectors {
c.Close()
}
close(t.close)
}
return nil
}
func (t *Tunnel) CloseOnIdle() bool {
t.mu.RLock()
defer t.mu.RUnlock()
@ -256,6 +286,22 @@ func (p *ConnectorPool) Get(network string, tid string) *Connector {
return t.GetConnector(network)
}
func (p *ConnectorPool) Close() error {
if p == nil {
return nil
}
p.mu.Lock()
defer p.mu.Unlock()
for k, v := range p.tunnels {
v.Close()
delete(p.tunnels, k)
}
return nil
}
func (p *ConnectorPool) closeIdles() {
ticker := time.NewTicker(1 * time.Hour)
defer ticker.Stop()

View File

@ -245,7 +245,7 @@ func (ing *localIngress) GetRule(ctx context.Context, host string, opts ...ingre
}
if ep != nil {
ing.options.logger.Debugf("ingress: %s -> %s", host, ep)
ing.options.logger.Debugf("ingress: %s -> %s:%s", host, ep.Hostname, ep.Endpoint)
}
return ep

View File

@ -13,3 +13,17 @@ func NewReadWriter(r io.Reader, w io.Writer) io.ReadWriter {
Writer: w,
}
}
type readWriteCloser struct {
io.Reader
io.Writer
io.Closer
}
func NewReadWriteCloser(r io.Reader, w io.Writer, c io.Closer) io.ReadWriteCloser {
return &readWriteCloser{
Reader: r,
Writer: w,
Closer: c,
}
}