Added feature to link a type to an address (#1219)

* Added feature to link a type to an address

src/Cutter.pro

@@ -302,7 +302,8 @@ SOURCES += \
     widgets/SdbWidget.cpp \
     common/PythonManager.cpp \
     plugins/PluginManager.cpp \
-    common/BasicBlockHighlighter.cpp
+    common/BasicBlockHighlighter.cpp \
+    dialogs/LinkTypeDialog.cpp
 
 HEADERS  += \
     core/Cutter.h \
@@ -418,7 +419,8 @@ HEADERS  += \
     widgets/SdbWidget.h \
     common/PythonManager.h \
     plugins/PluginManager.h \
-    common/BasicBlockHighlighter.h
+    common/BasicBlockHighlighter.h \
+    dialogs/LinkTypeDialog.h
 
 FORMS    += \
     dialogs/AboutDialog.ui \
@@ -479,7 +481,8 @@ FORMS    += \
     dialogs/WelcomeDialog.ui \
     dialogs/EditMethodDialog.ui \
     dialogs/LoadNewTypesDialog.ui \
-    widgets/SdbWidget.ui
+    widgets/SdbWidget.ui \
+    dialogs/LinkTypeDialog.ui
 
 RESOURCES += \
     resources.qrc \

src/core/Cutter.cpp

@@ -2315,6 +2315,12 @@ QString CutterCore::addTypes(const char *str)
     return QString();
 }
 
