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

     1  package netapi
     2  
     3  import (
     4  	"context"
     5  	"net"
     6  	"net/netip"
     7  	"sync"
     8  
     9  	"github.com/Asutorufa/yuhaiin/pkg/protos/statistic"
    10  )
    11  
    12  type ProcessDumper interface {
    13  	ProcessName(network string, src, dst Address) (string, error)
    14  }
    15  
    16  type Proxy interface {
    17  	StreamProxy
    18  	PacketProxy
    19  	Dispatch(context.Context, Address) (Address, error)
    20  }
    21  
    22  type StreamProxy interface {
    23  	Conn(context.Context, Address) (net.Conn, error)
    24  }
    25  
    26  type PacketProxy interface {
    27  	PacketConn(context.Context, Address) (net.PacketConn, error)
    28  }
    29  
    30  type Port interface {
    31  	Port() uint16
    32  	String() string
    33  }
    34  
    35  type Type uint8
    36  
    37  func (t Type) String() string {
    38  	switch t {
    39  	case FQDN:
    40  		return "FQDN"
    41  	case IP:
    42  		return "IP"
    43  	case UNIX:
    44  		return "UNIX"
    45  	case EMPTY:
    46  		return "EMPTY"
    47  	default:
    48  		return "UNKNOWN"
    49  	}
    50  }
    51  
    52  const (
    53  	FQDN  Type = 1
    54  	IP    Type = 2
    55  	UNIX  Type = 3
    56  	EMPTY Type = 4
    57  )
    58  
    59  type AddressSrc int32
    60  
    61  const (
    62  	AddressSrcEmpty AddressSrc = 0
    63  	AddressSrcDNS   AddressSrc = 1
    64  )
    65  
    66  type Result[T any] struct {
    67  	V   T
    68  	Err error
    69  }
    70  
    71  func NewErrResult[T any](err error) Result[T] {
    72  	return Result[T]{Err: err}
    73  }
    74  
    75  func NewResult[T any](v T) Result[T] {
    76  	return Result[T]{V: v}
    77  }
    78  
    79  type Address interface {
    80  	// Hostname return hostname of address, eg: www.example.com, 127.0.0.1, ff::ff
    81  	Hostname() string
    82  	IPs(context.Context) ([]net.IP, error)
    83  	// IP return net.IP, if address is ip else resolve the domain and return one of ips
    84  	IP(context.Context) (net.IP, error)
    85  	AddrPort(context.Context) Result[netip.AddrPort]
    86  	UDPAddr(context.Context) Result[*net.UDPAddr]
    87  	TCPAddr(context.Context) Result[*net.TCPAddr]
    88  	// Port return port of address
    89  	Port() Port
    90  	// Type return type of address, domain or ip
    91  	Type() Type
    92  	NetworkType() statistic.Type
    93  
    94  	net.Addr
    95  
    96  	SetSrc(AddressSrc)
    97  	// SetResolver will use call IP(), IPHost(), UDPAddr(), TCPAddr()
    98  	SetResolver(_ Resolver)
    99  	PreferIPv6(b bool)
   100  	PreferIPv4(b bool)
   101  	// OverrideHostname clone address(exclude Values) and change hostname
   102  	OverrideHostname(string) Address
   103  	OverridePort(Port) Address
   104  
   105  	IsFqdn() bool
   106  }
   107  
   108  type Store interface {
   109  	Add(k, v any) Store
   110  	Get(k any) (any, bool)
   111  	Range(func(k, v any) bool)
   112  	Map() map[any]any
   113  }
   114  
   115  type storeKey struct{}
   116  
   117  func StoreFromContext(ctx context.Context) Store {
   118  	store, ok := ctx.Value(storeKey{}).(Store)
   119  	if !ok {
   120  		return EmptyStore
   121  	}
   122  
   123  	return store
   124  }
   125  
   126  func NewStore(ctx context.Context) context.Context {
   127  	// ! it must return new context with store
   128  	// ! otherwise dns across proxy will use same store with proxy
   129  	return context.WithValue(ctx, storeKey{}, &store{store: make(map[any]any)})
   130  }
   131  
   132  func Get[T any](ctx context.Context, k any) (t T, _ bool) {
   133  	v, ok := StoreFromContext(ctx).Get(k)
   134  	if !ok {
   135  		return t, false
   136  	}
   137  
   138  	t, ok = v.(T)
   139  
   140  	return t, ok
   141  }
   142  
   143  func GetDefault[T any](ctx context.Context, k any, Default T) T {
   144  	v, ok := StoreFromContext(ctx).Get(k)
   145  	if !ok {
   146  		return Default
   147  	}
   148  
   149  	t, ok := v.(T)
   150  	if !ok {
   151  		return Default
   152  	}
   153  
   154  	return t
   155  }
   156  
   157  var EmptyStore = &emptyStore{}
   158  
   159  type emptyStore struct{}
   160  
   161  func (e *emptyStore) Add(k, v any) Store      { return e }
   162  func (*emptyStore) Get(k any) (any, bool)     { return nil, false }
   163  func (*emptyStore) Range(func(k, v any) bool) {}
   164  func (s *emptyStore) Map() map[any]any        { return map[any]any{} }
   165  
   166  type store struct {
   167  	mu    sync.RWMutex
   168  	store map[any]any
   169  }
   170  
   171  func (s *store) Add(key, value any) Store {
   172  	s.mu.Lock()
   173  	defer s.mu.Unlock()
   174  	if s.store == nil {
   175  		s.store = make(map[any]any)
   176  	}
   177  
   178  	s.store[key] = value
   179  
   180  	return s
   181  }
   182  
   183  func (s *store) Map() map[any]any {
   184  	s.mu.RLock()
   185  	defer s.mu.RUnlock()
   186  	if s.store == nil {
   187  		return map[any]any{}
   188  	}
   189  	return s.store
   190  }
   191  
   192  func (s *store) Get(key any) (any, bool) {
   193  	s.mu.RLock()
   194  	defer s.mu.RUnlock()
   195  	if s.store == nil {
   196  		return nil, false
   197  	}
   198  	v, ok := s.store[key]
   199  	return v, ok
   200  }
   201  
   202  func (s *store) Range(f func(k, v any) bool) {
   203  	s.mu.RLock()
   204  	defer s.mu.RUnlock()
   205  
   206  	for k, v := range s.store {
   207  		if !f(k, v) {
   208  			break
   209  		}
   210  	}
   211  }