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

     1  package sniffy
     2  
     3  import (
     4  	"net"
     5  
     6  	"github.com/Asutorufa/yuhaiin/pkg/net/netapi"
     7  	"github.com/Asutorufa/yuhaiin/pkg/net/sniffy/bittorrent"
     8  	"github.com/Asutorufa/yuhaiin/pkg/protos/config/bypass"
     9  	"github.com/Asutorufa/yuhaiin/pkg/utils/pool"
    10  )
    11  
    12  type entry[T any] struct {
    13  	name    string
    14  	checker func([]byte) (T, bool)
    15  }
    16  
    17  type Sniffier[T any] struct {
    18  	streamChecker []entry[T]
    19  	packetChecker []entry[T]
    20  }
    21  
    22  func New() *Sniffier[bypass.Mode] {
    23  	return &Sniffier[bypass.Mode]{
    24  		streamChecker: []entry[bypass.Mode]{
    25  			{
    26  				name: "bittorrent",
    27  				checker: func(b []byte) (bypass.Mode, bool) {
    28  					_, err := bittorrent.SniffBittorrent(b)
    29  					if err == nil {
    30  						return bypass.Mode_direct, true
    31  					}
    32  
    33  					return bypass.Mode_bypass, false
    34  				},
    35  			},
    36  		},
    37  
    38  		packetChecker: []entry[bypass.Mode]{
    39  			{
    40  				name: "bittorrent_utp",
    41  				checker: func(b []byte) (bypass.Mode, bool) {
    42  					_, err := bittorrent.SniffUTP(b)
    43  					if err == nil {
    44  						return bypass.Mode_direct, true
    45  					}
    46  
    47  					return bypass.Mode_bypass, false
    48  				},
    49  			},
    50  		},
    51  	}
    52  }
    53  
    54  func (s *Sniffier[T]) Packet(b []byte) (T, string, bool) {
    55  	for _, c := range s.packetChecker {
    56  		t, ok := c.checker(b)
    57  		if ok {
    58  			return t, c.name, ok
    59  		}
    60  	}
    61  
    62  	return *new(T), "", false
    63  }
    64  
    65  func (s *Sniffier[T]) Stream(c net.Conn) (net.Conn, T, string, bool) {
    66  	buf := pool.GetBytesBuffer(pool.DefaultSize)
    67  
    68  	n, _ := buf.ReadFrom(c)
    69  
    70  	if n <= 0 {
    71  		buf.Free()
    72  		return c, *new(T), "", false
    73  	}
    74  
    75  	c = netapi.NewPrefixBytesConn(c, buf)
    76  
    77  	for _, ck := range s.streamChecker {
    78  		t, ok := ck.checker(buf.Bytes())
    79  		if ok {
    80  			return c, t, ck.name, ok
    81  		}
    82  	}
    83  
    84  	return c, *new(T), "", false
    85  }