+bool CutterCore::isAddressMapped(RVA addr)
+{
+    // If value returned by "om. @ addr" is empty means that address is not mapped
+    return !Core()->cmd(QString("om. @ %1").arg(addr)).isEmpty();
+}
+
 QList<SearchDescription> CutterCore::getAllSearch(QString search_for, QString space)
 {
     CORE_LOCK();

src/core/Cutter.h

@@ -333,6 +333,13 @@ public:
     QString addTypes(const char *str);
     QString addTypes(const QString &str) { return addTypes(str.toUtf8().constData()); }
 
+    /*!
+     * \brief Checks if the given address is mapped to a region
+     * \param addr The address to be checked
+     * \return true if addr is mapped, false otherwise
+     */
+    bool isAddressMapped(RVA addr);
+
     QList<MemoryMapDescription> getMemoryMap();
     QList<SearchDescription> getAllSearch(QString search_for, QString space);
     BlockStatistics getBlockStatistics(unsigned int blocksCount);

src/dialogs/LinkTypeDialog.cpp

@@ -0,0 +1,106 @@
+#include "LinkTypeDialog.h"
+#include "ui_LinkTypeDialog.h"
+
+LinkTypeDialog::LinkTypeDialog(QWidget *parent) :
+    QDialog(parent),
+    ui(new Ui::LinkTypeDialog)
+{
+    ui->setupUi(this);
+
+    setWindowTitle(tr("Link type to address"));
+
+    // Populate the structureTypeComboBox
+    ui->structureTypeComboBox->addItem(tr("(No Type)"));
+    for (const TypeDescription &thisType : Core()->getAllStructs()) {
+        ui->structureTypeComboBox->addItem(thisType.type);
+    }
+}
+
+LinkTypeDialog::~LinkTypeDialog()
+{
+    delete ui;
+}
+
+void LinkTypeDialog::setDefaultType(const QString &type)
+{
+    int index = ui->structureTypeComboBox->findText(type);
+    if (index != -1) {
+        // index is valid so set is as the default
+        ui->structureTypeComboBox->setCurrentIndex(index);
+    }
+}
+
+void LinkTypeDialog::setDefaultAddress(QString address)
+{
+    ui->exprLineEdit->setText(address);
+
+    if (ui->addressLineEdit->text() == tr("Invalid Address")) {
+        return;
+    }
+
+    // check if the current address is already linked to a type and set it as default
+    QString type = findLinkedType(Core()->math(ui->addressLineEdit->text()));
+    if (!type.isEmpty()) {
+        setDefaultType(type);
+    }
+}
+
+
+void LinkTypeDialog::done(int r)
+{
+    if (r == QDialog::Accepted) {
+        QString address = ui->addressLineEdit->text();
+        if (Core()->isAddressMapped(Core()->math(address))) {
+            // Address is valid so link the type to the address
+            QString type = ui->structureTypeComboBox->currentText();
+            if (type == tr("(No Type)")) {
+                // Delete link
+                Core()->cmdRaw("tl- " + address);
+            } else {
+                // Create link
+                Core()->cmdRaw(QString("tl %1 = %2").arg(type).arg(address));
+            }
+            QDialog::done(r);
+
+            // Seek to the specified address
+            Core()->seek(address);
+
+            // Refresh the views
+            emit Core()->refreshCodeViews();
+            return;
+        }
+
+        // Address is invalid so display error message
+        QMessageBox::warning(this, tr("Error"), tr("The given address is invalid"));
+    } else {
+        QDialog::done(r);
+    }
+}
+
+QString LinkTypeDialog::findLinkedType(RVA address)
+{
+    if (address == RVA_INVALID) {
+        return QString();
+    }
+
+    QString ret = Core()->cmdRaw(QString("tls %1").arg(address));
+    if (ret.isEmpty()) {
+        // return empty string since the current address is not linked to a type
+        return QString();
+    }
+
+    // Extract the given type from returned data
+    // TODO: Implement "tlsj" in radare2 or some other function to directly get linked type
+    QString s = ret.split("\n").first();
+    return s.mid(1, s.size() - 2);
+}
+
+void LinkTypeDialog::on_exprLineEdit_textChanged(const QString &text)
+{
+    RVA addr = Core()->math(text);
+    if (Core()->isAddressMapped(addr)) {
+        ui->addressLineEdit->setText("0x" + QString::number(addr, 16));
+    } else {
+        ui->addressLineEdit->setText(tr("Invalid Address"));
+    }
+}

src/dialogs/LinkTypeDialog.h

@@ -0,0 +1,63 @@
+#ifndef LINKTYPEDIALOG_H
+#define LINKTYPEDIALOG_H
+
+#include "core/Cutter.h"
+#include <QDialog>
+
+namespace Ui {
+class LinkTypeDialog;
+}
+
+class LinkTypeDialog : public QDialog
+{
+    Q_OBJECT
+
+public:
+    explicit LinkTypeDialog(QWidget *parent = nullptr);
+    ~LinkTypeDialog();
+
+    /*!
+     * \brief Sets the default type which will be displayed in the combo box
+     * \param type Default type to be used as default type
+     */
+    void setDefaultType(const QString &type);
+
+    /*!
+     * \brief Sets the value of the default address which will be displayed
+     * If the given address is linked to a type, then it also sets the default
+     * type to the currently linked type
+     * \param address The address to be used as default address
+     */
+    void setDefaultAddress(QString address);
+
+private slots:
+
+    /*!
+     * \brief Overrides the done() method of QDialog
+     * On clicking the Ok button, it links a valid address to a type.
+     * If "(No Type)" is selected as type, it removes the link.
+     * In case of an invalid address, it displays error message
+     * \param r The value which will be returned by exec()
+     */
+    void done(int r) override;
+
+    /*!
+     * \brief Executed whenever the text inside exprLineEdit changes
+     * If expression evaluates to valid address, it is displayed in addressLineEdit
+     * Otherwise "Invalid Address" is shown in addressLineEdit
+     * \param text The current value of exprLineEdit
+     */
+    void on_exprLineEdit_textChanged(const QString &text);
+
+private:
+    Ui::LinkTypeDialog *ui;
+
+    /*!
+     * \brief Used for finding the type which is linked to the given address
+     * \param address
+     * \return The type linked to "address" if it exists, or empty string otherwise
+     */
+    QString findLinkedType(RVA address);
+};
+
+#endif // LINKTYPEDIALOG_H

src/dialogs/LinkTypeDialog.ui

@@ -0,0 +1,117 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<ui version="4.0">
+ <class>LinkTypeDialog</class>
+ <widget class="QDialog" name="LinkTypeDialog">
+  <property name="geometry">
+   <rect>
+    <x>0</x>
+    <y>0</y>
+    <width>500</width>
+    <height>105</height>
+   </rect>
+  </property>
+  <property name="minimumSize">
+   <size>
+    <width>500</width>
+    <height>0</height>
+   </size>
+  </property>
+  <property name="windowTitle">
+   <string>Dialog</string>
+  </property>
+  <layout class="QGridLayout" name="gridLayout">
+   <item row="1" column="2">
+    <widget class="QLabel" name="equalSymbolLabel">
+     <property name="text">
+      <string>=</string>
+     </property>
+    </widget>
+   </item>
+   <item row="1" column="3">
+    <widget class="QLineEdit" name="addressLineEdit">
+     <property name="enabled">
+      <bool>true</bool>
+     </property>
+     <property name="readOnly">
+      <bool>true</bool>
+     </property>
+    </widget>
+   </item>
+   <item row="1" column="1">
+    <widget class="QLineEdit" name="exprLineEdit">
+     <property name="placeholderText">
+      <string>Enter Address</string>
+     </property>
+    </widget>
+   </item>
+   <item row="0" column="0">
+    <widget class="QLabel" name="structureTypeLabel">
+     <property name="text">
+      <string>Structure Type</string>
+     </property>
+     <property name="buddy">
+      <cstring>structureTypeComboBox</cstring>
+     </property>
+    </widget>
+   </item>
+   <item row="0" column="1" colspan="3">
+    <widget class="QComboBox" name="structureTypeComboBox"/>
+   </item>
+   <item row="2" column="3">
+    <widget class="QDialogButtonBox" name="buttonBox">
+     <property name="orientation">
+      <enum>Qt::Horizontal</enum>
+     </property>
+     <property name="standardButtons">
+      <set>QDialogButtonBox::Cancel|QDialogButtonBox::Ok</set>
+     </property>
+    </widget>
+   </item>
+   <item row="1" column="0">
+    <widget class="QLabel" name="addressLabel">
+     <property name="text">
+      <string>Address/Flag</string>
+     </property>
+     <property name="buddy">
+      <cstring>exprLineEdit</cstring>
+     </property>
+    </widget>
+   </item>
+  </layout>
+ </widget>
+ <resources/>
+ <connections>
+  <connection>
+   <sender>buttonBox</sender>
+   <signal>accepted()</signal>
+   <receiver>LinkTypeDialog</receiver>
+   <slot>accept()</slot>
+   <hints>
+    <hint type="sourcelabel">
+     <x>248</x>
+     <y>254</y>
+    </hint>
+    <hint type="destinationlabel">
+     <x>157</x>
+     <y>274</y>
+    </hint>
+   </hints>
+  </connection>
+  <connection>
+   <sender>buttonBox</sender>
+   <signal>rejected()</signal>
+   <receiver>LinkTypeDialog</receiver>
+   <slot>reject()</slot>
+   <hints>
+    <hint type="sourcelabel">
+     <x>316</x>
+     <y>260</y>
+    </hint>
+    <hint type="destinationlabel">
+     <x>286</x>
+     <y>274</y>
+    </hint>
+   </hints>
+  </connection>
+ </connections>
+</ui>

src/menus/DisassemblyContextMenu.cpp

@@ -8,6 +8,7 @@
 #include "dialogs/EditVariablesDialog.h"
 #include "dialogs/SetToDataDialog.h"
 #include "dialogs/EditFunctionDialog.h"
+#include "dialogs/LinkTypeDialog.h"
 #include <QtCore>
 #include <QShortcut>
 #include <QJsonArray>
@@ -75,6 +76,10 @@ DisassemblyContextMenu::DisassemblyContextMenu(QWidget *parent)
     connect(structureOffsetMenu, SIGNAL(triggered(QAction*)),
             this, SLOT(on_actionStructureOffsetMenu_triggered(QAction*)));
 
+    initAction(&actionLinkType, tr("Link Type to Address"),
+               SLOT(on_actionLinkType_triggered()), getLinkTypeSequence());
+    addAction(&actionLinkType);
+
     initAction(&actionSetToCode, tr("Set as Code"),
                SLOT(on_actionSetToCode_triggered()), getSetToCodeSequence());
     addAction(&actionSetToCode);
@@ -239,6 +244,11 @@ void DisassemblyContextMenu::setCanCopy(bool enabled)
     this->canCopy = enabled;
 }
 
