Rework bytes sent/received metrics

metrics/service.go

@@ -28,7 +28,6 @@ var _ proxy.Listener = (*Service)(nil)
 
 type Service struct {
 	proxyRequestDuration              *prometheus.HistogramVec
-	proxyRequestData                  *prometheus.CounterVec
 	proxyNumberOfLiveConnecions       *prometheus.GaugeVec
 	proxyNumberOfIncommingConnections *prometheus.CounterVec
 	proxyNumberOfProcessedConnections *prometheus.CounterVec
@@ -44,15 +43,6 @@ func NewMetricsService() (*Service, error) {
 		return nil, err
 	}
 
-	proxyRequestData := prometheus.NewCounterVec(prometheus.CounterOpts{
-		Name: "proxy_request_data",
-		Help: "Proxy request data in bytes",
-	}, []string{"request_type", "direction"})
-
-	if err := prometheus.Register(proxyRequestData); err != nil {
-		return nil, err
-	}
-
 	proxyNumberOfLiveConnections := prometheus.NewGaugeVec(prometheus.GaugeOpts{
 		Name: "proxy_number_of_live_connections",
 		Help: "Number of currently live connections",
@@ -82,7 +72,6 @@ func NewMetricsService() (*Service, error) {
 
 	return &Service{
 		proxyRequestDuration:              proxyRequestDuration,
-		proxyRequestData:                  proxyRequestData,
 		proxyNumberOfLiveConnecions:       proxyNumberOfLiveConnections,
 		proxyNumberOfIncommingConnections: proxyNumberOfIncommingConnections,
 		proxyNumberOfProcessedConnections: proxyNumberOfProcessedConnections,
@@ -107,16 +96,6 @@ func (s *Service) ProxyHandlerMiddleware(next func(c *proxy.Context), proxyHandl
 			"request_type": proxyHandlerType,
 		}).Observe(time.Since(startTime).Seconds())
 
-		s.proxyRequestData.With(prometheus.Labels{
-			"request_type": proxyHandlerType,
-			"direction":    "sent",
-		}).Add(float64(c.BytesSent()))
-
-		s.proxyRequestData.With(prometheus.Labels{
-			"request_type": proxyHandlerType,
-			"direction":    "received",
-		}).Add(float64(c.BytesReceived()))
-
 		s.proxyNumberOfProcessedConnections.With(prometheus.Labels{
 			"request_type": proxyHandlerType,
 		}).Inc()

proxy/conn.go

@@ -0,0 +1,72 @@
+/*
+ * Copyright (C) 2024 The "MysteriumNetwork/openvpn-forwarder" Authors.
+ *
+ * This program is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU General Public License as published by
+ * the Free Software Foundation, either version 3 of the License, or
+ * (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU General Public License for more details.
+ *
+ * You should have received a copy of the GNU General Public License
+ * along with this program.  If not, see <http://www.gnu.org/licenses/>.
+ */
+
+package proxy
+
+import (
+	"net"
+
+	"github.com/prometheus/client_golang/prometheus"
+)
+
+var proxyRequestData = prometheus.NewCounterVec(prometheus.CounterOpts{
+	Name: "proxy_request_data",
+	Help: "Proxy request data in bytes",
+}, []string{"request_type", "direction"})
+
+func init() {
+	prometheus.MustRegister(proxyRequestData)
+}
+
+// NewConn returns net.Conn wrapped with metrics.
+func NewConn(conn net.Conn) *Conn {
+	return &Conn{
+		Conn: conn,
+	}
+}
+
+// Conn wraps net.Conn with intercepts of read/write events.
+type Conn struct {
+	net.Conn
+
+	bytesSent     int64
+	bytesReceived int64
+}
+
+// Read bytes from net.Conn and count read bytes
+func (c *Conn) Read(b []byte) (n int, err error) {
+	count, err := c.Conn.Read(b)
+
+	proxyRequestData.MustCurryWith(prometheus.Labels{
+		"request_type": "HTTPS",
+	}).WithLabelValues("received").Add(float64(count))
+
+	return count, err
+}
+
+// Write bytes to net.Conn and counts written bytes
+func (c *Conn) Write(b []byte) (n int, err error) {
+	count, err := c.Conn.Write(b)
+
+	proxyRequestData.MustCurryWith(prometheus.Labels{
+		"request_type": "HTTPS",
+	}).WithLabelValues("sent").Add(float64(count))
+
+	return count, err
+}
+
+var _ net.Conn = (*Conn)(nil)

proxy/context.go

@@ -29,15 +29,4 @@ type Context struct {
 
 	destinationHost    string
 	destinationAddress string
-
-	bytesSent     int64
-	bytesReceived int64
-}
-
-func (c *Context) BytesSent() int64 {
-	return c.bytesSent
-}
-
-func (c *Context) BytesReceived() int64 {
-	return c.bytesReceived
 }

proxy/server.go

@@ -263,11 +263,11 @@ func (s *proxyServer) serveTLS(c *Context) {
 	io.Copy(tlsConn, conn)
 }
 
-func (s *proxyServer) connectTo(c *Context, remoteHost string) (conn io.ReadWriteCloser, err error) {
+func (s *proxyServer) connectTo(c *Context, remoteHost string) (*Conn, error) {
 	domain := strings.Split(remoteHost, ":")
 	s.dt.Inc(domain[0])
 
-	conn, err = s.dialer.Dial("tcp", remoteHost)
+	conn, err := s.dialer.Dial("tcp", remoteHost)
 	if err != nil {
 		return nil, errors.Wrap(err, "failed to establish connection")
 	}
@@ -282,7 +282,7 @@ func (s *proxyServer) connectTo(c *Context, remoteHost string) (conn io.ReadWrit
 		}
 	}
 
-	return conn, nil
+	return NewConn(conn), nil
 }
 
 func (s *proxyServer) sendOnProxyConnectionAccept() {