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

     1  package statistics
     2  
     3  import (
     4  	"context"
     5  	"net"
     6  
     7  	"github.com/Asutorufa/yuhaiin/pkg/components/shunt"
     8  	"github.com/Asutorufa/yuhaiin/pkg/log"
     9  	"github.com/Asutorufa/yuhaiin/pkg/net/netapi"
    10  	"github.com/Asutorufa/yuhaiin/pkg/net/proxy/direct"
    11  	"github.com/Asutorufa/yuhaiin/pkg/protos/statistic"
    12  	gs "github.com/Asutorufa/yuhaiin/pkg/protos/statistic/grpc"
    13  	"github.com/Asutorufa/yuhaiin/pkg/utils/cache"
    14  	"github.com/Asutorufa/yuhaiin/pkg/utils/convert"
    15  	"github.com/Asutorufa/yuhaiin/pkg/utils/id"
    16  	"github.com/Asutorufa/yuhaiin/pkg/utils/slice"
    17  	"github.com/Asutorufa/yuhaiin/pkg/utils/syncmap"
    18  	"google.golang.org/protobuf/types/known/emptypb"
    19  )
    20  
    21  type Connections struct {
    22  	gs.UnimplementedConnectionsServer
    23  
    24  	netapi.Proxy
    25  	idSeed id.IDGenerator
    26  
    27  	connStore syncmap.SyncMap[uint64, connection]
    28  
    29  	Cache *Cache
    30  
    31  	notify *notify
    32  }
    33  
    34  func NewConnStore(cache *cache.Cache, dialer netapi.Proxy) *Connections {
    35  	if dialer == nil {
    36  		dialer = direct.Default
    37  	}
    38  
    39  	return &Connections{
    40  		Proxy:  dialer,
    41  		Cache:  NewCache(cache),
    42  		notify: newNotify(),
    43  	}
    44  }
    45  
    46  func (c *Connections) Notify(_ *emptypb.Empty, s gs.Connections_NotifyServer) error {
    47  	id := c.notify.register(s, c.connStore.ValueSlice()...)
    48  	defer c.notify.unregister(id)
    49  	log.Debug("new notify client", "id", id)
    50  	<-s.Context().Done()
    51  	log.Debug("remove notify client", "id", id)
    52  	return s.Context().Err()
    53  }
    54  
    55  func (c *Connections) Conns(context.Context, *emptypb.Empty) (*gs.NotifyNewConnections, error) {
    56  	return &gs.NotifyNewConnections{
    57  		Connections: slice.To(c.connStore.ValueSlice(),
    58  			func(c connection) *statistic.Connection { return c.Info() }),
    59  	}, nil
    60  }
    61  
    62  func (c *Connections) CloseConn(_ context.Context, x *gs.NotifyRemoveConnections) (*emptypb.Empty, error) {
    63  	for _, x := range x.Ids {
    64  		if z, ok := c.connStore.Load(x); ok {
    65  			z.Close()
    66  		}
    67  	}
    68  	return &emptypb.Empty{}, nil
    69  }
    70  
    71  func (c *Connections) Close() error {
    72  	c.notify.Close()
    73  	c.connStore.Range(func(key uint64, v connection) bool {
    74  		v.Close()
    75  		return true
    76  	})
    77  
    78  	c.Cache.Close()
    79  	return nil
    80  }
    81  
    82  func (c *Connections) Total(context.Context, *emptypb.Empty) (*gs.TotalFlow, error) {
    83  	return &gs.TotalFlow{
    84  		Download: c.Cache.LoadDownload(),
    85  		Upload:   c.Cache.LoadUpload(),
    86  	}, nil
    87  }
    88  
    89  func (c *Connections) Remove(id uint64) {
    90  	if z, ok := c.connStore.LoadAndDelete(id); ok {
    91  		log.Debug("close conn", "id", z.ID())
    92  	}
    93  
    94  	c.notify.pubRemoveConns(id)
    95  }
    96  
    97  func (c *Connections) storeConnection(o connection) {
    98  	c.connStore.Store(o.ID(), o)
    99  	c.notify.pubNewConns(o)
   100  	log.Debug("new conn",
   101  		"id", o.ID(),
   102  		"addr", o.Info().Addr,
   103  		"src", o.Info().Extra[(netapi.SourceKey{}).String()],
   104  		"network", o.Info().Type.ConnType,
   105  		"outbound", o.Info().Extra["Outbound"],
   106  	)
   107  }
   108  
   109  func (c *Connections) PacketConn(ctx context.Context, addr netapi.Address) (net.PacketConn, error) {
   110  	con, err := c.Proxy.PacketConn(ctx, addr)
   111  	if err != nil {
   112  		return nil, err
   113  	}
   114  
   115  	z := &packetConn{con, c.getConnection(ctx, con, addr), c}
   116  
   117  	c.storeConnection(z)
   118  	return z, nil
   119  }
   120  
   121  func getRemote(con any) string {
   122  	r, ok := con.(interface{ RemoteAddr() net.Addr })
   123  	if ok {
   124  		// https://github.com/google/gvisor/blob/a9bdef23522b5a2ff2a7ec07c3e0573885b46ecb/pkg/tcpip/adapters/gonet/gonet.go#L457
   125  		// gvisor TCPConn will return nil remoteAddr
   126  		if addr := r.RemoteAddr(); addr != nil {
   127  			return addr.String()
   128  		}
   129  	}
   130  
   131  	return ""
   132  }
   133  
   134  func getRealAddr(store netapi.Store, addr netapi.Address) string {
   135  	z, ok := store.Get(shunt.DOMAIN_MARK_KEY{})
   136  	if ok {
   137  		if s, ok := convert.ToString(z); ok {
   138  			return s
   139  		}
   140  	}
   141  
   142  	return addr.String()
   143  }
   144  
   145  func (c *Connections) getConnection(ctx context.Context, conn interface{ LocalAddr() net.Addr }, addr netapi.Address) *statistic.Connection {
   146  	store := netapi.StoreFromContext(ctx)
   147  
   148  	connection := &statistic.Connection{
   149  		Id:   c.idSeed.Generate(),
   150  		Addr: getRealAddr(store, addr),
   151  		Type: &statistic.NetType{
   152  			ConnType:       addr.NetworkType(),
   153  			UnderlyingType: statistic.Type(statistic.Type_value[conn.LocalAddr().Network()]),
   154  		},
   155  		Extra: convert.ToStringMap(store),
   156  	}
   157  
   158  	if out := getRemote(conn); out != "" {
   159  		connection.Extra["Outbound"] = out
   160  	}
   161  	return connection
   162  }
   163  
   164  func (c *Connections) Conn(ctx context.Context, addr netapi.Address) (net.Conn, error) {
   165  	con, err := c.Proxy.Conn(ctx, addr)
   166  	if err != nil {
   167  		return nil, err
   168  	}
   169  
   170  	z := &conn{con, c.getConnection(ctx, con, addr), c}
   171  	c.storeConnection(z)
   172  	return z, nil
   173  }