github.com/Asutorufa/yuhaiin@v0.3.6-0.20240502055049-7984da7023a0/pkg/components/inbound/handler.go (about)

     1  package inbound
     2  
     3  import (
     4  	"context"
     5  	"errors"
     6  	"fmt"
     7  	"time"
     8  
     9  	"github.com/Asutorufa/yuhaiin/pkg/components/shunt"
    10  	"github.com/Asutorufa/yuhaiin/pkg/log"
    11  	"github.com/Asutorufa/yuhaiin/pkg/net/nat"
    12  	"github.com/Asutorufa/yuhaiin/pkg/net/netapi"
    13  	"github.com/Asutorufa/yuhaiin/pkg/net/sniffy"
    14  	"github.com/Asutorufa/yuhaiin/pkg/protos/config/bypass"
    15  	"github.com/Asutorufa/yuhaiin/pkg/utils/relay"
    16  )
    17  
    18  var Timeout = time.Second * 20
    19  
    20  type handler struct {
    21  	dialer     netapi.Proxy
    22  	dnsHandler netapi.DNSServer
    23  	table      *nat.Table
    24  
    25  	sniffyEnabled bool
    26  	sniffer       *sniffy.Sniffier[bypass.Mode]
    27  }
    28  
    29  func NewHandler(dialer netapi.Proxy, dnsHandler netapi.DNSServer) *handler {
    30  	h := &handler{
    31  		dialer:        dialer,
    32  		table:         nat.NewTable(dialer),
    33  		dnsHandler:    dnsHandler,
    34  		sniffer:       sniffy.New(),
    35  		sniffyEnabled: true,
    36  	}
    37  
    38  	return h
    39  }
    40  
    41  func (s *handler) Stream(ctx context.Context, meta *netapi.StreamMeta) {
    42  	go func() {
    43  		if err := s.stream(ctx, meta); err != nil {
    44  			if errors.Is(err, netapi.ErrBlocked) {
    45  				log.Debug("blocked", "msg", err)
    46  			} else {
    47  				log.Error("stream", "error", err)
    48  			}
    49  		}
    50  	}()
    51  }
    52  
    53  func (s *handler) stream(ctx context.Context, meta *netapi.StreamMeta) error {
    54  	ctx, cancel := context.WithTimeout(ctx, Timeout)
    55  	defer cancel()
    56  
    57  	ctx = netapi.NewStore(ctx)
    58  	defer meta.Src.Close()
    59  
    60  	dst := meta.Address
    61  	store := netapi.StoreFromContext(ctx)
    62  
    63  	store.Add(netapi.SourceKey{}, meta.Source).
    64  		Add(netapi.DestinationKey{}, meta.Destination)
    65  	if meta.Inbound != nil {
    66  		store.Add(netapi.InboundKey{}, meta.Inbound)
    67  	}
    68  
    69  	if s.sniffyEnabled {
    70  		src, mode, name, ok := s.sniffer.Stream(meta.Src)
    71  		if ok {
    72  			store.
    73  				Add("Protocol", name).
    74  				Add(shunt.ForceModeKey{}, mode)
    75  		}
    76  		defer src.Close()
    77  
    78  		meta.Src = src
    79  	}
    80  
    81  	remote, err := s.dialer.Conn(ctx, dst)
    82  	if err != nil {
    83  		return fmt.Errorf("dial %s failed: %w", dst, err)
    84  	}
    85  	defer remote.Close()
    86  
    87  	relay.Relay(meta.Src, remote)
    88  	return nil
    89  }
    90  
    91  func (s *handler) Packet(ctx context.Context, pack *netapi.Packet) {
    92  	go func() {
    93  		ctx, cancel := context.WithTimeout(ctx, Timeout)
    94  		defer cancel()
    95  
    96  		ctx = netapi.NewStore(ctx)
    97  		store := netapi.StoreFromContext(ctx)
    98  
    99  		if s.sniffyEnabled {
   100  			mode, name, ok := s.sniffer.Packet(pack.Payload.Bytes())
   101  			if ok {
   102  				store.
   103  					Add("Protocol", name).
   104  					Add(shunt.ForceModeKey{}, mode)
   105  			}
   106  		}
   107  
   108  		if err := s.table.Write(ctx, pack); err != nil {
   109  			log.Error("packet", "error", err)
   110  		}
   111  	}()
   112  }
   113  
   114  func (s *handler) Close() error { return s.table.Close() }