support stats counter in dispatcher

app/dispatcher/default.go

@@ -11,6 +11,7 @@ import (
 	"v2ray.com/core/common"
 	"v2ray.com/core/common/buf"
 	"v2ray.com/core/common/net"
+	"v2ray.com/core/common/protocol"
 	"v2ray.com/core/proxy"
 	"v2ray.com/core/transport/ray"
 )
@@ -23,6 +24,8 @@ var (
 type DefaultDispatcher struct {
 	ohm    core.OutboundHandlerManager
 	router core.Router
+	policy core.PolicyManager
+	stats  core.StatManager
 }
 
 // NewDefaultDispatcher create a new DefaultDispatcher.
@@ -31,6 +34,8 @@ func NewDefaultDispatcher(ctx context.Context, config *Config) (*DefaultDispatch
 	d := &DefaultDispatcher{
 		ohm:    v.OutboundHandlerManager(),
 		router: v.Router(),
+		policy: v.PolicyManager(),
+		stats:  v.Stats(),
 	}
 
 	if err := v.RegisterFeature((*core.Dispatcher)(nil), d); err != nil {
@@ -47,6 +52,10 @@ func (*DefaultDispatcher) Start() error {
 // Close implements app.Application.
 func (*DefaultDispatcher) Close() error { return nil }
 
+func getStatsName(u *protocol.User) string {
+	return "user>traffic>" + u.Email
+}
+
 // Dispatch implements core.Dispatcher.
 func (d *DefaultDispatcher) Dispatch(ctx context.Context, destination net.Destination) (ray.InboundRay, error) {
 	if !destination.IsValid() {
@@ -54,7 +63,26 @@ func (d *DefaultDispatcher) Dispatch(ctx context.Context, destination net.Destin
 	}
 	ctx = proxy.ContextWithTarget(ctx, destination)
 
-	outbound := ray.New(ctx)
+	var rayOptions []ray.Option
+
+	user := protocol.UserFromContext(ctx)
+	if user != nil && len(user.Email) > 0 {
+		name := getStatsName(user)
+		c, err := d.stats.RegisterCounter(name)
+		if err != nil {
+			c = d.stats.GetCounter(name)
+		}
+		if c == nil {
+			newError("failed to get stats counter ", name).AtWarning().WithContext(ctx).WriteToLog()
+		}
+
+		p := d.policy.ForLevel(user.Level)
+		if p.Stats.EnablePerUser {
+			rayOptions = append(rayOptions, ray.WithStatCounter(c))
+		}
+	}
+
+	outbound := ray.New(ctx, rayOptions...)
 	snifferList := proxyman.ProtocolSniffersFromContext(ctx)
 	if destination.Address.Family().IsDomain() || len(snifferList) == 0 {
 		go d.routedDispatch(ctx, outbound, destination)