github.com/chwjbn/xclash@v0.2.0/transport/snell/pool.go (about)

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