From 451c90b93207dd85c72f220061c41bfb9b5ba7d4 Mon Sep 17 00:00:00 2001
From: Nicholas Hastings <skamonkey@gmail.com>
Date: Fri, 21 Mar 2014 22:09:15 -0400
Subject: [PATCH] Improve logical entity support in SDKHooks natives (bug 6069,
 r=asherkin).

---
 extensions/sdkhooks/extension.cpp |  2 +-
 extensions/sdkhooks/natives.cpp   | 33 ++++++++++++++++++-------------
 extensions/sdkhooks/util.cpp      | 30 ----------------------------
 extensions/sdkhooks/util.h        |  1 -
 4 files changed, 20 insertions(+), 46 deletions(-)

diff --git a/extensions/sdkhooks/extension.cpp b/extensions/sdkhooks/extension.cpp
index b5af51ca..987d77fe 100644
--- a/extensions/sdkhooks/extension.cpp
+++ b/extensions/sdkhooks/extension.cpp
@@ -528,7 +528,7 @@ HookReturn SDKHooks::Hook(int entity, SDKHookType type, IPluginFunction *callbac
 	if(!g_HookTypes[type].supported)
 		return HookRet_NotSupported;
 
-	CBaseEntity *pEnt = UTIL_GetCBaseEntity(entity);
+	CBaseEntity *pEnt = gamehelpers->ReferenceToEntity(entity);
 	if(!pEnt)
 		return HookRet_InvalidEntity;
 	if(type < 0 || type >= SDKHook_MAXHOOKS)
diff --git a/extensions/sdkhooks/natives.cpp b/extensions/sdkhooks/natives.cpp
index 76018e08..b1555557 100644
--- a/extensions/sdkhooks/natives.cpp
+++ b/extensions/sdkhooks/natives.cpp
@@ -58,14 +58,15 @@ cell_t Native_Hook(IPluginContext *pContext, const cell_t *params)
 			break;
 		case HookRet_BadEntForHookType:
 		{
-			const char * pClassname = gamehelpers->GetEntityClassname(PEntityOfEntIndex(gamehelpers->ReferenceToIndex(params[1])));
-			if (!pClassname)
-			{
-				pContext->ThrowNativeError("Hook type not valid for this type of entity (%i).", entity);
-			}
-			else
-			{
-				pContext->ThrowNativeError("Hook type not valid for this type of entity (%s)", pClassname);
+			CBaseEntity *pEnt = gamehelpers->ReferenceToEntity(params[1]);
+			const char *pClassname = pEnt ? gamehelpers->GetEntityClassname(pEnt) : NULL;
+			if (!pClassname)
+			{
+				pContext->ThrowNativeError("Hook type not valid for this type of entity (%i).", entity);
+			}
+			else
+			{
+				pContext->ThrowNativeError("Hook type not valid for this type of entity (%i/%s)", entity, pClassname);
 			}
 			
 			break;
@@ -108,18 +109,18 @@ cell_t Native_TakeDamage(IPluginContext *pContext, const cell_t *params)
 #if !defined SH_DECL_MANUALEXTERN1
 	pContext->ThrowNativeError("SDKHooks_TakeDamage is not supported on this engine.");
 #else
-	CBaseEntity *pVictim = UTIL_GetCBaseEntity(params[1]);
+	CBaseEntity *pVictim = gamehelpers->ReferenceToEntity(params[1]);
 	if (!pVictim)
 		return pContext->ThrowNativeError("Invalid entity index %d for victim", params[1]);
 	
-	CBaseEntity *pInflictor = UTIL_GetCBaseEntity(params[2]);
+	CBaseEntity *pInflictor = gamehelpers->ReferenceToEntity(params[2]);
 	if (!pInflictor)
 		return pContext->ThrowNativeError("Invalid entity index %d for inflictor", params[2]);
 
 	CBaseEntity *pAttacker;
 	if (params[3] != -1)
 	{
-		pAttacker = UTIL_GetCBaseEntity(params[3]);
+		pAttacker = gamehelpers->ReferenceToEntity(params[3]);
 		if (!pAttacker)
 		{
 			return pContext->ThrowNativeError("Invalid entity index %d for attackerr", params[3]);
@@ -136,7 +137,7 @@ cell_t Native_TakeDamage(IPluginContext *pContext, const cell_t *params)
 	CBaseEntity *pWeapon;
 	if (params[6] != -1)
 	{
-		pWeapon = UTIL_GetCBaseEntity(params[6]);
+		pWeapon = gamehelpers->ReferenceToEntity(params[6]);
 		if (!pWeapon)
 		{
 			return pContext->ThrowNativeError("Invalid entity index %d for weapon", params[6]);
@@ -192,11 +193,15 @@ cell_t Native_DropWeapon(IPluginContext *pContext, const cell_t *params)
 #if !defined SH_DECL_MANUALEXTERN1
 	pContext->ThrowNativeError("SDKHooks_DropWeapon is not supported on this engine.");
 #else
-	CBaseEntity *pPlayer = UTIL_GetCBaseEntity(params[1], true);
+	CBaseEntity *pPlayer = gamehelpers->ReferenceToEntity(params[1]);
 	if (!pPlayer)
 		return pContext->ThrowNativeError("Invalid client index %d", params[1]);
+
+	IGamePlayer *pGamePlayer = playerhelpers->GetGamePlayer(gamehelpers->ReferenceToIndex(params[1]));
+	if (!pGamePlayer || !pGamePlayer->IsInGame())
+		return pContext->ThrowNativeError("Client index %d not in game", params[1]);
 	
-	CBaseEntity *pWeapon = UTIL_GetCBaseEntity(params[2]);
+	CBaseEntity *pWeapon = gamehelpers->ReferenceToEntity(params[2]);
 	if (!pWeapon)
 		return pContext->ThrowNativeError("Invalid entity index %d for weapon", params[2]);
 
diff --git a/extensions/sdkhooks/util.cpp b/extensions/sdkhooks/util.cpp
index 8370943d..03abd1de 100644
--- a/extensions/sdkhooks/util.cpp
+++ b/extensions/sdkhooks/util.cpp
@@ -34,36 +34,6 @@
 #include "extension.h"
 #include "util.h"
 
-CBaseEntity *UTIL_GetCBaseEntity(int num, bool onlyPlayers)
-{
-	edict_t *pEdict = PEntityOfEntIndex(num);
-	if (!pEdict || pEdict->IsFree())
-	{
-		return NULL;
-	}
-
-	if (num > 0 && num <= playerhelpers->GetMaxClients())
-	{
-		IGamePlayer *pPlayer = playerhelpers->GetGamePlayer(pEdict);
-		if (!pPlayer || !pPlayer->IsConnected())
-		{
-			return NULL;
-		}
-	}
-	else if (onlyPlayers)
-	{
-		return NULL;
-	}
-
-	IServerUnknown *pUnk;
-	if ((pUnk=pEdict->GetUnknown()) == NULL)
-	{
-		return NULL;
-	}
-
-	return pUnk->GetBaseEntity();
-}
-
 bool UTIL_ContainsDataTable(SendTable *pTable, const char *name)
 {
 	const char *pname = pTable->GetName();
diff --git a/extensions/sdkhooks/util.h b/extensions/sdkhooks/util.h
index 3a4b46f5..9d1a919d 100644
--- a/extensions/sdkhooks/util.h
+++ b/extensions/sdkhooks/util.h
@@ -34,7 +34,6 @@
 
 #include <compat_wrappers.h>
 
-CBaseEntity *UTIL_GetCBaseEntity(int num, bool onlyPlayers=false);
 bool UTIL_ContainsDataTable(SendTable *pTable, const char *name);
 
 #endif  //_INCLUDE_TF2TOOLS_UTIL_H_