From 86b2b0fcec8a1444241c033d87cfa0497173436a Mon Sep 17 00:00:00 2001
From: Adam Snyder <armsnyder@gmail.com>
Date: Sat, 10 Dec 2022 15:41:58 -0800
Subject: [PATCH] New argument --max-packet-size

---
 Makefile                        |  1 +
 README.md                       |  1 +
 internal/collector/collector.go | 20 ++++++++++++--------
 main.go                         |  8 +++++++-
 4 files changed, 21 insertions(+), 9 deletions(-)

diff --git a/Makefile b/Makefile
index fbfacc7..c4b26ea 100644
--- a/Makefile
+++ b/Makefile
@@ -7,6 +7,7 @@ fix:
 	golangci-lint run --fix
 
 build:
+	go mod tidy
 	go build ./...
 
 test:
diff --git a/README.md b/README.md
index 87198b1..b17d1f6 100644
--- a/README.md
+++ b/README.md
@@ -38,6 +38,7 @@ Flag | Variable | Default | Help
 --path | A2S_EXPORTER_PATH | /metrics | Path for the metrics exporter.
 --namespace | A2S_EXPORTER_NAMESPACE | a2s | Namespace prefix for all exported a2s metrics.
 --a2s-only-metrics | A2S_EXPORTER_A2S_ONLY_METRICS | false | If true, excludes Go runtime and promhttp metrics.
+--max-packet-size | A2S_EXPORTER_MAX_PACKET_SIZE | 0 | Advanced option to set a non-standard max packet size of the A2S query server. Set to 0 to use the default max packet size.
 
 #### Special
 
diff --git a/internal/collector/collector.go b/internal/collector/collector.go
index e607344..9162f94 100644
--- a/internal/collector/collector.go
+++ b/internal/collector/collector.go
@@ -9,14 +9,15 @@ import (
 )
 
 type Collector struct {
-	addr   string
-	client *a2s.Client
-	descs  map[string]*prometheus.Desc
+	addr          string
+	clientOptions []func(*a2s.Client) error
+	client        *a2s.Client
+	descs         map[string]*prometheus.Desc
 }
 
 type adder func(name string, value float64, labelValues ...string)
 
-func New(namespace, addr string) *Collector {
+func New(namespace, addr string, clientOptions ...func(*a2s.Client) error) *Collector {
 	descs := make(map[string]*prometheus.Desc)
 
 	fullDesc := func(name, help string, labels []string) {
@@ -53,8 +54,9 @@ func New(namespace, addr string) *Collector {
 	playerDesc("player_the_ship_money", "Player's money in a The Ship server.")
 
 	return &Collector{
-		addr:  addr,
-		descs: descs,
+		addr:          addr,
+		clientOptions: clientOptions,
+		descs:         descs,
 	}
 }
 
@@ -97,7 +99,7 @@ func (c *Collector) queryInfo() (serverInfo *a2s.ServerInfo, playerInfo *a2s.Pla
 
 	// Lazy initialization of UDP client.
 	if c.client == nil {
-		c.client, err = a2s.NewClient(c.addr)
+		c.client, err = a2s.NewClient(c.addr, c.clientOptions...)
 		if err != nil {
 			fmt.Println("Could not create A2S client:", err)
 			return
@@ -115,7 +117,9 @@ func (c *Collector) queryInfo() (serverInfo *a2s.ServerInfo, playerInfo *a2s.Pla
 	// constructed with The Ship App ID.
 	playerClient := c.client
 	if a2s.AppID(serverInfo.ID) == a2s.App_TheShip {
-		playerClient, err = a2s.NewClient(c.addr, a2s.SetAppID(int32(serverInfo.ID)))
+		options := []func(*a2s.Client) error{a2s.SetAppID(int32(serverInfo.ID))}
+		options = append(options, c.clientOptions...)
+		playerClient, err = a2s.NewClient(c.addr, options...)
 		if err != nil {
 			fmt.Println("Could not create A2S client for The Ship player query:", err)
 			return
diff --git a/main.go b/main.go
index 75de2a3..ce77e4d 100644
--- a/main.go
+++ b/main.go
@@ -10,6 +10,7 @@ import (
 
 	"github.com/prometheus/client_golang/prometheus"
 	"github.com/prometheus/client_golang/prometheus/promhttp"
+	"github.com/rumblefrog/go-a2s"
 
 	"github.com/armsnyder/a2s-exporter/internal/collector"
 )
@@ -24,6 +25,7 @@ func main() {
 	path := flag.String("path", envOrDefault("A2S_EXPORTER_PATH", "/metrics"), "Path for the metrics exporter.")
 	namespace := flag.String("namespace", envOrDefault("A2S_EXPORTER_NAMESPACE", "a2s"), "Namespace prefix for all exported a2s metrics.")
 	a2sOnlyMetrics := flag.Bool("a2s-only-metrics", envOrDefaultBool("A2S_EXPORTER_A2S_ONLY_METRICS", false), "If true, skips exporting Go runtime metrics.")
+	maxPacketSize := flag.Int("max-packet-size", envOrDefaultInt("A2S_EXPORTER_MAX_PACKET_SIZE", 0), "Advanced option to set a non-standard max packet size of the A2S query server. Set to 0 to use the default max packet size.")
 	help := flag.Bool("h", false, "Show help.")
 	version := flag.Bool("version", false, "Show build version.")
 
@@ -57,7 +59,11 @@ func main() {
 	}
 
 	// Register A2S metrics.
-	registry.MustRegister(collector.New(*namespace, *address))
+	var clientOptions []func(*a2s.Client) error
+	if *maxPacketSize > 0 {
+		clientOptions = append(clientOptions, a2s.SetMaxPacketSize(uint32(*maxPacketSize)))
+	}
+	registry.MustRegister(collector.New(*namespace, *address, clientOptions...))
 
 	// Set up http handler.
 	handler := promhttp.HandlerFor(registry, promhttp.HandlerOpts{})