From e3c56e1855600f099d740e814f3d7f2f197cce40 Mon Sep 17 00:00:00 2001
From: lhridder <lhridder@gmail.com>
Date: Thu, 14 Jul 2022 00:53:07 +0200
Subject: [PATCH] add bandwith track feature

---
 README.md            |  7 +++++--
 cmd/infrared/main.go |  8 ++++++++
 config.go            |  2 ++
 gateway.go           | 21 +++++++++++++++++++++
 proxy.go             | 14 +++++++++++---
 5 files changed, 47 insertions(+), 5 deletions(-)

diff --git a/README.md b/README.md
index 633957e..d891c19 100644
--- a/README.md
+++ b/README.md
@@ -19,6 +19,7 @@ fork from [haveachin/infrared](https://github.com/haveachin/infrared)
 - Global .json config
 - Removed docker and callback features
 - Status packet caching
+- Bandwith usage tracking for proxy configs through prometheus
 
 ## Command-Line Flags
 
@@ -51,7 +52,8 @@ fork from [haveachin/infrared](https://github.com/haveachin/infrared)
   "rejoinMessage": "Please rejoin to verify your connection.",
   "underAttack": false,
   "debug": false,
-  "connectionTreshold": 50
+  "connectionTreshold": 50,
+  "trackBandwith": false
 }
 ```
 Values can be left out if they don't deviate from the default, a config.json with just `{}` is still required for startup.
@@ -75,7 +77,8 @@ Values can be left out if they don't deviate from the default, a config.json wit
 - `rejoinMessage` what text response should be sent when a player needs to rejoin to verify they're not a bot.
 - `underAttack` if the instance should permanently be in attack mode.
 - `debug` if debug logs should be enabled.
-- `connectionTreshold` at what amount of packets per second the underAttack mode should trigger.`
+- `connectionTreshold` at what amount of packets per second the underAttack mode should trigger.
+- `trackBandwith` whether or not bandwith usage should be tracked in prometheus (requires prometheusEnabled).
 
 
 ## Proxy Config
diff --git a/cmd/infrared/main.go b/cmd/infrared/main.go
index c4c99f1..58b57a8 100644
--- a/cmd/infrared/main.go
+++ b/cmd/infrared/main.go
@@ -123,6 +123,14 @@ func main() {
 			log.Println(err)
 			return
 		}
+
+		if infrared.Config.TrackBandwith {
+			go func() {
+				for {
+					gateway.TrackBandwith()
+				}
+			}()
+		}
 	}
 
 	if !infrared.Config.UnderAttack {
diff --git a/config.go b/config.go
index 26615b5..cacf3d5 100644
--- a/config.go
+++ b/config.go
@@ -60,6 +60,7 @@ type GlobalConfig struct {
 	UnderAttack            bool     `json:"underAttack"`
 	Debug                  bool     `json:"debug"`
 	ConnectionTreshold     int      `json:"connectionTreshold"`
+	TrackBandwith          bool     `json:"trackBandwith"`
 }
 
 var Config GlobalConfig
@@ -85,6 +86,7 @@ var DefaultConfig = GlobalConfig{
 	UnderAttack:            false,
 	Debug:                  false,
 	ConnectionTreshold:     50,
+	TrackBandwith:          false,
 }
 
 func (cfg *ProxyConfig) Dialer() (*Dialer, error) {
diff --git a/gateway.go b/gateway.go
index 1d7ebf5..29153eb 100644
--- a/gateway.go
+++ b/gateway.go
@@ -43,6 +43,10 @@ var (
 		Name: "infrared_underAttack",
 		Help: "Is the proxy under attack",
 	})
+	usedBandwith = promauto.NewCounterVec(prometheus.CounterOpts{
+		Name: "infrared_used_bandwith",
+		Help: "The total number of used bytes of bandwith per proxy",
+	}, []string{"host"})
 	ctx = context.Background()
 )
 
@@ -223,6 +227,10 @@ func (gateway *Gateway) RegisterProxy(proxy *Proxy) error {
 
 	playersConnected.WithLabelValues(proxy.DomainName())
 
+	if Config.TrackBandwith {
+		usedBandwith.WithLabelValues(proxy.DomainName())
+	}
+
 	// Check if a gate is already listening to the Proxy address
 	addr := proxy.ListenTo()
 	if _, ok := gateway.listeners.Load(addr); ok {
@@ -754,6 +762,19 @@ func (gateway *Gateway) ClearCps() {
 	time.Sleep(time.Second)
 }
 
+func (gateway *Gateway) TrackBandwith() {
+	gateway.Proxies.Range(func(k, v interface{}) bool {
+		proxy := v.(*Proxy)
+		name := proxy.DomainName()
+		proxy.mu.Lock()
+		usedBandwith.WithLabelValues(name).Add(float64(proxy.usedBandwith))
+		proxy.usedBandwith = 0
+		proxy.mu.Unlock()
+		return false
+	})
+	time.Sleep(5 * time.Second)
+}
+
 func contains(s []string, str string) bool {
 	for _, v := range s {
 		if v == str {
diff --git a/proxy.go b/proxy.go
index 981f5d2..1012cd6 100644
--- a/proxy.go
+++ b/proxy.go
@@ -39,6 +39,8 @@ type Proxy struct {
 	cacheStatusTime   time.Time
 	cacheResponse     status.ClientBoundResponse
 	cacheOnlineStatus bool
+
+	usedBandwith int
 }
 
 func (proxy *Proxy) DomainNames() []string {
@@ -185,8 +187,8 @@ func (proxy *Proxy) handleLoginConnection(conn Conn, session Session) error {
 	playersConnected.With(prometheus.Labels{"host": proxyDomain}).Inc()
 	defer playersConnected.With(prometheus.Labels{"host": proxyDomain}).Dec()
 
-	go pipe(rconn, conn)
-	pipe(conn, rconn)
+	go pipe(rconn, conn, proxy)
+	pipe(conn, rconn, proxy)
 	return nil
 }
 
@@ -339,7 +341,7 @@ func (proxy *Proxy) handleStatusConnection(conn Conn, session Session) error {
 	return nil
 }
 
-func pipe(src, dst Conn) {
+func pipe(src, dst Conn, proxy *Proxy) {
 	buffer := make([]byte, 0xffff)
 
 	for {
@@ -354,6 +356,12 @@ func pipe(src, dst Conn) {
 		if err != nil {
 			return
 		}
+
+		if Config.TrackBandwith {
+			proxy.mu.Lock()
+			proxy.usedBandwith = proxy.usedBandwith + len(data)
+			proxy.mu.Unlock()
+		}
 	}
 }