Partial fix

--HG--
extra : convert_revision : svn%3Ac2935e3e-5518-0410-8daf-afa5dab7d4e3/trunk%40122

sourcehook/sourcehook.cpp

@@ -341,12 +341,26 @@ namespace SourceHook
 
 			bool erase;
 			for (List<HookInfo>::iterator hookiter = hooks.begin();
-				hookiter != hooks.end(); erase ? hookiter = hooks.erase(hookiter) : ++hookiter)
+				hookiter != hooks.end(); )
 			{
 				erase = hookiter->plug == plug && hookiter->handler->IsEqual(handler) &&
 					hookiter->thisptr_offs == thisptr_offs;
 				if (erase)
+				{
 					hookiter->handler->DeleteThis();			// Make the _plugin_ delete the handler object
+
+					// Move all iterators pointing at this
+					List<HookInfo>::iterator oldhookiter = hookiter;
+					hookiter = hooks.erase(hookiter);
+					List<HookInfo>::iterator newhookiter = hookiter;
+					--newhookiter; // The hook loop will ++ it then
+					CHookList::CIter *pItIter;
+					for (pItIter = iface_iter->m_PreHooks.m_UsedIters; pItIter; pItIter = pItIter->m_pNext)
+						if (pItIter->m_Iter == oldhookiter)
+							pItIter->m_Iter = newhookiter;
+				}
+				else
+					++hookiter;
 			}
 			if (iface_iter->m_PostHooks.m_List.empty() && iface_iter->m_PreHooks.m_List.empty())
 			{
@@ -701,10 +715,10 @@ namespace SourceHook
 	// CHookList
 	////////////////////////////
 
-	CSourceHookImpl::CHookList::CHookList() : m_FreeIters(NULL)
+	CSourceHookImpl::CHookList::CHookList() : m_FreeIters(NULL), m_UsedIters(NULL)
 	{
 	}
-	CSourceHookImpl::CHookList::CHookList(const CHookList &other) : m_List(other.m_List), m_FreeIters(NULL)
+	CSourceHookImpl::CHookList::CHookList(const CHookList &other) : m_List(other.m_List), m_FreeIters(NULL), m_UsedIters(NULL)
 	{
 	}
 	CSourceHookImpl::CHookList::~CHookList()
@@ -718,19 +732,41 @@ namespace SourceHook
 	}
 	IHookList::IIter *CSourceHookImpl::CHookList::GetIter()
 	{
+		CIter *ret;
 		if (m_FreeIters)
 		{
-			CIter *ret = m_FreeIters;
+			ret = m_FreeIters;
 			m_FreeIters = ret->m_pNext;
 			ret->GoToBegin();
-			return ret;
 		}
-		return new CIter(this);
+		else
+		{
+			ret = new CIter(this);
+		}
+		
+		ret->m_pNext = m_UsedIters;
+		ret->m_pPrev = NULL;
+		if (m_UsedIters)
+			m_UsedIters->m_pPrev = ret;
+		m_UsedIters = ret;
+
+		return ret;
 	}
 	void CSourceHookImpl::CHookList::ReleaseIter(IIter *pIter)
 	{
 		CIter *pIter2 = static_cast<CIter*>(pIter);
+
+		// Unlink from m_UsedIters
+
+		if (pIter2->m_pNext)
+			pIter2->m_pNext->m_pPrev = pIter2->m_pPrev;
+		if (pIter2->m_pPrev)
+			pIter2->m_pPrev->m_pNext = pIter2->m_pNext;
+
+		// Link to m_FreeIters
+
 		pIter2->m_pNext = m_FreeIters;
+
 		m_FreeIters = pIter2;
 	}
 
@@ -750,13 +786,21 @@ namespace SourceHook
 
 	bool CSourceHookImpl::CHookList::CIter::End()
 	{
+		if (!m_pList)
+			return false;
 		return m_Iter == m_pList->m_List.end();
 	}
 	void CSourceHookImpl::CHookList::CIter::Next()
 	{
+		if (!m_pList)
+			return;
 		++m_Iter;
 		SkipPaused();
 	}
+	void CSourceHookImpl::CHookList::CIter::Clear()
+	{
+		m_pList = NULL;
+	}
 	void CSourceHookImpl::CHookList::CIter::SkipPaused()
 	{
 		while (m_Iter != m_pList->m_List.end() && m_Iter->paused)

sourcehook/sourcehook_impl.h

@@ -68,10 +68,12 @@ namespace SourceHook
 				friend class CHookList;
 
 				CHookList *m_pList;
-				List<HookInfo>::iterator m_Iter;
 
 				void SkipPaused();
 			public:
+
+				List<HookInfo>::iterator m_Iter;
+
 				CIter(CHookList *pList);
 
 				virtual ~CIter();
@@ -83,10 +85,14 @@ namespace SourceHook
 				ISHDelegate *Handler();
 				int ThisPtrOffs();
 
-				CIter *m_pNext;		// When stored in m_FreeIters
+				void Clear();
+
+				CIter *m_pNext;		// When stored in m_FreeIters and m_UsedIters
+				CIter *m_pPrev;		// Only used when stored in m_UsedIters
 			};
 
 			CIter *m_FreeIters;
+			CIter *m_UsedIters;
 
 			CHookList();
 			CHookList(const CHookList &other);

sourcehook/test/testbail.cpp

@@ -5,6 +5,28 @@
 
 void *___testbail_gabgab;
 
+namespace
+{
+	class zomg_lolz
+	{
+	public:
+		virtual void zomg()
+		{
+		}
+	};
+	SH_DECL_HOOK0_void(zomg_lolz, zomg, SH_NOATTRIB, 0);
+	void Handler()
+	{
+		printf("H1\n");
+		SH_REMOVE_HOOK_STATICFUNC(zomg_lolz, zomg, reinterpret_cast<zomg_lolz*>(META_IFACEPTR),
+			Handler, false);
+	}
+	void Handler2()
+	{
+		printf("H2\n");
+	}
+}
+
 bool TestBail(std::string &error)
 {
 	SourceHook::CSourceHookImpl g_SHImpl;
@@ -32,6 +54,23 @@ bool TestBail(std::string &error)
 
 	// If it didn't crash, it's ok
 
+	// NEW TEST: Remove hook from handler
+	zomg_lolz inst;
+	SH_ADD_HOOK_STATICFUNC(zomg_lolz, zomg, &inst, Handler, false);
+	SH_ADD_HOOK_STATICFUNC(zomg_lolz, zomg, &inst, Handler2, false);
+
+	zomg_lolz *mwah = &inst;
+	mwah->zomg();
+	mwah->zomg();
+
+	SH_ADD_HOOK_STATICFUNC(zomg_lolz, zomg, &inst, Handler, false);
+	SH_REMOVE_HOOK_STATICFUNC(zomg_lolz, zomg, &inst, Handler2, false);
+
+	mwah->zomg();
+	mwah->zomg();
+
+	// Shouldn't crash again...
+
 	return true;
 }