github.com/metacubex/mihomo@v1.18.5/transport/snell/pool.go (about)

     1  package snell
     2  
     3  import (
     4  	"context"
     5  	"net"
     6  	"time"
     7  
     8  	"github.com/metacubex/mihomo/component/pool"
     9  	"github.com/metacubex/mihomo/transport/shadowsocks/shadowaead"
    10  )
    11  
    12  type Pool struct {
    13  	pool *pool.Pool[*Snell]
    14  }
    15  
    16  func (p *Pool) Get() (net.Conn, error) {
    17  	return p.GetContext(context.Background())
    18  }
    19  
    20  func (p *Pool) GetContext(ctx context.Context) (net.Conn, error) {
    21  	elm, err := p.pool.GetContext(ctx)
    22  	if err != nil {
    23  		return nil, err
    24  	}
    25  
    26  	return &PoolConn{elm, p}, nil
    27  }
    28  
    29  func (p *Pool) Put(conn *Snell) {
    30  	if err := HalfClose(conn); err != nil {
    31  		_ = conn.Close()
    32  		return
    33  	}
    34  
    35  	p.pool.Put(conn)
    36  }
    37  
    38  type PoolConn struct {
    39  	*Snell
    40  	pool *Pool
    41  }
    42  
    43  func (pc *PoolConn) Read(b []byte) (int, error) {
    44  	// save old status of reply (it mutable by Read)
    45  	reply := pc.Snell.reply
    46  
    47  	n, err := pc.Snell.Read(b)
    48  	if err == shadowaead.ErrZeroChunk {
    49  		// if reply is false, it should be client halfclose.
    50  		// ignore error and read data again.
    51  		if !reply {
    52  			pc.Snell.reply = false
    53  			return pc.Snell.Read(b)
    54  		}
    55  	}
    56  	return n, err
    57  }
    58  
    59  func (pc *PoolConn) Write(b []byte) (int, error) {
    60  	return pc.Snell.Write(b)
    61  }
    62  
    63  func (pc *PoolConn) Close() error {
    64  	// mihomo use SetReadDeadline to break bidirectional copy between client and server.
    65  	// reset it before reuse connection to avoid io timeout error.
    66  	_ = pc.Snell.Conn.SetReadDeadline(time.Time{})
    67  	pc.pool.Put(pc.Snell)
    68  	return nil
    69  }
    70  
    71  func NewPool(factory func(context.Context) (*Snell, error)) *Pool {
    72  	p := pool.New[*Snell](
    73  		func(ctx context.Context) (*Snell, error) {
    74  			return factory(ctx)
    75  		},
    76  		pool.WithAge[*Snell](15000),
    77  		pool.WithSize[*Snell](10),
    78  		pool.WithEvict[*Snell](func(item *Snell) {
    79  			_ = item.Close()
    80  		}),
    81  	)
    82  
    83  	return &Pool{pool: p}
    84  }