+void DisassemblyContextMenu::setCurHighlightedWord(const QString &text)
+{
+    this->curHighlightedWord = text;
+}
+
 void DisassemblyContextMenu::aboutToShowSlot()
 {
     // check if set immediate base menu makes sense
@@ -427,6 +437,11 @@ QKeySequence DisassemblyContextMenu::getDisplayOptionsSequence() const
     return {}; //TODO insert correct sequence
 }
 
+QKeySequence DisassemblyContextMenu::getLinkTypeSequence() const
+{
+    return {Qt::Key_L};
+}
+
 QList<QKeySequence> DisassemblyContextMenu::getAddBPSequence() const
 {
     return {Qt::Key_F2, Qt::CTRL + Qt::Key_B};
@@ -750,6 +765,13 @@ void DisassemblyContextMenu::on_actionStructureOffsetMenu_triggered(QAction *act
     Core()->applyStructureOffset(action->data().toString(), offset);
 }
 
+void DisassemblyContextMenu::on_actionLinkType_triggered()
+{
+    LinkTypeDialog dialog(this);
+    dialog.setDefaultAddress(curHighlightedWord);
+    dialog.exec();
+}
+
 void DisassemblyContextMenu::on_actionDeleteComment_triggered()
 {
     Core()->delComment(offset);

src/menus/DisassemblyContextMenu.h

@@ -20,6 +20,12 @@ public slots:
     void setOffset(RVA offset);
     void setCanCopy(bool enabled);
 
+    /*!
+     * \brief Sets the value of curHighlightedWord
+     * \param text The current highlighted word
+     */
+    void setCurHighlightedWord(const QString &text);
+
 private slots:
     void aboutToShowSlot();
 
@@ -63,6 +69,13 @@ private slots:
      */
     void on_actionStructureOffsetMenu_triggered(QAction *action);
 
+    /*!
+     * \brief Executed on selecting the "Link Type to Address" option
+     * Opens the LinkTypeDialog box from where the user can link the address
+     * to a type
+     */
+    void on_actionLinkType_triggered();
+
 private:
     QKeySequence getCopySequence() const;
     QKeySequence getCommentSequence() const;
@@ -78,8 +91,15 @@ private:
     QKeySequence getDisplayOptionsSequence() const;
     QList<QKeySequence> getAddBPSequence() const;
 
+    /*!
+     * \return the shortcut key for "Link Type to Address" option
+     */
+    QKeySequence getLinkTypeSequence() const;
+
+
     RVA offset;
     bool canCopy;
+    QString curHighlightedWord; // The current highlighted word
 
     QList<QAction *> anonymousActions;
 
@@ -110,6 +130,8 @@ private:
 
     QMenu *structureOffsetMenu;
 
+    QAction actionLinkType;
+
     QMenu *setBaseMenu;
     QAction actionSetBaseBinary;
     QAction actionSetBaseOctal;

src/widgets/DisassemblerGraphView.cpp

@@ -859,6 +859,9 @@ void DisassemblerGraphView::blockClicked(GraphView::GraphBlock &block, QMouseEve
 
     mMenu->setOffset(addr);
     mMenu->setCanCopy(highlight_token);
+    if (highlight_token) {
+        mMenu->setCurHighlightedWord(highlight_token->content);
+    }
     if (event->button() == Qt::RightButton) {
         mMenu->exec(event->globalPos());
     }

src/widgets/DisassemblyWidget.cpp

@@ -335,6 +335,7 @@ void DisassemblyWidget::highlightCurrentLine()
     QTextCursor cursor = mDisasTextEdit->textCursor();
     cursor.select(QTextCursor::WordUnderCursor);
     QString searchString = cursor.selectedText();
+    curHighlightedWord = searchString;
 
     cursor.movePosition(QTextCursor::StartOfLine);
     int listStartPos = cursor.position();
@@ -517,6 +518,13 @@ void DisassemblyWidget::cursorPositionChanged()
     seekFromCursor = false;
     highlightCurrentLine();
     mCtxMenu->setCanCopy(mDisasTextEdit->textCursor().hasSelection());
+    if (mDisasTextEdit->textCursor().hasSelection()) {
+        // A word is selected so use it
+        mCtxMenu->setCurHighlightedWord(mDisasTextEdit->textCursor().selectedText());
+    } else {
+        // No word is selected so use the word under the cursor
+        mCtxMenu->setCurHighlightedWord(curHighlightedWord);
+    }
 }
 
 void DisassemblyWidget::moveCursorRelative(bool up, bool page)

src/widgets/DisassemblyWidget.h

@@ -53,6 +53,8 @@ private:
     RVA bottomOffset;
     int maxLines;
 
+    QString curHighlightedWord;
+
     /*!
      * offset of lines below the first line of the current seek
      */

src/widgets/TypesWidget.cpp

@@ -4,6 +4,7 @@
 #include "common/Helpers.h"
 
 #include "dialogs/LoadNewTypesDialog.h"
+#include "dialogs/LinkTypeDialog.h"
 
 #include <QMenu>
 #include <QFileDialog>
@@ -220,11 +221,21 @@ void TypesWidget::setScrollMode()
 
 void TypesWidget::showTypesContextMenu(const QPoint &pt)
 {
+    QModelIndex index = ui->typesTreeView->indexAt(pt);
+
     QMenu menu(ui->typesTreeView);
     menu.addAction(ui->actionLoad_New_Types);
+
+    if (index.isValid()) {
+        TypeDescription t = index.data(TypesModel::TypeDescriptionRole).value<TypeDescription>();
+        if (t.category == "Struct") {
+            // Add "Link To Address" option
+            menu.addAction(ui->actionLink_Type_To_Address);
+        }
+    }
+
     menu.addAction(ui->actionExport_Types);
 
-    QModelIndex index = ui->typesTreeView->indexAt(pt);
     if (index.isValid()) {
         TypeDescription t = index.data(TypesModel::TypeDescriptionRole).value<TypeDescription>();
         if (t.category != "Typedef") {
@@ -279,3 +290,16 @@ void TypesWidget::on_actionDelete_Type_triggered()
         types_model->removeRow(index.row());
     }
 }
+
+void TypesWidget::on_actionLink_Type_To_Address_triggered()
+{
+    LinkTypeDialog dialog(this);
+
+    QModelIndex index = ui->typesTreeView->currentIndex();
+    if (index.isValid()) {
+        TypeDescription t = index.data(TypesModel::TypeDescriptionRole).value<TypeDescription>();
+        dialog.setDefaultType(t.type);
+    }
+
+    dialog.exec();
+}

src/widgets/TypesWidget.h

@@ -104,6 +104,12 @@ private slots:
      */
     void on_actionDelete_Type_triggered();
 
+    /*!
+     * \brief Executed on clicking the Link To Address option in the context menu
+     * Opens the LinkTypeDialog box from where the user can link a address to a type
+     */
+    void on_actionLink_Type_To_Address_triggered();
+
 private:
     std::unique_ptr<Ui::TypesWidget> ui;
 

src/widgets/TypesWidget.ui

@@ -95,6 +95,11 @@
     <string>Delete Type</string>
    </property>
   </action>
+  <action name="actionLink_Type_To_Address">
+   <property name="text">
+    <string>Link Type to Address</string>
+   </property>
+  </action>
  </widget>
  <customwidgets>
   <customwidget>