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() + } } }