Fix memory errors caused by invalid pointer access in Release() and broken ComPtr smart pointer implementation.

Author: PCMan

ComPtr.h

@@ -33,8 +33,8 @@ class ComPtr {
 	ComPtr(void): p_(NULL) {
 	}
 
-	ComPtr(T* p): p_(p) {
-		if(p_)
+	ComPtr(T* p, bool ref = true): p_(p) {
+		if(p_ && ref)
 			p_->AddRef();
 	}
 
@@ -81,12 +81,17 @@ class ComPtr {
 		return p_ < p;
 	}
 
+	ComPtr& operator = (const ComPtr& other) {
+		return operator = (other.p_);
+	}
+
 	ComPtr& operator = (T* p) {
-		if(p_)
-			p_->Release();
+		T* old = p_;
 		p_ = p;
 		if(p_)
 			p_->AddRef();
+		if (old)
+			old->Release();
 		return *this;
 	}
 

DisplayAttributeInfo.cpp

@@ -62,10 +62,10 @@ STDMETHODIMP_(ULONG) DisplayAttributeInfo::AddRef(void) {
 
 STDMETHODIMP_(ULONG) DisplayAttributeInfo::Release(void) {
 	assert(refCount_ > 0);
-	--refCount_;
+	const ULONG newCount = --refCount_;
 	if(0 == refCount_)
 		delete this;
-	return refCount_;
+	return newCount;
 }
 
 // ITfDisplayAttributeInfo

DisplayAttributeInfoEnum.cpp

@@ -62,10 +62,10 @@ STDMETHODIMP_(ULONG) DisplayAttributeInfoEnum::AddRef(void) {
 
 STDMETHODIMP_(ULONG) DisplayAttributeInfoEnum::Release(void) {
 	assert(refCount_ > 0);
-	--refCount_;
-	if(0 == refCount_)
+	const ULONG newCount = --refCount_;
+	if (0 == refCount_)
 		delete this;
-	return refCount_;
+	return newCount;
 }
 
 

DisplayAttributeProvider.cpp

@@ -61,10 +61,10 @@ STDMETHODIMP_(ULONG) DisplayAttributeProvider::AddRef(void) {
 
 STDMETHODIMP_(ULONG) DisplayAttributeProvider::Release(void) {
 	assert(refCount_ > 0);
-	--refCount_;
-	if(0 == refCount_)
+	const ULONG newCount = --refCount_;
+	if (0 == refCount_)
 		delete this;
-	return refCount_;
+	return newCount;
 }
 
 

EditSession.cpp

@@ -67,10 +67,10 @@ STDMETHODIMP_(ULONG) EditSession::AddRef(void) {
 
 STDMETHODIMP_(ULONG) EditSession::Release(void) {
 	assert(refCount_ > 0);
-	--refCount_;
-	if(0 == refCount_)
+	const ULONG newCount = --refCount_;
+	if (0 == refCount_)
 		delete this;
-	return refCount_;
+	return newCount;
 }
 
 STDMETHODIMP EditSession::DoEditSession(TfEditCookie ec) {

ImeModule.cpp

@@ -505,10 +505,11 @@ STDMETHODIMP_(ULONG) ImeModule::Release(void) {
 	// The last reference is released in DllMain() when unloading.
 	// Hence interlocked operations are enough here, I guess.
 	assert(refCount_ > 0);
+	const ULONG newCount = --refCount_;
 	if(::InterlockedExchangeSubtract(&refCount_, 1) == 1) {
 		delete this;
 	}
-	return refCount_;
+	return newCount;
 }
 
 // IClassFactory

LangBarButton.cpp

@@ -212,10 +212,10 @@ STDMETHODIMP_(ULONG) LangBarButton::AddRef(void) {
 
 STDMETHODIMP_(ULONG) LangBarButton::Release(void) {
 	assert(refCount_ > 0);
-	--refCount_;
-	if(0 == refCount_)
+	const ULONG newCount = --refCount_;
+	if (0 == refCount_)
 		delete this;
-	return refCount_;
+	return newCount;
 }
 
 // ITfLangBarItem

TextService.cpp

@@ -68,12 +68,6 @@ TextService::~TextService(void) {
 		}
 	}
 
-	if(!langBarButtons_.empty()) {
-		for(vector<LangBarButton*>::iterator it = langBarButtons_.begin(); it != langBarButtons_.end(); ++it) {
-			LangBarButton* button = *it;
-			button->Release();
-		}
-	}
 	if(langBarMgr_) {
 		langBarMgr_->UnadviseEventSink(langBarSinkCookie_);
 	}
@@ -108,31 +102,26 @@ DWORD TextService::langBarStatus() const {
 
 void TextService::addButton(LangBarButton* button) {
 	if(button) {
-		langBarButtons_.push_back(button);
-		button->AddRef();
+		langBarButtons_.emplace_back(button);
 		if(isActivated()) {
-			ITfLangBarItemMgr* langBarItemMgr;
+			ComPtr<ITfLangBarItemMgr> langBarItemMgr;
 			if(threadMgr_->QueryInterface(IID_ITfLangBarItemMgr, (void**)&langBarItemMgr) == S_OK) {
 				langBarItemMgr->AddItem(button);
-				langBarItemMgr->Release();
 			}
 		}
 	}
 }
 
 void TextService::removeButton(LangBarButton* button) {
 	if(button) {
-		vector<LangBarButton*>::iterator it;
-		it = find(langBarButtons_.begin(), langBarButtons_.end(), button);
+		auto it = find(langBarButtons_.begin(), langBarButtons_.end(), button);
 		if(it != langBarButtons_.end()) {
 			if(isActivated()) {
-				ITfLangBarItemMgr* langBarItemMgr;
+				ComPtr<ITfLangBarItemMgr> langBarItemMgr;
 				if(threadMgr_->QueryInterface(IID_ITfLangBarItemMgr, (void**)&langBarItemMgr) == S_OK) {
 					langBarItemMgr->RemoveItem(button);
-					langBarItemMgr->Release();
 				}
 			}
-			button->Release();
 			langBarButtons_.erase(it);
 		}
 	}
@@ -621,13 +610,13 @@ STDMETHODIMP_(ULONG) TextService::AddRef(void) {
 
 STDMETHODIMP_(ULONG) TextService::Release(void) {
 	assert(refCount_ > 0);
-	--refCount_;
+	const ULONG newCount = --refCount_;
 	if(0 == refCount_) {
 		// ImeModule needs to do some clean up before deleting the TextService object.
 		module_->removeTextService(this);
 		delete this;
 	}
-	return refCount_;
+	return newCount;
 }
 
 // ITfTextInputProcessor
@@ -700,8 +689,7 @@ STDMETHODIMP TextService::Activate(ITfThreadMgr *pThreadMgr, TfClientId tfClient
 	if(!langBarButtons_.empty()) {
 		ComPtr<ITfLangBarItemMgr> langBarItemMgr;
 		if(threadMgr_->QueryInterface(IID_ITfLangBarItemMgr, (void**)&langBarItemMgr) == S_OK) {
-			for(vector<LangBarButton*>::iterator it = langBarButtons_.begin(); it != langBarButtons_.end(); ++it) {
-				LangBarButton* button = *it;
+			for(auto& button: langBarButtons_) {
 				langBarItemMgr->AddItem(button);
 			}
 		}
@@ -729,8 +717,7 @@ STDMETHODIMP TextService::Deactivate() {
 	if(!langBarButtons_.empty()) {
 		ComPtr<ITfLangBarItemMgr> langBarItemMgr;
 		if(threadMgr_->QueryInterface(IID_ITfLangBarItemMgr, (void**)&langBarItemMgr) == S_OK) {
-			for(vector<LangBarButton*>::iterator it = langBarButtons_.begin(); it != langBarButtons_.end(); ++it) {
-				LangBarButton* button = *it;
+			for(auto& button: langBarButtons_) {
 				langBarItemMgr->RemoveItem(button);
 			}
 		}

TextService.h

@@ -305,7 +305,7 @@ class TextService:
 
 	ITfComposition* composition_; // acquired when starting composition, released when ending composition
 	ComPtr<ITfLangBarMgr> langBarMgr_;
-	std::vector<LangBarButton*> langBarButtons_;
+	std::vector<ComPtr<LangBarButton>> langBarButtons_;
 	std::vector<PreservedKey> preservedKeys_;
 	std::vector<CompartmentMonitor> compartmentMonitors_;