pool.go 1.7 KB

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