WT-8392: crash when using WToolBar::widget()

It is possible to add any widget using WToolBar::addWidget(), but count()
and widget() didn't expect it. Fixed by tracking all added widgets
instead of relying on dynamic_cast

Author: korneel-emweb

src/Wt/WToolBar.C

@@ -10,6 +10,8 @@
 #include "Wt/WSplitButton.h"
 #include "Wt/WToolBar.h"
 
+#include <algorithm>
+
 namespace Wt {
 
 LOGGER("WToolBar");
@@ -34,6 +36,8 @@ void WToolBar::setOrientation(Orientation orientation)
 void WToolBar::addButton(std::unique_ptr<WPushButton> button,
 			 AlignmentFlag alignmentFlag)
 {
+  widgets_.push_back(button.get());
+
   if (compact_){
     if (alignmentFlag == AlignmentFlag::Right)
       button->setAttributeValue("style", "float:right;");
@@ -48,6 +52,8 @@ void WToolBar::addButton(std::unique_ptr<WPushButton> button,
 void WToolBar::addButton(std::unique_ptr<WSplitButton> button,
 			 AlignmentFlag alignmentFlag)
 {
+  widgets_.push_back(button.get());
+
   setCompact(false);
   lastGroup_ = nullptr;
   if (alignmentFlag == AlignmentFlag::Right)
@@ -58,6 +64,8 @@ void WToolBar::addButton(std::unique_ptr<WSplitButton> button,
 void WToolBar::addWidget(std::unique_ptr<WWidget> widget,
 			 AlignmentFlag alignmentFlag)
 {
+  widgets_.push_back(widget.get());
+
   setCompact(false);
   lastGroup_ = nullptr;
   if (alignmentFlag == AlignmentFlag::Right)
@@ -67,23 +75,19 @@ void WToolBar::addWidget(std::unique_ptr<WWidget> widget,
 
 std::unique_ptr<WWidget> WToolBar::removeWidget(WWidget *widget)
 {
-  WWidget *p = widget->parent();
-  if (p == impl_)
-    return impl_->removeWidget(widget);
-  else {
-    int i = impl_->indexOf(p);
-    if (i >= 0) {
-      WContainerWidget *cw = dynamic_cast<WContainerWidget *>(p);
-      if (cw) {
-	std::unique_ptr<WWidget> result = cw->removeWidget(widget);
-	if (cw->count() == 0)
-	  cw->parent()->removeWidget(cw);
-	return result;
-      }
-    }
-  }
-
-  return std::unique_ptr<WWidget>();
+  auto it = std::find(widgets_.begin(), widgets_.end(), widget);
+  if (it != widgets_.end()) {
+    auto widget = *it;
+    auto parent = static_cast<WContainerWidget*>(widget->parent());
+
+    auto retval = parent->removeWidget(widget);
+    if (parent != impl_  && parent->count() == 0)
+      impl_->removeWidget(parent);
+    widgets_.erase(it);
+
+    return retval;
+  } else
+    return nullptr;
 }
 
 void WToolBar::addSeparator()
@@ -129,51 +133,15 @@ WContainerWidget *WToolBar::lastGroup()
 
 int WToolBar::count() const
 {
-  if (compact_)
-    return impl_->count();
-  else {
-    int result = 0;
-    for (int i = 0; i < impl_->count(); ++i) {
-      WWidget *w = impl_->widget(i);
-
-      if (dynamic_cast<WSplitButton *>(w))
-	++result;
-      else {
-	WContainerWidget *group = dynamic_cast<WContainerWidget *>(w);
-
-	result += group->count();
-      }
-    }
-
-    return result;
-  }
+  return widgets_.size();
 }
 
 WWidget *WToolBar::widget(int index) const
 {
-  if (compact_)
-    return impl_->widget(index);
-  else {
-    int current = 0;
-    for (int i = 0; i < impl_->count(); ++i) {
-      WWidget *w = impl_->widget(i);
-
-      if (dynamic_cast<WSplitButton *>(w)) {
-	if (index == current)
-	  return w;
-	++current;
-      } else {
-	WContainerWidget *group = dynamic_cast<WContainerWidget *>(w);
-
-	if (index < current + group->count())
-	  return group->widget(index - current);
-
-	current += group->count();
-      }
-    }
-
+  if (index < widgets_.size())
+    return widgets_[index];
+  else
     return nullptr;
-  }
 }
 
 }

src/Wt/WToolBar.h

@@ -67,7 +67,10 @@ class WT_API WToolBar : public WCompositeWidget
    */
   void addSeparator();
 
-  /*! \brief Returns the number of buttons.
+  /*! \brief Returns the number of widgets.
+   *
+   * The counted widgets are either a WPushButton or WSplitButton added by
+   * addButton() or a widget added by addWidget().
    *
    * \sa widget()
    */
@@ -99,6 +102,8 @@ class WT_API WToolBar : public WCompositeWidget
   WContainerWidget *impl_;
   WContainerWidget *lastGroup_;
 
+  std::vector<WWidget *> widgets_;
+
   WContainerWidget *lastGroup();
 };