Add stats counter for outbounds

app/proxyman/outbound/handler.go

@@ -10,6 +10,8 @@ import (
 	"v2ray.com/core/common/net"
 	"v2ray.com/core/common/session"
 	"v2ray.com/core/features/outbound"
+	"v2ray.com/core/features/policy"
+	"v2ray.com/core/features/stats"
 	"v2ray.com/core/proxy"
 	"v2ray.com/core/transport"
 	"v2ray.com/core/transport/internet"
@@ -17,6 +19,31 @@ import (
 	"v2ray.com/core/transport/pipe"
 )
 
+func getStatCounter(v *core.Instance, tag string) (stats.Counter, stats.Counter) {
+	var uplinkCounter stats.Counter
+	var downlinkCounter stats.Counter
+
+	policy := v.GetFeature(policy.ManagerType()).(policy.Manager)
+	if len(tag) > 0 && policy.ForSystem().Stats.OutboundUplink {
+		statsManager := v.GetFeature(stats.ManagerType()).(stats.Manager)
+		name := "outbound>>>" + tag + ">>>traffic>>>uplink"
+		c, _ := stats.GetOrRegisterCounter(statsManager, name)
+		if c != nil {
+			uplinkCounter = c
+		}
+	}
+	if len(tag) > 0 && policy.ForSystem().Stats.OutboundDownlink {
+		statsManager := v.GetFeature(stats.ManagerType()).(stats.Manager)
+		name := "outbound>>>" + tag + ">>>traffic>>>downlink"
+		c, _ := stats.GetOrRegisterCounter(statsManager, name)
+		if c != nil {
+			downlinkCounter = c
+		}
+	}
+
+	return uplinkCounter, downlinkCounter
+}
+
 // Handler is an implements of outbound.Handler.
 type Handler struct {
 	tag             string
@@ -25,14 +52,19 @@ type Handler struct {
 	proxy           proxy.Outbound
 	outboundManager outbound.Manager
 	mux             *mux.ClientManager
+	uplinkCounter   stats.Counter
+	downlinkCounter stats.Counter
 }
 
 // NewHandler create a new Handler based on the given configuration.
 func NewHandler(ctx context.Context, config *core.OutboundHandlerConfig) (outbound.Handler, error) {
 	v := core.MustFromContext(ctx)
+	uplinkCounter, downlinkCounter := getStatCounter(v, config.Tag)
 	h := &Handler{
 		tag:             config.Tag,
 		outboundManager: v.GetFeature(outbound.ManagerType()).(outbound.Manager),
+		uplinkCounter:   uplinkCounter,
+		downlinkCounter: downlinkCounter,
 	}
 
 	if config.SenderSettings != nil {
@@ -148,7 +180,7 @@ func (h *Handler) Dial(ctx context.Context, dest net.Destination) (internet.Conn
 					conn = tls.Client(conn, tlsConfig)
 				}
 
-				return conn, nil
+				return h.getStatCouterConnection(conn), nil
 			}
 
 			newError("failed to get outbound handler with tag: ", tag).AtWarning().WriteToLog(session.ExportIDToError(ctx))
@@ -164,7 +196,19 @@ func (h *Handler) Dial(ctx context.Context, dest net.Destination) (internet.Conn
 		}
 	}
 
-	return internet.Dial(ctx, dest, h.streamSettings)
+	conn, err := internet.Dial(ctx, dest, h.streamSettings)
+	return h.getStatCouterConnection(conn), err
+}
+
+func (h *Handler) getStatCouterConnection(conn internet.Connection) (internet.Connection) {
+	if h.uplinkCounter != nil || h.downlinkCounter != nil {
+		return &internet.StatCouterConnection{
+			Connection:   conn,
+			ReadCounter:  h.downlinkCounter,
+			WriteCounter: h.uplinkCounter,
+		}
+	}
+	return conn
 }
 
 // GetOutbound implements proxy.GetOutbound.