From 39bbff0fbc4cf3017b63046366c0626dcb6dd529 Mon Sep 17 00:00:00 2001
From: Lewis Waddicor <lewisw@unity3d.com>
Date: Wed, 7 Jun 2023 15:42:32 +0100
Subject: [PATCH] feat: Add relay health (#35)

This adds a relay health flag for TF2E
---
 lib/svrquery/protocol/titanfall/types.go      | 7 +++++++
 lib/svrquery/protocol/titanfall/types_test.go | 9 ++++++++-
 2 files changed, 15 insertions(+), 1 deletion(-)

diff --git a/lib/svrquery/protocol/titanfall/types.go b/lib/svrquery/protocol/titanfall/types.go
index 2c04ca5..1eeb909 100644
--- a/lib/svrquery/protocol/titanfall/types.go
+++ b/lib/svrquery/protocol/titanfall/types.go
@@ -113,6 +113,7 @@ func (a HealthFlags) MarshalJSON() ([]byte, error) {
 		SlowServerFrames bool
 		Hitching         bool
 		DOS              bool
+		Relay            bool
 	}{}
 	obj.None = a.None()
 	obj.PacketLossIn = a.PacketLossIn()
@@ -122,6 +123,7 @@ func (a HealthFlags) MarshalJSON() ([]byte, error) {
 	obj.SlowServerFrames = a.SlowServerFrames()
 	obj.Hitching = a.Hitching()
 	obj.DOS = a.DOS()
+	obj.Relay = a.Relay()
 
 	return json.Marshal(obj)
 }
@@ -166,6 +168,11 @@ func (a HealthFlags) DOS() bool {
 	return (a>>6)&1 == 1
 }
 
+// Relay health flag
+func (a HealthFlags) Relay() bool {
+	return (a>>7)&1 == 1
+}
+
 // BasicInfo represents basic information contained in a query response.
 type BasicInfo struct {
 	Port                      uint16
diff --git a/lib/svrquery/protocol/titanfall/types_test.go b/lib/svrquery/protocol/titanfall/types_test.go
index dbbc3c0..ee87f88 100644
--- a/lib/svrquery/protocol/titanfall/types_test.go
+++ b/lib/svrquery/protocol/titanfall/types_test.go
@@ -18,6 +18,7 @@ func TestHealthFlags(t *testing.T) {
 		expSlowServerFrames bool
 		expHitching         bool
 		expDOS              bool
+		expRelay            bool
 	}{
 		{
 			input:   0,
@@ -51,10 +52,14 @@ func TestHealthFlags(t *testing.T) {
 			input:  1 << 6,
 			expDOS: true,
 		},
+		{
+			input:    1 << 7,
+			expRelay: true,
+		},
 	}
 
 	for _, tc := range testCases {
-		t.Run(fmt.Sprintf("verify 0b%b", tc.input), func(t *testing.T) {
+		t.Run(fmt.Sprintf("verify 0b%.32b", tc.input), func(t *testing.T) {
 			hf := HealthFlags(tc.input)
 			require.Equal(t, tc.expNone, hf.None())
 			require.Equal(t, tc.expPacketLossIn, hf.PacketLossIn())
@@ -63,6 +68,8 @@ func TestHealthFlags(t *testing.T) {
 			require.Equal(t, tc.expPacketChokedOut, hf.PacketChokedOut())
 			require.Equal(t, tc.expSlowServerFrames, hf.SlowServerFrames())
 			require.Equal(t, tc.expHitching, hf.Hitching())
+			require.Equal(t, tc.expDOS, hf.DOS())
+			require.Equal(t, tc.expRelay, hf.Relay())
 		})
 	}
 }