feat: add preemptive size-limiting to UDP import

Mahlet-Inc · Sep 22, 2020 · 9417f8e · 9417f8e
1 parent bc050f7
commit 9417f8e
Show file tree

Hide file tree

Showing 4 changed files with 109 additions and 7 deletions.
diff --git a/src/hobbits-plugins/importerexporters/UdpData/udpdata.cpp b/src/hobbits-plugins/importerexporters/UdpData/udpdata.cpp
@@ -35,6 +35,13 @@ bool UdpData::canImport()
 QString UdpData::getImportLabelForState(QJsonObject pluginState)
 {
     if (pluginState.contains("port") && pluginState.value("port").isDouble()) {
+
+        if (pluginState.contains("limit") && pluginState.value("limit").isDouble()) {
+            int value = pluginState.value("limit").toInt();
+            if (value > 0) {
+                return QString("UDP Listen for %2 KB on port %1").arg(pluginState.value("port").toInt()).arg(double(value)/1000.0);
+            }
+        }
         return QString("UDP Listen on port %1").arg(pluginState.value("port").toInt());
     }
     return "";
@@ -54,14 +61,24 @@ QSharedPointer<ImportResult> UdpData::importBits(QJsonObject pluginState)
     QSharedPointer<UdpReceiver> receiver = QSharedPointer<UdpReceiver>(new UdpReceiver());
     if (pluginState.contains("port") && pluginState.value("port").isDouble()) {
         receiver->setPort(pluginState.value("port").toInt());
+        if (pluginState.contains("limit") && pluginState.value("limit").isDouble()) {
+            receiver->setLimit(pluginState.value("limit").toInt());
+        }
         receiver->startListening();
     }
 
     if (receiver->exec()) {
-        QSharedPointer<BitContainer> container = BitContainer::create(receiver->getDownloadedData());
+        QSharedPointer<BitContainer> container;
+        if (receiver->getLimit() > 0) {
+            container = BitContainer::create(receiver->getDownloadedData(), receiver->getLimit() * 8);
+        }
+        else {
+            container = BitContainer::create(receiver->getDownloadedData());
+        }
         container->setName("UDP Import");
         QJsonObject state;
         state.insert("port", receiver->getPort());
+        state.insert("limit", receiver->getLimit());
         return ImportResult::result(container, state);
     }
     else {

diff --git a/src/hobbits-plugins/importerexporters/UdpData/udpreceiver.cpp b/src/hobbits-plugins/importerexporters/UdpData/udpreceiver.cpp
@@ -1,6 +1,7 @@
 #include "udpreceiver.h"
 #include "ui_udpreceiver.h"
 #include <QNetworkDatagram>
+#include <QTimer>
 
 UdpReceiver::UdpReceiver(QWidget *parent) :
     QDialog(parent),
@@ -13,6 +14,7 @@ UdpReceiver::UdpReceiver(QWidget *parent) :
     setWindowTitle("Import via UDP");
 
     ui->pb_stop->setEnabled(false);
+
 }
 
 UdpReceiver::~UdpReceiver()
@@ -36,6 +38,27 @@ void UdpReceiver::setPort(int port)
     ui->sb_port->setValue(port);
 }
 
+int UdpReceiver::getLimit()
+{
+    if (ui->ck_limit->isChecked()) {
+        return ui->sb_limit->value() * 1000;
+    }
+    else {
+        return -1;
+    }
+}
+
+void UdpReceiver::setLimit(int bytes)
+{
+    if (bytes > 0) {
+        ui->ck_limit->setChecked(true);
+        ui->sb_limit->setValue(bytes / 1000);
+    }
+    else {
+        ui->ck_limit->setChecked(false);
+    }
+}
+
 QString UdpReceiver::getError()
 {
     if (m_socket) {
@@ -98,6 +121,11 @@ void UdpReceiver::socketRead()
         m_downloadFile.write(datagram.data());
     }
     ui->lb_progress->setText(QString("%1 KBs Received").arg(double(m_downloadFile.size())/1000.0, 1, 'g', 3));
+
+    if (getLimit() > 0 && m_downloadFile.size() >= getLimit()) {
+        disconnect(m_socket, &QUdpSocket::readyRead, this, &UdpReceiver::socketRead);
+        QTimer::singleShot(10, this, &UdpReceiver::stopListening);
+    }
 }
 
 void UdpReceiver::on_pb_start_pressed()

