Improve trident lightning protection (TownyAdvanced#6964)

jwkerr · Oct 9, 2023 · c45eafe · c45eafe
1 parent e64feaa
commit c45eafe
Show file tree

Hide file tree

Showing 3 changed files with 52 additions and 5 deletions.
diff --git a/Towny/src/main/java/com/palmergames/bukkit/towny/listeners/TownyEntityListener.java b/Towny/src/main/java/com/palmergames/bukkit/towny/listeners/TownyEntityListener.java
@@ -28,6 +28,7 @@
 import org.bukkit.entity.DragonFireball;
 import org.bukkit.entity.Entity;
 import org.bukkit.entity.EntityType;
+import org.bukkit.entity.LightningStrike;
 import org.bukkit.entity.LivingEntity;
 import org.bukkit.entity.Mob;
 import org.bukkit.entity.Monster;
@@ -63,6 +64,7 @@
 import org.bukkit.event.hanging.HangingPlaceEvent;
 import org.bukkit.potion.PotionEffect;
 import org.bukkit.projectiles.BlockProjectileSource;
+import org.bukkit.projectiles.ProjectileSource;
 
 import java.util.ArrayList;
 import java.util.Collections;
@@ -635,10 +637,10 @@ public void onEntityCombustByEntityEvent(EntityCombustByEntityEvent event) {
 		LivingEntity attacker = null;
 		if (combuster instanceof Projectile projectile) {
 
-			Object source = projectile.getShooter();
+			final ProjectileSource source = projectile.getShooter();
 
-			if (source instanceof BlockProjectileSource) {
-				if (CombatUtil.preventDispenserDamage(((BlockProjectileSource) source).getBlock(), defender, DamageCause.PROJECTILE)) {
+			if (source instanceof BlockProjectileSource blockSource) {
+				if (CombatUtil.preventDispenserDamage(blockSource.getBlock(), defender, DamageCause.PROJECTILE)) {
 					combuster.remove();
 					event.setCancelled(true);
 					return;
@@ -657,8 +659,11 @@ public void onEntityCombustByEntityEvent(EntityCombustByEntityEvent event) {
 					event.setCancelled(true);
 				}
 			}
+		} else if (combuster instanceof LightningStrike lightning) {
+			// Protect entities from being lit on fire by player caused lightning
+			if (CombatUtil.getLightningCausingEntity(lightning) instanceof Player player && CombatUtil.preventDamageCall(player, defender, DamageCause.LIGHTNING))
+				event.setCancelled(true);
 		}
-
 	}
 
 	/**

diff --git a/Towny/src/main/java/com/palmergames/bukkit/towny/listeners/TownyWorldListener.java b/Towny/src/main/java/com/palmergames/bukkit/towny/listeners/TownyWorldListener.java
@@ -11,6 +11,7 @@
 import com.palmergames.bukkit.towny.object.Translatable;
 import com.palmergames.bukkit.towny.utils.BorderUtil;
 
+import com.palmergames.bukkit.towny.utils.CombatUtil;
 import org.bukkit.Bukkit;
 import org.bukkit.Material;
 import org.bukkit.World;
@@ -107,6 +108,13 @@ public void onLightningStrike(LightningStrikeEvent event) {
 			event.setCancelled(true);
 			return;
 		}
+
+		/*
+		 * Paper provides a method to get the entity who caused a lightning strike, if the return is non-null then we don't need to add it to the list,
+		 * since we can get the player directly in the CombatUtil.
+		 */
+		if (event.getCause().equals(LightningStrikeEvent.Cause.TRIDENT) && CombatUtil.getLightningCausingEntity(event.getLightning()) != null)
+			return;
 
 		final TownyWorld townyWorld = TownyAPI.getInstance().getTownyWorld(event.getWorld());
 		if (townyWorld == null || !townyWorld.isUsingTowny())

diff --git a/Towny/src/main/java/com/palmergames/bukkit/towny/utils/CombatUtil.java b/Towny/src/main/java/com/palmergames/bukkit/towny/utils/CombatUtil.java
@@ -34,8 +34,13 @@
 import org.bukkit.entity.Wolf;
 import org.bukkit.event.entity.EntityDamageEvent.DamageCause;
 import org.bukkit.projectiles.BlockProjectileSource;
+import org.bukkit.projectiles.ProjectileSource;
+import org.jetbrains.annotations.ApiStatus;
 import org.jetbrains.annotations.NotNull;
+import org.jetbrains.annotations.Nullable;
 
+import java.lang.invoke.MethodHandle;
+import java.lang.invoke.MethodHandles;
 import java.util.List;
 
 /**
@@ -74,14 +79,18 @@ public static boolean preventDamageCall(Entity attacker, Entity defender, Damage
 		 */
 		if (attacker instanceof Projectile projectile) {
 
-			Object source = projectile.getShooter();
+			final ProjectileSource source = projectile.getShooter();
 
 			if (source instanceof Entity entity)
 				directSource = entity;
 			else if (source instanceof BlockProjectileSource blockProjectileSource) {
 				if (CombatUtil.preventDispenserDamage(blockProjectileSource.getBlock(), defender, cause))
 					return true;
 			}
+		} else if (attacker instanceof LightningStrike lightning) {
+			final Entity causingEntity = getLightningCausingEntity(lightning);
+			if (causingEntity != null)
+				directSource = causingEntity;
 		}
 
 		if (directSource instanceof Player player)
@@ -668,4 +677,29 @@ public static boolean preventDispenserDamage(Block dispenser, Entity entity, Dam
 	private static boolean isNotNPC(Entity entity) {
 		return !PluginIntegrations.getInstance().checkCitizens(entity);
 	}
+
+	private static final @Nullable MethodHandle GET_LIGHTNING_CAUSING_ENTITY;
+
+	static {
+		MethodHandle temp = null;
+		try {
+			// https://jd.papermc.io/paper/1.20/org/bukkit/entity/LightningStrike.html#getCausingEntity()
+			//noinspection JavaReflectionMemberAccess
+			temp = MethodHandles.publicLookup().unreflect(LightningStrike.class.getMethod("getCausingEntity"));
+		} catch (Throwable ignored) {}
+
+		GET_LIGHTNING_CAUSING_ENTITY = temp;
+	}
+
+	@ApiStatus.Internal
+	public static @Nullable Entity getLightningCausingEntity(@NotNull LightningStrike lightning) {
+		if (GET_LIGHTNING_CAUSING_ENTITY == null)
+			return null;
+
+		try {
+			return (Entity) GET_LIGHTNING_CAUSING_ENTITY.invokeExact(lightning);
+		} catch (Throwable thr) {
+			return null;
+		}
+	}
 }