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

     1  package reject
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     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/protos/statistic"
    13  	"github.com/Asutorufa/yuhaiin/pkg/utils/lru"
    14  )
    15  
    16  var _ netapi.Proxy = (*reject)(nil)
    17  
    18  type rejectImmediately struct{ netapi.EmptyDispatch }
    19  
    20  func (rejectImmediately) Conn(_ context.Context, addr netapi.Address) (net.Conn, error) {
    21  	return nil, netapi.NewBlockError(statistic.Type_tcp, addr.Hostname())
    22  }
    23  func (rejectImmediately) PacketConn(_ context.Context, addr netapi.Address) (net.PacketConn, error) {
    24  	return nil, netapi.NewBlockError(statistic.Type_udp, addr.Hostname())
    25  }
    26  
    27  type reject struct {
    28  	cache         *lru.LRU[string, object]
    29  	max, internal int
    30  	netapi.EmptyDispatch
    31  }
    32  
    33  type object struct {
    34  	times int8
    35  	time  time.Time
    36  	delay time.Duration
    37  }
    38  
    39  func init() {
    40  	point.RegisterProtocol(func(*protocol.Protocol_Reject) point.WrapProxy {
    41  		return func(netapi.Proxy) (netapi.Proxy, error) { return Default, nil }
    42  	})
    43  }
    44  
    45  var Default = rejectImmediately{}
    46  
    47  func NewReject(maxDelay, interval int) netapi.Proxy {
    48  	return &reject{cache: lru.New(lru.WithCapacity[string, object](100)), max: maxDelay, internal: interval}
    49  }
    50  
    51  func (r *reject) delay(addr netapi.Address) time.Duration {
    52  	if r.max == 0 {
    53  		return 0
    54  	}
    55  	z, ok := r.cache.Load(addr.Hostname())
    56  	if !ok || !z.time.Add(time.Duration(r.internal)*time.Second).After(time.Now()) {
    57  		r.cache.Add(addr.Hostname(), object{time: time.Now(), delay: 0, times: 0})
    58  		return 0
    59  	}
    60  
    61  	if z.times < 7 {
    62  		z.times++
    63  	}
    64  
    65  	if z.times >= 7 && z.delay < time.Second*time.Duration(r.max) {
    66  		z.delay = z.delay + time.Second
    67  	}
    68  
    69  	time.Sleep(z.delay)
    70  	r.cache.Add(addr.Hostname(), object{z.times, time.Now(), z.delay})
    71  	return z.delay
    72  }
    73  
    74  func (r *reject) Conn(_ context.Context, addr netapi.Address) (net.Conn, error) {
    75  	return nil, fmt.Errorf("blocked address tcp[%v], delay %v", addr, r.delay(addr))
    76  }
    77  
    78  func (r *reject) PacketConn(_ context.Context, addr netapi.Address) (net.PacketConn, error) {
    79  	return nil, fmt.Errorf("blocked address udp[%v]. delay %v", addr, r.delay(addr))
    80  }