diff --git a/src/hobbits-plugins/importerexporters/UdpData/udpreceiver.h b/src/hobbits-plugins/importerexporters/UdpData/udpreceiver.h
@@ -20,6 +20,8 @@ class UdpReceiver : public QDialog
     QTemporaryFile* getDownloadedData();
     int getPort();
     void setPort(int);
+    int getLimit();
+    void setLimit(int bytes);
     QString getError();
 
 public slots:

diff --git a/src/hobbits-plugins/importerexporters/UdpData/udpreceiver.ui b/src/hobbits-plugins/importerexporters/UdpData/udpreceiver.ui
@@ -6,8 +6,8 @@
    <rect>
     <x>0</x>
     <y>0</y>
-    <width>358</width>
-    <height>142</height>
+    <width>373</width>
+    <height>219</height>
    </rect>
   </property>
   <property name="windowTitle">
@@ -18,6 +18,12 @@
     <layout class="QFormLayout" name="formLayout">
      <item row="0" column="0">
       <widget class="QLabel" name="label">
+       <property name="sizePolicy">
+        <sizepolicy hsizetype="Preferred" vsizetype="Expanding">
+         <horstretch>0</horstretch>
+         <verstretch>0</verstretch>
+        </sizepolicy>
+       </property>
        <property name="text">
         <string>Port</string>
        </property>
@@ -36,15 +42,47 @@
        </property>
       </widget>
      </item>
-     <item row="1" column="0" colspan="2">
-      <widget class="QPushButton" name="pb_start">
+     <item row="1" column="1">
+      <widget class="QSpinBox" name="sb_limit">
+       <property name="enabled">
+        <bool>false</bool>
+       </property>
+       <property name="suffix">
+        <string/>
+       </property>
+       <property name="minimum">
+        <number>1</number>
+       </property>
+       <property name="maximum">
+        <number>10000000</number>
+       </property>
+       <property name="value">
+        <number>1000</number>
+       </property>
+      </widget>
+     </item>
+     <item row="1" column="0">
+      <widget class="QCheckBox" name="ck_limit">
+       <property name="sizePolicy">
+        <sizepolicy hsizetype="Minimum" vsizetype="Expanding">
+         <horstretch>0</horstretch>
+         <verstretch>0</verstretch>
+        </sizepolicy>
+       </property>
        <property name="text">
-        <string>Begin Listening</string>
+        <string>KB Limit</string>
        </property>
       </widget>
      </item>
     </layout>
    </item>
+   <item>
+    <widget class="QPushButton" name="pb_start">
+     <property name="text">
+      <string>Begin Listening</string>
+     </property>
+    </widget>
+   </item>
    <item>
     <layout class="QHBoxLayout" name="horizontalLayout">
      <item>
@@ -66,5 +104,22 @@
   </layout>
  </widget>
  <resources/>
- <connections/>
+ <connections>
+  <connection>
+   <sender>ck_limit</sender>
+   <signal>toggled(bool)</signal>
+   <receiver>sb_limit</receiver>
+   <slot>setEnabled(bool)</slot>
+   <hints>
+    <hint type="sourcelabel">
+     <x>66</x>
+     <y>79</y>
+    </hint>
+    <hint type="destinationlabel">
+     <x>171</x>
+     <y>83</y>
+    </hint>
+   </hints>
+  </connection>
+ </connections>
 </ui>