diff --git a/handler/forward/local/handler.go b/handler/forward/local/handler.go index b27e352..eedbeb7 100644 --- a/handler/forward/local/handler.go +++ b/handler/forward/local/handler.go @@ -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) } }() diff --git a/handler/forward/remote/handler.go b/handler/forward/remote/handler.go index 23ec8e9..0231df0 100644 --- a/handler/forward/remote/handler.go +++ b/handler/forward/remote/handler.go @@ -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) } }() diff --git a/handler/http/handler.go b/handler/http/handler.go index 0829274..acc6c85 100644 --- a/handler/http/handler.go +++ b/handler/http/handler.go @@ -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 } diff --git a/handler/socks/v5/handler.go b/handler/socks/v5/handler.go index ece0b83..f24358c 100644 --- a/handler/socks/v5/handler.go +++ b/handler/socks/v5/handler.go @@ -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 diff --git a/handler/tunnel/entrypoint.go b/handler/tunnel/entrypoint.go index b582431..75cd9ed 100644 --- a/handler/tunnel/entrypoint.go +++ b/handler/tunnel/entrypoint.go @@ -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) } }() diff --git a/handler/tunnel/handler.go b/handler/tunnel/handler.go index 90d4b37..b177520 100644 --- a/handler/tunnel/handler.go +++ b/handler/tunnel/handler.go @@ -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 } diff --git a/handler/tunnel/tunnel.go b/handler/tunnel/tunnel.go index 7d3d432..ebc9041 100644 --- a/handler/tunnel/tunnel.go +++ b/handler/tunnel/tunnel.go @@ -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() diff --git a/ingress/ingress.go b/ingress/ingress.go index 96027dc..e6487af 100644 --- a/ingress/ingress.go +++ b/ingress/ingress.go @@ -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 diff --git a/internal/io/io.go b/internal/io/io.go index d27bc0f..87c72f8 100644 --- a/internal/io/io.go +++ b/internal/io/io.go @@ -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, + } +}