Limit on incoming kdb messages

notes.md

@@ -3,6 +3,7 @@
     * Fix logic to detect server disconnect
     * Connection status added to the status bar
     * Capturing connection statistics (new menu action)
+    * Limit on incoming message (new settings)
 * File drag and drop (no support for a folder drag and drop) 
 * Persist location of editor windows
 * Keep config properties sorted on disk

src/kx/ConnectionStateListener.java

@@ -1,6 +1,10 @@
 package kx;
 
+import java.io.IOException;
+
 public interface ConnectionStateListener {
 
     void connectionStateChange(boolean connected);
+    void checkIncomingLimit(long msgLength) throws IOException;
+
 }

src/kx/KConnection.java

@@ -208,6 +208,10 @@ public void run() {
                     IPC ipc = new IPC(buffer, 4, false, isLittleEndian);
                     final int msgLength = ipc.ri() - 8;
 
+                    if (response && connectionStateListener != null) {
+                        connectionStateListener.checkIncomingLimit(msgLength);
+                    }
+
                     stats.receivedBytes(msgLength);
                     ProgressCallback progress = getProgressCallback();
                     if (progress!=null) {

src/studio/kdb/Config.java

@@ -6,6 +6,7 @@
 import studio.core.Credentials;
 import studio.core.DefaultAuthenticationMechanism;
 import studio.kdb.config.ActionOnExit;
+import studio.kdb.config.KdbMessageLimitAction;
 import studio.ui.ServerList;
 import studio.ui.Util;
 import studio.utils.*;
@@ -97,6 +98,8 @@ private enum ConfigType { STRING, INT, DOUBLE, BOOLEAN, FONT, BOUNDS, COLOR, ENU
     public static final String SESSION_INVALIDATION_ENABLED = configDefault("sessionInvalidationEnabled", ConfigType.BOOLEAN, false);
     public static final String SESSION_INVALIDATION_TIMEOUT_IN_HOURS = configDefault("sessionInvalidationTimeoutInHours", ConfigType.INT, 12);
     public static final String SESSION_REUSE = configDefault("sessionsReuse", ConfigType.BOOLEAN, true);
+    public static final String KDB_MESSAGE_SIZE_LIMIT_MB = configDefault("kdbMessageSizeLimitMB", ConfigType.INT, 10);
+    public static final String KDB_MESSAGE_SIZE_LIMIT_ACTION = configDefault("kdbMessageSizeLimitAction", ConfigType.ENUM, KdbMessageLimitAction.ASK);
 
     private enum FontStyle {
         Plain(Font.PLAIN), Bold(Font.BOLD), Italic(Font.ITALIC), ItalicAndBold(Font.BOLD|Font.ITALIC);

src/studio/kdb/Session.java

@@ -6,13 +6,18 @@
 import studio.core.AuthenticationManager;
 import studio.core.Credentials;
 import studio.core.IAuthenticationMechanism;
+import studio.kdb.config.KdbMessageLimitAction;
 import studio.ui.EditorTab;
+import studio.ui.StudioOptionPane;
+import studio.ui.StudioWindow;
 
+import javax.swing.*;
 import java.io.IOException;
 import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.concurrent.atomic.AtomicInteger;
 
 public class Session implements ConnectionStateListener, KAuthentication {
     private KConnection kConn;
@@ -57,6 +62,35 @@ public void connectionStateChange(boolean connected) {
         }
     }
 
+    @Override
+    public void checkIncomingLimit(long msgLength) throws IOException {
+        KdbMessageLimitAction action = Config.getInstance().getEnum(Config.KDB_MESSAGE_SIZE_LIMIT_ACTION);
+        long limit = 1_000_000L * Config.getInstance().getInt(Config.KDB_MESSAGE_SIZE_LIMIT_MB);
+        if (msgLength < limit) return;
+        final String msg = "Incoming message size " + msgLength + " breached the limit of " + limit + ".";
+        if (action == KdbMessageLimitAction.BLOCK) {
+            throw new IOException(msg);
+        }
+
+        AtomicInteger choice = new AtomicInteger();
+        try {
+            SwingUtilities.invokeAndWait(() -> {
+                choice.set(StudioOptionPane.showYesNoDialog(StudioWindow.getActiveStudioWindow(),
+                        msg + "\n\nDownloading large amount of data can result in OutOfMemoryError.\n" +
+                                "Note: you can change the limit as well as action in the Settings menu.\n\n" +
+                                "Download?",
+                        "Too Big Incoming Message"
+                ));
+            });
+        } catch (Exception e) {
+            throw new IOException("Error during showing incoming message limit dialog", e);
+        }
+
+        if (choice.get() != JOptionPane.YES_OPTION) {
+            throw new IOException(msg);
+        }
+    }
+
     private Session(Server server) {
         this.server = server;
         init();

src/studio/kdb/config/KdbMessageLimitAction.java

@@ -0,0 +1,18 @@
+package studio.kdb.config;
+
+public enum KdbMessageLimitAction {
+    ASK("show dialog to download"),
+    BLOCK("close session (do not download)");
+
+    private final String description;
+
+    KdbMessageLimitAction(String description) {
+        this.description = description;
+    }
+
+    @Override
+    public String toString() {
+        return description;
+    }
+
+}

src/studio/ui/SettingsDialog.java

@@ -5,6 +5,7 @@
 import studio.kdb.Config;
 import studio.kdb.KFormatContext;
 import studio.kdb.config.ActionOnExit;
+import studio.kdb.config.KdbMessageLimitAction;
 import studio.ui.settings.FontSelectionPanel;
 import studio.utils.LineEnding;
 
@@ -24,6 +25,8 @@ public class SettingsDialog extends EscapeDialog {
     private JCheckBox chBoxSessionInvalidation;
     private JFormattedTextField txtSessionInvalidation;
     private JCheckBox chBoxSessionReuse;
+    private JFormattedTextField txtKdbMessageSizeLimit;
+    private JComboBox<KdbMessageLimitAction> comboBoxKdbMessageSizeAction;
     private JCheckBox chBoxShowServerCombo;
     private JCheckBox chBoxAutoSave;
     private JComboBox<ActionOnExit> comboBoxActionOnExit;
@@ -179,6 +182,14 @@ private void initComponents() {
         chBoxSessionReuse = new JCheckBox("Reuse kdb connection between tabs");
         chBoxSessionReuse.setSelected(CONFIG.getBoolean(Config.SESSION_REUSE));
 
+        JLabel lblKdbMessageSize = new JLabel("When the incoming message is greater than ");
+        txtKdbMessageSizeLimit = new JFormattedTextField(formatter);
+        txtKdbMessageSizeLimit.setValue(CONFIG.getInt(Config.KDB_MESSAGE_SIZE_LIMIT_MB));
+
+        JLabel lblKdbMessageSizeSuffix = new JLabel("MB  ");
+        comboBoxKdbMessageSizeAction = new JComboBox<>(KdbMessageLimitAction.values());
+        comboBoxKdbMessageSizeAction.setSelectedItem(CONFIG.getEnum(Config.KDB_MESSAGE_SIZE_LIMIT_ACTION));
+
         chBoxEmulateTab = new JCheckBox("Emulate tab with spaces");
         chBoxEmulateTab.setSelected(CONFIG.getBoolean(Config.EDITOR_TAB_EMULATED));
         txtEmulatedTabSize = new JFormattedTextField(formatter);
@@ -265,6 +276,7 @@ private void initComponents() {
                         .addLineAndGlue(lblActionOnExit, comboBoxActionOnExit)
                         .addLineAndGlue(chBoxSessionInvalidation, txtSessionInvalidation, lblSessionInvalidationSuffix)
                         .addLineAndGlue(chBoxSessionReuse)
+                        .addLineAndGlue(lblKdbMessageSize, txtKdbMessageSizeLimit, lblKdbMessageSizeSuffix, comboBoxKdbMessageSizeAction)
                         .addLine(lblAuthMechanism, comboBoxAuthMechanism, lblUser, txtUser, lblPassword, txtPassword)
         );
 
@@ -320,6 +332,8 @@ public void saveSettings() {
         CONFIG.setBoolean(Config.SESSION_INVALIDATION_ENABLED, chBoxSessionInvalidation.isSelected());
         CONFIG.setInt(Config.SESSION_INVALIDATION_TIMEOUT_IN_HOURS, Math.max(1,(Integer)txtSessionInvalidation.getValue()));
         CONFIG.setBoolean(Config.SESSION_REUSE, chBoxSessionReuse.isSelected());
+        CONFIG.setInt(Config.KDB_MESSAGE_SIZE_LIMIT_MB, Math.max(1, (Integer)txtKdbMessageSizeLimit.getValue()));
+        CONFIG.setEnum(Config.KDB_MESSAGE_SIZE_LIMIT_ACTION, (KdbMessageLimitAction) comboBoxKdbMessageSizeAction.getSelectedItem());
         CONFIG.setDefaultAuthMechanism(auth);
         CONFIG.setDefaultCredentials(auth, new Credentials(getUser(), getPassword()));