fix: hostage leader possibly nil

akiver/cs-demo-manager#839
markus-wa · May 16, 2024 · 31b54b4 · 31b54b4
1 parent a58c851
commit 31b54b4
Show file tree

Hide file tree

Showing 2 changed files with 35 additions and 1 deletion.
diff --git a/pkg/demoinfocs/common/hostage.go b/pkg/demoinfocs/common/hostage.go
@@ -3,6 +3,7 @@ package common
 import (
 	"github.com/golang/geo/r3"
 
+	"github.com/markus-wa/demoinfocs-golang/v4/pkg/demoinfocs/constants"
 	st "github.com/markus-wa/demoinfocs-golang/v4/pkg/demoinfocs/sendtables"
 )
 
@@ -54,7 +55,12 @@ func (hostage *Hostage) Health() int {
 // Returns nil if the hostage is not following a player.
 func (hostage *Hostage) Leader() *Player {
 	if hostage.demoInfoProvider.IsSource2() {
-		return hostage.demoInfoProvider.FindPlayerByPawnHandle(getUInt64(hostage.Entity, "m_leader"))
+		leaderHandle := getUInt64(hostage.Entity, "m_leader")
+		if leaderHandle != constants.InvalidEntityHandleSource2 {
+			return hostage.demoInfoProvider.FindPlayerByPawnHandle(leaderHandle)
+		}
+
+		return hostage.demoInfoProvider.FindPlayerByPawnHandle(getUInt64(hostage.Entity, "m_hHostageGrabber"))
 	}
 
 	return hostage.demoInfoProvider.FindPlayerByHandle(uint64(getInt(hostage.Entity, "m_leader")))

diff --git a/pkg/demoinfocs/common/hostage_test.go b/pkg/demoinfocs/common/hostage_test.go
@@ -5,6 +5,7 @@ import (
 
 	"github.com/stretchr/testify/assert"
 
+	"github.com/markus-wa/demoinfocs-golang/v4/pkg/demoinfocs/constants"
 	st "github.com/markus-wa/demoinfocs-golang/v4/pkg/demoinfocs/sendtables"
 )
 
@@ -19,6 +20,29 @@ func TestHostage_Leader(t *testing.T) {
 	assert.Equal(t, player, hostage.Leader())
 }
 
+func TestHostage_LeaderWithInvalidHandleS2(t *testing.T) {
+	player := new(Player)
+	player.EntityID = 10
+	provider := demoInfoProviderMock{
+		playersByHandle: map[uint64]*Player{10: player},
+		isSource2:       true,
+	}
+	hostage := hostageWithProperties([]fakeProp{
+		{
+			propName: "m_leader",
+			value:    st.PropertyValue{Any: uint64(constants.InvalidEntityHandleSource2)},
+			isNil:    false,
+		},
+		{
+			propName: "m_hHostageGrabber",
+			value:    st.PropertyValue{Any: uint64(10)},
+			isNil:    false,
+		},
+	}, provider)
+
+	assert.Equal(t, player, hostage.Leader())
+}
+
 func TestHostage_State(t *testing.T) {
 	hostage := hostageWithProperty("m_nHostageState", st.PropertyValue{IntVal: int(HostageStateFollowingPlayer)}, demoInfoProviderMock{})
 
@@ -34,3 +58,7 @@ func TestHostage_Health(t *testing.T) {
 func hostageWithProperty(propName string, value st.PropertyValue, provider demoInfoProviderMock) *Hostage {
 	return &Hostage{Entity: entityWithProperty(propName, value), demoInfoProvider: provider}
 }
+
+func hostageWithProperties(properties []fakeProp, provider demoInfoProviderMock) *Hostage {
+	return &Hostage{Entity: entityWithProperties(properties), demoInfoProvider: provider}
+}