github.com/Asutorufa/yuhaiin@v0.3.6-0.20240502055049-7984da7023a0/pkg/net/proxy/drop/drop.go (about)

     1  package drop
     2  
     3  import (
     4  	"context"
     5  	"io"
     6  	"net"
     7  	"time"
     8  
     9  	"github.com/Asutorufa/yuhaiin/pkg/net/netapi"
    10  	"github.com/Asutorufa/yuhaiin/pkg/protos/node/point"
    11  	"github.com/Asutorufa/yuhaiin/pkg/protos/node/protocol"
    12  	"github.com/Asutorufa/yuhaiin/pkg/utils/lru"
    13  	"github.com/Asutorufa/yuhaiin/pkg/utils/singleflight"
    14  )
    15  
    16  func init() {
    17  	point.RegisterProtocol(func(*protocol.Protocol_Drop) point.WrapProxy {
    18  		return func(netapi.Proxy) (netapi.Proxy, error) { return Drop, nil }
    19  	})
    20  }
    21  
    22  var Drop = &drop{
    23  	lru: lru.New(
    24  		lru.WithCapacity[string, time.Duration](512),
    25  		lru.WithExpireTimeout[string, time.Duration](time.Second*5),
    26  	),
    27  	sf: &singleflight.Group[string, time.Duration]{},
    28  }
    29  
    30  type drop struct {
    31  	netapi.EmptyDispatch
    32  	lru *lru.LRU[string, time.Duration]
    33  	sf  *singleflight.Group[string, time.Duration]
    34  }
    35  
    36  func (d *drop) Conn(ctx context.Context, addr netapi.Address) (net.Conn, error) {
    37  	time := d.waitTime(addr)
    38  	return NewDrop(ctx, time), nil
    39  }
    40  
    41  func (d *drop) waitTime(addr netapi.Address) time.Duration {
    42  	time, _, _ := d.sf.Do(addr.String(), func() (time.Duration, error) {
    43  		en, ok := d.lru.Load(addr.String())
    44  		if ok {
    45  			if en == 0 {
    46  				en = time.Second
    47  			} else if en < time.Second*30 {
    48  				en *= 2
    49  			}
    50  		}
    51  
    52  		d.lru.Add(addr.String(), en)
    53  
    54  		return en, nil
    55  	})
    56  
    57  	return time
    58  }
    59  
    60  func (d *drop) PacketConn(ctx context.Context, addr netapi.Address) (net.PacketConn, error) {
    61  	time := d.waitTime(addr)
    62  	return NewDrop(ctx, time), nil
    63  }
    64  
    65  var _ net.Conn = (*DropConn)(nil)
    66  
    67  type DropConn struct {
    68  	ctx    context.Context
    69  	cancel context.CancelFunc
    70  }
    71  
    72  func NewDrop(ctx context.Context, timeout time.Duration) *DropConn {
    73  	if timeout == 0 {
    74  		return &DropConn{}
    75  	}
    76  
    77  	ctx, cancel := context.WithTimeout(ctx, timeout)
    78  	return &DropConn{ctx, cancel}
    79  }
    80  
    81  func (d *DropConn) ReadFrom(b []byte) (n int, addr net.Addr, err error) {
    82  	if d.ctx != nil {
    83  		<-d.ctx.Done()
    84  	}
    85  	return 0, nil, io.EOF
    86  }
    87  
    88  func (d *DropConn) WriteTo(b []byte, addr net.Addr) (n int, err error) {
    89  	return len(b), nil
    90  }
    91  
    92  func (d *DropConn) Read(b []byte) (n int, err error) {
    93  	if d.ctx != nil {
    94  		<-d.ctx.Done()
    95  	}
    96  
    97  	return 0, io.EOF
    98  }
    99  
   100  func (d *DropConn) Write(b []byte) (n int, err error) {
   101  	return len(b), nil
   102  }
   103  
   104  func (d *DropConn) Close() error {
   105  	if d.cancel != nil {
   106  		d.cancel()
   107  	}
   108  	return nil
   109  }
   110  
   111  func (d *DropConn) LocalAddr() net.Addr {
   112  	return &net.TCPAddr{IP: net.IP{0, 0, 0, 0}}
   113  }
   114  
   115  func (d *DropConn) RemoteAddr() net.Addr {
   116  	return &net.TCPAddr{IP: net.IP{0, 0, 0, 0}}
   117  }
   118  
   119  func (d *DropConn) SetDeadline(t time.Time) error { return nil }
   120  
   121  func (d *DropConn) SetReadDeadline(t time.Time) error { return nil }
   122  
   123  func (d *DropConn) SetWriteDeadline(t time.Time) error { return nil }