Modify coroutines and updating

+ auto scroll
TypeError · Dec 14, 2019 · 9d4f906 · 9d4f906
1 parent ee54127
commit 9d4f906
Show file tree

Hide file tree

Showing 2 changed files with 30 additions and 45 deletions.
diff --git a/src/BookmarkOptions.kt b/src/BookmarkOptions.kt
@@ -1,9 +1,5 @@
 package burp
 
-import kotlinx.coroutines.Dispatchers
-import kotlinx.coroutines.GlobalScope
-import kotlinx.coroutines.launch
-import kotlinx.coroutines.swing.Swing
 import java.awt.FlowLayout
 import javax.swing.*
 
@@ -27,10 +23,7 @@ class BookmarkOptions(
         clearButton.addActionListener { clearBookmarks() }
         searchBar.addActionListener { searchBookmarks() }
         searchButton.addActionListener { searchBookmarks() }
-        resetButton.addActionListener {
-            searchBar.text = ""
-            resetSearch()
-        }
+        resetButton.addActionListener { resetSearch() }
         searchPanel.add(searchLabel)
         searchPanel.add(searchBar)
         searchPanel.add(searchButton)
@@ -43,44 +36,34 @@ class BookmarkOptions(
     }
 
     private fun loadHighlightedRequests() {
-        GlobalScope.launch {
-            val bookmarks = bookmarksPanel.bookmarks
-            val highlightedProxyHistory = callbacks.proxyHistory.filter { it.highlight != null }
-            val bookmarkRequests = bookmarks.map { callbacks.helpers.bytesToString(it.requestResponse.request) }
-            val bookmarksToAdd = highlightedProxyHistory
-                .filter { !bookmarkRequests.contains(callbacks.helpers.bytesToString(it.request)) }.toTypedArray()
-            launch(Dispatchers.Swing) {
-                bookmarksPanel.addBookmark(bookmarksToAdd)
-                bookmarksPanel.model.filteredBookmarks()
-            }
-        }
+        val bookmarks = bookmarksPanel.bookmarks
+        val highlightedProxyHistory = callbacks.proxyHistory.filter { it.highlight != null }
+        val bookmarkRequests = bookmarks.map { callbacks.helpers.bytesToString(it.requestResponse.request) }
+        val bookmarksToAdd = highlightedProxyHistory
+            .filter { !bookmarkRequests.contains(callbacks.helpers.bytesToString(it.request)) }.toTypedArray()
+        bookmarksPanel.addBookmark(bookmarksToAdd)
     }
 
     private fun searchBookmarks() {
         val searchText = searchBar.text
         val bookmarks = this.bookmarksPanel.bookmarks
         if (searchText.isNotEmpty()) {
-            GlobalScope.launch {
-                val filteredBookmarks = bookmarks
-                    .filter {
-                        callbacks.helpers.bytesToString(it.requestResponse.request).contains(searchText) ||
-                                callbacks.helpers.bytesToString(it.requestResponse.response).contains(searchText)
-                    }.toMutableList()
-                launch(Dispatchers.Swing) {
-                    bookmarksPanel.model.filteredBookmarks(filteredBookmarks)
-                }
-            }
-        } else {
-            bookmarksPanel.model.filteredBookmarks()
+            val filteredBookmarks = bookmarks
+                .filter {
+                    callbacks.helpers.bytesToString(it.requestResponse.request).contains(searchText) ||
+                            callbacks.helpers.bytesToString(it.requestResponse.response).contains(searchText)
+                }.toMutableList()
+            bookmarksPanel.model.refreshBookmarks(filteredBookmarks)
         }
     }
 
     private fun resetSearch() {
-        bookmarksPanel.model.filteredBookmarks()
+        searchBar.text = ""
+        bookmarksPanel.model.refreshBookmarks()
     }
 
     private fun clearBookmarks() {
         bookmarksPanel.model.bookmarks.clear()
-        bookmarksPanel.model.filteredBookmarks()
+        bookmarksPanel.model.refreshBookmarks()
     }
 }
diff --git a/src/BookmarkTab.kt b/src/BookmarkTab.kt
@@ -4,6 +4,7 @@ import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.GlobalScope
 import kotlinx.coroutines.launch
 import kotlinx.coroutines.swing.Swing
+import kotlinx.coroutines.withContext
 import java.awt.FlowLayout
 import java.net.URL
 import java.time.LocalDateTime
@@ -144,6 +145,10 @@ class BookmarksPanel(private val callbacks: IBurpExtenderCallbacks) {
             requestResponse.highlight = "magenta"
             requestResponse.comment = "[^]"
         }
+
+        SwingUtilities.invokeLater {
+            table.scrollRectToVisible(table.getCellRect(table.rowCount - 1, table.columnCount, true))
+        }
     }
 
     private fun getTitle(response: ByteArray?): String {
@@ -158,11 +163,10 @@ class BookmarksPanel(private val callbacks: IBurpExtenderCallbacks) {
     private fun repeatRequest() {
         GlobalScope.launch(Dispatchers.IO) {
             val requestResponse = callbacks.makeHttpRequest(messageEditor.httpService, requestViewer?.message)
-            responseViewer?.setMessage(requestResponse.response, false)
-            if (repeatInTable.isSelected) {
-                launch(Dispatchers.Swing) {
-                    createBookmark(requestResponse, true, false)
-                    model.filteredBookmarks()
+            withContext(Dispatchers.Swing) {
+                responseViewer?.setMessage(requestResponse.response, false)
+                if (repeatInTable.isSelected) {
+                    createBookmark(requestResponse, repeated = true, proxyHistory = false)
                 }
             }
         }
@@ -250,23 +254,21 @@ class BookmarksModel : AbstractTableModel() {
 
     fun addBookmark(bookmark: Bookmark) {
         bookmarks.add(bookmark)
-        filteredBookmarks()
-        fireTableDataChanged()
+        displayedBookmarks = bookmarks
+        fireTableRowsInserted(displayedBookmarks.lastIndex, displayedBookmarks.lastIndex)
     }
 
     fun removeBookmarks(selectedBookmarks: MutableList<Bookmark>) {
         bookmarks.removeAll(selectedBookmarks)
-        filteredBookmarks()
-        fireTableDataChanged()
+        refreshBookmarks()
     }
 
     fun clearBookmarks() {
         bookmarks.clear()
-        filteredBookmarks()
-        fireTableDataChanged()
+        refreshBookmarks()
     }
 
-    fun filteredBookmarks(updatedBookmarks: MutableList<Bookmark> = bookmarks) {
+    fun refreshBookmarks(updatedBookmarks: MutableList<Bookmark> = bookmarks) {
         displayedBookmarks = updatedBookmarks
         fireTableDataChanged()
     }