github.com/metacubex/mihomo@v1.18.5/listener/autoredir/tcp.go (about)

     1  package autoredir
     2  
     3  import (
     4  	"net"
     5  	"net/netip"
     6  
     7  	"github.com/metacubex/mihomo/adapter/inbound"
     8  	N "github.com/metacubex/mihomo/common/net"
     9  	C "github.com/metacubex/mihomo/constant"
    10  	"github.com/metacubex/mihomo/log"
    11  	"github.com/metacubex/mihomo/transport/socks5"
    12  )
    13  
    14  type Listener struct {
    15  	listener   net.Listener
    16  	addr       string
    17  	closed     bool
    18  	additions  []inbound.Addition
    19  	lookupFunc func(netip.AddrPort) (socks5.Addr, error)
    20  }
    21  
    22  // RawAddress implements C.Listener
    23  func (l *Listener) RawAddress() string {
    24  	return l.addr
    25  }
    26  
    27  // Address implements C.Listener
    28  func (l *Listener) Address() string {
    29  	return l.listener.Addr().String()
    30  }
    31  
    32  // Close implements C.Listener
    33  func (l *Listener) Close() error {
    34  	l.closed = true
    35  	return l.listener.Close()
    36  }
    37  
    38  func (l *Listener) TCPAddr() netip.AddrPort {
    39  	return l.listener.Addr().(*net.TCPAddr).AddrPort()
    40  }
    41  
    42  func (l *Listener) SetLookupFunc(lookupFunc func(netip.AddrPort) (socks5.Addr, error)) {
    43  	l.lookupFunc = lookupFunc
    44  }
    45  
    46  func (l *Listener) handleRedir(conn net.Conn, tunnel C.Tunnel) {
    47  	if l.lookupFunc == nil {
    48  		log.Errorln("[Auto Redirect] lookup function is nil")
    49  		return
    50  	}
    51  
    52  	target, err := l.lookupFunc(conn.RemoteAddr().(*net.TCPAddr).AddrPort())
    53  	if err != nil {
    54  		log.Warnln("[Auto Redirect] %v", err)
    55  		_ = conn.Close()
    56  		return
    57  	}
    58  
    59  	N.TCPKeepAlive(conn)
    60  
    61  	tunnel.HandleTCPConn(inbound.NewSocket(target, conn, C.REDIR, l.additions...))
    62  }
    63  
    64  func New(addr string, tunnel C.Tunnel, additions ...inbound.Addition) (*Listener, error) {
    65  	if len(additions) == 0 {
    66  		additions = []inbound.Addition{
    67  			inbound.WithInName("DEFAULT-REDIR"),
    68  			inbound.WithSpecialRules(""),
    69  		}
    70  	}
    71  	l, err := net.Listen("tcp", addr)
    72  	if err != nil {
    73  		return nil, err
    74  	}
    75  	rl := &Listener{
    76  		listener:  l,
    77  		addr:      addr,
    78  		additions: additions,
    79  	}
    80  
    81  	go func() {
    82  		for {
    83  			c, err := l.Accept()
    84  			if err != nil {
    85  				if rl.closed {
    86  					break
    87  				}
    88  				continue
    89  			}
    90  			go rl.handleRedir(c, tunnel)
    91  		}
    92  	}()
    93  
    94  	return rl, nil
    95  }