github.com/laof/lite-speed-test@v0.0.0-20230930011949-1f39b7037845/proxy/proxy.go (about)

     1  package proxy
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  	"io"
     7  	"net"
     8  	"strconv"
     9  	"time"
    10  
    11  	"github.com/laof/lite-speed-test/common"
    12  	N "github.com/laof/lite-speed-test/common/net"
    13  	"github.com/laof/lite-speed-test/common/pool"
    14  	"github.com/laof/lite-speed-test/log"
    15  	"github.com/laof/lite-speed-test/tunnel"
    16  	"github.com/laof/lite-speed-test/utils"
    17  )
    18  
    19  // proxy http/scocks to vmess
    20  const Name = "PROXY"
    21  
    22  type Proxy struct {
    23  	sources []tunnel.Server
    24  	sink    tunnel.Client
    25  	ctx     context.Context
    26  	cancel  context.CancelFunc
    27  	pool    *utils.WorkerPool
    28  }
    29  
    30  func (p *Proxy) Run() error {
    31  	p.relayConnLoop()
    32  	// p.relayPacketLoop()
    33  	<-p.ctx.Done()
    34  	return nil
    35  }
    36  
    37  func (p *Proxy) Close() error {
    38  	p.cancel()
    39  	p.sink.Close()
    40  	for _, source := range p.sources {
    41  		source.Close()
    42  	}
    43  	if p.pool != nil {
    44  		p.pool.Stop()
    45  	}
    46  	return nil
    47  }
    48  
    49  // forward from socks/http connection to vmess/trojan
    50  // TODO: bypass cn
    51  func (p *Proxy) relayConnLoop() {
    52  	pool := utils.WorkerPool{
    53  		WorkerFunc: func(inbound tunnel.Conn) error {
    54  			defer inbound.Close()
    55  			addr := inbound.Metadata().Address
    56  			var outbound net.Conn
    57  			var err error
    58  			start := time.Now()
    59  			if addr.IP != nil && N.IsPrivateAddress(addr.IP) {
    60  				networkType := addr.NetworkType
    61  				if networkType == "" {
    62  					networkType = "tcp"
    63  				}
    64  				add := net.JoinHostPort(addr.IP.String(), strconv.Itoa(addr.Port))
    65  				outbound, err = net.Dial(networkType, add)
    66  			} else {
    67  				outbound, err = p.sink.DialConn(addr, nil)
    68  			}
    69  			if err != nil {
    70  				log.Error(common.NewError("proxy failed to dial connection").Base(err))
    71  				return err
    72  			}
    73  			elapsed := fmt.Sprintf("%dms", time.Since(start).Milliseconds())
    74  			log.D("connect to:", addr, elapsed)
    75  			defer outbound.Close()
    76  			// relay
    77  			return relay(outbound, inbound)
    78  		},
    79  		MaxWorkersCount:       2000,
    80  		LogAllErrors:          false,
    81  		MaxIdleWorkerDuration: 2 * time.Minute,
    82  	}
    83  	p.pool = &pool
    84  	pool.Start()
    85  	for _, source := range p.sources {
    86  		go func(source tunnel.Server) {
    87  			for {
    88  				inbound, err := source.AcceptConn(nil)
    89  				if err != nil {
    90  					select {
    91  					case <-p.ctx.Done():
    92  						log.D("exiting")
    93  						return
    94  					default:
    95  					}
    96  					log.Error(common.NewError("failed to accept connection").Base(err))
    97  					continue
    98  				}
    99  				pool.Serve(inbound)
   100  			}
   101  		}(source)
   102  	}
   103  }
   104  
   105  func relay(leftConn, rightConn net.Conn) error {
   106  	ch := make(chan error)
   107  
   108  	go func() {
   109  		buf := pool.Get(pool.RelayBufferSize)
   110  		// Wrapping to avoid using *net.TCPConn.(ReadFrom)
   111  		// See also https://github.com/Dreamacro/clash/pull/1209
   112  		_, err := io.CopyBuffer(N.WriteOnlyWriter{Writer: leftConn}, N.ReadOnlyReader{Reader: rightConn}, buf)
   113  		if err != nil {
   114  			log.E(err.Error())
   115  		}
   116  		pool.Put(buf)
   117  		leftConn.SetReadDeadline(time.Now())
   118  		ch <- err
   119  	}()
   120  
   121  	buf := pool.Get(pool.RelayBufferSize)
   122  	_, err := io.CopyBuffer(N.WriteOnlyWriter{Writer: rightConn}, N.ReadOnlyReader{Reader: leftConn}, buf)
   123  	if err != nil {
   124  		log.E(err.Error())
   125  	}
   126  	pool.Put(buf)
   127  	rightConn.SetReadDeadline(time.Now())
   128  	err = <-ch
   129  	return err
   130  }
   131  
   132  func NewProxy(ctx context.Context, cancel context.CancelFunc, sources []tunnel.Server, sink tunnel.Client) *Proxy {
   133  	return &Proxy{
   134  		sources: sources,
   135  		sink:    sink,
   136  		ctx:     ctx,
   137  		cancel:  cancel,
   138  	}
   139  }
   140  
   141  // A Dialer is a means to establish a connection.
   142  // Custom dialers should also implement ContextDialer.
   143  type Dialer interface {
   144  	// Dial connects to the given address via the proxy.
   145  	Dial(network, addr string) (c net.Conn, err error)
   146  }