From 12c8845fef0dd7dda42f38936fb02fc8a59e165f Mon Sep 17 00:00:00 2001
From: aler9 <46489434+aler9@users.noreply.github.com>
Date: Thu, 10 Nov 2022 22:15:34 +0100
Subject: [PATCH] client: add variables to store received and sent bytes

---
 client.go     | 40 ++++++++++++++++++++++++++--------------
 clientudpl.go | 19 +++++++++++++++----
 2 files changed, 41 insertions(+), 18 deletions(-)

diff --git a/client.go b/client.go
index 24eb7b0f..ab9e7804 100644
--- a/client.go
+++ b/client.go
@@ -22,6 +22,7 @@ import (
 
 	"github.com/aler9/gortsplib/pkg/auth"
 	"github.com/aler9/gortsplib/pkg/base"
+	"github.com/aler9/gortsplib/pkg/bytecounter"
 	"github.com/aler9/gortsplib/pkg/conn"
 	"github.com/aler9/gortsplib/pkg/headers"
 	"github.com/aler9/gortsplib/pkg/liberrors"
@@ -215,6 +216,10 @@ type Client struct {
 	UserAgent string
 	// disable automatic RTCP sender reports.
 	DisableRTCPSenderReports bool
+	// pointer to a variable that stores received bytes.
+	BytesReceived *uint64
+	// pointer to a variable that stores sent bytes.
+	BytesSent *uint64
 
 	//
 	// system functions (all optional)
@@ -324,6 +329,12 @@ func (c *Client) Start(scheme string, host string) error {
 	if c.UserAgent == "" {
 		c.UserAgent = "gortsplib"
 	}
+	if c.BytesReceived == nil {
+		c.BytesReceived = new(uint64)
+	}
+	if c.BytesSent == nil {
+		c.BytesSent = new(uint64)
+	}
 
 	// system functions
 	if c.DialContext == nil {
@@ -966,23 +977,22 @@ func (c *Client) connOpen() error {
 		return err
 	}
 
-	c.nconn = func() net.Conn {
-		if c.scheme == "rtsps" {
-			tlsConfig := c.TLSConfig
+	if c.scheme == "rtsps" {
+		tlsConfig := c.TLSConfig
 
-			if tlsConfig == nil {
-				tlsConfig = &tls.Config{}
-			}
-
-			host, _, _ := net.SplitHostPort(c.host)
-			tlsConfig.ServerName = host
-
-			return tls.Client(nconn, tlsConfig)
+		if tlsConfig == nil {
+			tlsConfig = &tls.Config{}
 		}
-		return nconn
-	}()
 
-	c.conn = conn.NewConn(c.nconn)
+		host, _, _ := net.SplitHostPort(c.host)
+		tlsConfig.ServerName = host
+
+		nconn = tls.Client(nconn, tlsConfig)
+	}
+
+	c.nconn = nconn
+	bc := bytecounter.New(c.nconn, c.BytesReceived, c.BytesSent)
+	c.conn = conn.NewConn(bc)
 
 	c.connCloserStart()
 	return nil
@@ -1856,6 +1866,8 @@ func (c *Client) runWriter() {
 		}
 		data := tmp.(trackTypePayload)
 
+		atomic.AddUint64(c.BytesSent, uint64(len(data.payload)))
+
 		writeFunc(data.trackID, data.isRTP, data.payload)
 	}
 }
diff --git a/clientudpl.go b/clientudpl.go
index ae8a1162..be4e5136 100644
--- a/clientudpl.go
+++ b/clientudpl.go
@@ -191,7 +191,11 @@ func (u *clientUDPListener) runReader(forPlay bool) {
 }
 
 func (u *clientUDPListener) processPlayRTP(now time.Time, payload []byte) {
-	if len(payload) == (maxPacketSize + 1) {
+	plen := len(payload)
+
+	atomic.AddUint64(u.c.BytesReceived, uint64(plen))
+
+	if plen == (maxPacketSize + 1) {
 		u.c.OnDecodeError(fmt.Errorf("RTP packet is too big to be read with UDP"))
 		return
 	}
@@ -222,7 +226,11 @@ func (u *clientUDPListener) processPlayRTP(now time.Time, payload []byte) {
 }
 
 func (u *clientUDPListener) processPlayRTCP(now time.Time, payload []byte) {
-	if len(payload) == (maxPacketSize + 1) {
+	plen := len(payload)
+
+	atomic.AddUint64(u.c.BytesReceived, uint64(plen))
+
+	if plen == (maxPacketSize + 1) {
 		u.c.OnDecodeError(fmt.Errorf("RTCP packet is too big to be read with UDP"))
 		return
 	}
@@ -243,7 +251,11 @@ func (u *clientUDPListener) processPlayRTCP(now time.Time, payload []byte) {
 }
 
 func (u *clientUDPListener) processRecordRTCP(now time.Time, payload []byte) {
-	if len(payload) == (maxPacketSize + 1) {
+	plen := len(payload)
+
+	atomic.AddUint64(u.c.BytesReceived, uint64(plen))
+
+	if plen == (maxPacketSize + 1) {
 		u.c.OnDecodeError(fmt.Errorf("RTCP packet is too big to be read with UDP"))
 		return
 	}
@@ -265,7 +277,6 @@ func (u *clientUDPListener) processRecordRTCP(now time.Time, payload []byte) {
 func (u *clientUDPListener) write(payload []byte) error {
 	// no mutex is needed here since Write() has an internal lock.
 	// https://github.com/golang/go/issues/27203#issuecomment-534386117
-
 	u.pc.SetWriteDeadline(time.Now().Add(u.c.WriteTimeout))
 	_, err := u.pc.WriteTo(payload, u.writeAddr)
 	return err