From 9f10aa20168160ed76dc76cde08267fb7d3b39be Mon Sep 17 00:00:00 2001
From: JonBooth78 <jon.booth.uk@gmail.com>
Date: Tue, 17 Oct 2023 12:29:22 +1300
Subject: [PATCH] Fix the undefined behaviour when addiing bodies during the
 update loops that iterate through them

The undefined behaviour comes because when we add an item to the vector, it could have to reallocate to grow the backing storage and therefore the end iterator becomes potentially invalid.

Defer adding bodies until the end of an Space update tick.

A few game related scenarios (emerging from hyperspace, saving a game and loading a game) require this update to be invoked too.
---
 src/Game.cpp  |  6 ++++++
 src/Space.cpp | 39 +++++++++++++++++++++++----------------
 src/Space.h   | 12 +++++++++---
 3 files changed, 38 insertions(+), 19 deletions(-)

diff --git a/src/Game.cpp b/src/Game.cpp
index e32e5b9401a..b70abf1098f 100644
--- a/src/Game.cpp
+++ b/src/Game.cpp
@@ -68,6 +68,8 @@ Game::Game(const SystemPath &path, const double startDateTime, const char *shipT
 
 	m_space.reset(new Space(this, m_galaxy, path));
 
+	m_space->UpdateBodies();
+
 	Body *b = m_space->FindBodyForPath(&path);
 	assert(b);
 
@@ -525,6 +527,8 @@ void Game::SwitchToNormalSpace()
 	m_player->SetFrame(m_space->GetRootFrame());
 	m_space->AddBody(m_player.get());
 
+	m_space->UpdateBodies();
+
 	// place it
 	vector3d pos, vel;
 	m_space->GetHyperspaceExitParams(m_hyperspaceSource, m_hyperspaceDest, pos, vel);
@@ -937,6 +941,8 @@ void Game::SaveGame(const std::string &filename, Game *game)
 		throw CouldNotOpenFileException();
 	}
 
+	game->m_space->UpdateBodies();
+
 	Json rootNode;
 	game->ToJson(rootNode); // Encode the game data as JSON and give to the root value.
 	std::vector<uint8_t> jsonData;
diff --git a/src/Space.cpp b/src/Space.cpp
index 5be29a6c89c..1709973a107 100644
--- a/src/Space.cpp
+++ b/src/Space.cpp
@@ -438,7 +438,10 @@ void Space::RebuildSystemBodyIndex()
 
 void Space::AddBody(Body *b)
 {
-	m_bodies.push_back(b);
+#ifndef NDEBUG
+	assert(!m_processingFinalizationQueue);
+#endif
+	m_assignedBodies.emplace_back(b, BodyAssignation::CREATE);
 }
 
 void Space::RemoveBody(Body *b)
@@ -1073,22 +1076,26 @@ void Space::UpdateBodies()
 	m_processingFinalizationQueue = true;
 #endif
 
-	// removing or deleting bodies from space
+	// adding, removing or deleting bodies from space
 	for (const auto &b : m_assignedBodies) {
-		auto remove_iterator = m_bodies.end();
-		for (auto it = m_bodies.begin(); it != m_bodies.end(); ++it) {
-			if (*it != b.first)
-				(*it)->NotifyRemoved(b.first);
-			else
-				remove_iterator = it;
-		}
-		if (remove_iterator != m_bodies.end()) {
-			*remove_iterator = m_bodies.back();
-			m_bodies.pop_back();
-			if (b.second == BodyAssignation::KILL)
-				delete b.first;
-			else
-				b.first->SetFrame(FrameId::Invalid);
+		if (b.second == BodyAssignation::CREATE) {
+			m_bodies.push_back(b.first);
+		} else {
+			auto remove_iterator = m_bodies.end();
+			for (auto it = m_bodies.begin(); it != m_bodies.end(); ++it) {
+				if (*it != b.first)
+					(*it)->NotifyRemoved(b.first);
+				else
+					remove_iterator = it;
+			}
+			if (remove_iterator != m_bodies.end()) {
+				*remove_iterator = m_bodies.back();
+				m_bodies.pop_back();
+				if (b.second == BodyAssignation::KILL)
+					delete b.first;
+				else
+					b.first->SetFrame(FrameId::Invalid);
+			}
 		}
 	}
 
diff --git a/src/Space.h b/src/Space.h
index b92007dd21e..dd53d6e9f7c 100644
--- a/src/Space.h
+++ b/src/Space.h
@@ -105,6 +105,12 @@ class Space {
 	//it is not designed to be called in each game loop!
 	std::vector<BodyDist> BodiesInAngle(const Body *b, const vector3d &offset, const vector3d &dir, double cosOfMaxAngle) const;
 
+	// Don't call this unless you really have to
+	// A bit of a hack but required so we can add a player
+	// into space as we emerge from hyperspace
+	// and then find stuff.
+	void UpdateBodies();
+
 private:
 	void GenSectorCache(RefCountedPtr<Galaxy> galaxy, const SystemPath *here);
 	void UpdateStarSystemCache(const SystemPath *here);
@@ -112,7 +118,6 @@ class Space {
 	// make sure SystemBody* is in Pi::currentSystem
 	FrameId GetFrameWithSystemBody(const SystemBody *b) const;
 
-	void UpdateBodies();
 
 	void CollideFrame(FrameId fId);
 
@@ -128,10 +133,11 @@ class Space {
 	// all the bodies we know about
 	std::vector<Body *> m_bodies;
 
-	// bodies that were removed/killed this timestep and need pruning at the end
+	// bodies that were removed/killed/added this timestep and need pruning/adding at the end
 	enum class BodyAssignation {
 		KILL = 0,
-		REMOVE = 1
+		REMOVE = 1,
+		CREATE = 2
 	};
 
 	std::vector<std::pair<Body *, BodyAssignation>> m_assignedBodies;