Decouple checking ddl from classifying h2 stmts (#28510)

* Decouple checking ddl from classifying h2 stmts

- should enable followup for easily blocking more kinds of queries
- check all statements to make sure they aren't "ddl".

* fix classify-query

* linter fixes + get-field refactor

* return the CommandInterface values as ints

Reach into the CommandList when needed

* docstring wording

* catch invalid queries -- they can't be classified

* respond to review comments

- fix see: comment
- remove CommandList typehint
- rename command-list -> command
- skip checking disallowed ddl commands when query is `nil`

then also:
- fix a linter error
- add a test for `nil` query ddl command checking

* rename command-list-type -> first-command-type

* fix schema + handle h2-parser build failure

src/metabase/driver/h2.clj

@@ -19,6 +19,7 @@
    [metabase.util.honey-sql-2 :as h2x]
    [metabase.util.i18n :refer [deferred-tru tru]]
    [metabase.util.log :as log]
+   [metabase.util.malli :as mu]
    [metabase.util.ssh :as ssh])
   (:import
    (java.sql Clob ResultSet ResultSetMetaData)
@@ -111,10 +112,15 @@
 (defn- get-field
   "Returns value of private field. This function is used to bypass field protection to instantiate
    a low-level H2 Parser object in order to detect DDL statements in queries."
-  [obj field]
-  (.get (doto (.getDeclaredField (class obj) field)
-          (.setAccessible true))
-        obj))
+  ([obj field]
+   (.get (doto (.getDeclaredField (class obj) field)
+           (.setAccessible true))
+         obj))
+  ([obj field or-else]
+   (try (get-field obj field)
+        (catch java.lang.NoSuchFieldException _e
+          ;; when there are no fields: return or-else
+          or-else))))
 
 (defn- make-h2-parser
   "Returns an H2 Parser object for the given (H2) database ID"
@@ -127,22 +133,57 @@
       (when (instance? SessionLocal session)
         (Parser. session)))))
 
-(defn- contains-ddl?
+(mu/defn ^:private classify-query :- [:maybe
+                                      [:map
+                                       [:command-types [:vector pos-int?]]
+                                       [:remaining-sql [:maybe :string]]]]
+  "Takes an h2 db id, and a query, returns the command-types from `query` and any remaining sql.
+   More info on command types here:
+   https://github.com/h2database/h2database/blob/master/h2/src/main/org/h2/command/CommandInterface.java
+
+  If the h2 parser cannot be built, returns `nil`.
+
+  - Each `command-type` corresponds to a value in org.h2.command.CommandInterface, and match the commands from `query` in order.
+  - `remaining-sql` is a nillable sql string that is unable to be classified without running preceding queries first.
+    Usually if `remaining-sql` exists we will deny the query."
   [database query]
   (when-let [h2-parser (make-h2-parser database)]
     (try
-      (let [command      (.prepareCommand h2-parser query)
-            command-type (.getCommandType command)]
-        ;; TODO: do we need to handle CommandList?
-        ;; Command types are organized with all DDL commands listed first
-        ;; see https://github.com/h2database/h2database/blob/master/h2/src/main/org/h2/command/CommandInterface.java
-        (< command-type CommandInterface/ALTER_SEQUENCE))
-      ;; if the query is invalid, then it isn't DDL
-      (catch Throwable _ false))))
+      (let [command            (.prepareCommand h2-parser query)
+            first-command-type (.getCommandType command)
+            command-types      (cond-> [first-command-type]
+                                 (not (instance? org.h2.command.CommandContainer command))
+                                 (into
+                                  (map #(.getType ^org.h2.command.Prepared %))
+                                  ;; when there are no fields: return no commands
+                                  (get-field command "commands" [])))]
+        {:command-types command-types
+         ;; when there is no remaining sql: return nil for remaining-sql
+         :remaining-sql (get-field command "remaining" nil)})
+      ;; only valid queries can be classified.
+      (catch org.h2.message.DbException _
+        {:command-types [] :remaining-sql nil}))))
+
+(defn- cmd-type-ddl? [cmd-type]
+  ;; Command types are organized with all DDL commands listed first, so all ddl commands are before ALTER_SEQUENCE.
+  ;; see https://github.com/h2database/h2database/blob/master/h2/src/main/org/h2/command/CommandInterface.java#L297
+  (< cmd-type CommandInterface/ALTER_SEQUENCE))
+
+(defn- contains-ddl? [{:keys [command-types remaining-sql]}]
+  (let [cmd-type-nums command-types]
+    (boolean
+     (or (some cmd-type-ddl? cmd-type-nums)
+         (some? remaining-sql)))))
+
+;; TODO: black-list RUNSCRIPT, and a bunch more -- but they're not technically ddl. Should be simple to build off of [[classify-query]].
+;; e.g.: similar to contains-ddl? but instead of cmd-type-ddl? use: #(#{CommandInterface/RUNSCRIPT} %)
 
 (defn- check-disallow-ddl-commands [{:keys [database] {:keys [query]} :native}]
-  (when (and query (contains-ddl? database query))
-    (throw (IllegalArgumentException. "DDL commands are not allowed to be used with h2."))))
+  (when query
+    (when-let [query-classification (classify-query database query)]
+      (when (contains-ddl? query-classification)
+        (throw (ex-info "IllegalArgument: DDL commands are not allowed to be used with h2."
+                        {:classification query-classification}))))))
 
 (defmethod driver/execute-reducible-query :h2
   [driver query chans respond]

test/metabase/driver/h2_test.clj

@@ -167,15 +167,29 @@
                       "ORDER BY ATTEMPTS.DATE ASC")
                  (some-> (qp/compile query) :query pretty-sql))))))))
 
+
 (deftest classify-ddl-test
   (mt/test-driver :h2
-    (are [query] (= false (#'h2/contains-ddl? (u/the-id (mt/db)) query))
+    (are [query] (= false (#'h2/contains-ddl? (#'h2/classify-query (u/the-id (mt/db)) query)))
       "select 1"
       "update venues set name = 'bill'"
       "delete venues"
       "select 1;
        update venues set name = 'bill';
-       delete venues;")
+       delete venues;"
+      "update venues set name = 'stomp';"
+      "select * from venues; update venues set name = 'stomp';"
+      "update venues set name = 'stomp'; select * from venues;")
+
+    (are [query] (= true (#'h2/contains-ddl? (#'h2/classify-query (u/the-id (mt/db)) query)))
+      "select * from venues; update venues set name = 'stomp';
+       CREATE ALIAS EXEC AS 'String shellexec(String cmd) throws java.io.IOException {Runtime.getRuntime().exec(cmd);return \"y4tacker\";}';
+       EXEC ('open -a Calculator.app')"
+      "select * from venues; update venues set name = 'stomp';
+       CREATE ALIAS EXEC AS 'String shellexec(String cmd) throws java.io.IOException {Runtime.getRuntime().exec(cmd);return \"y4tacker\";}';"
+      "CREATE ALIAS EXEC AS 'String shellexec(String cmd) throws java.io.IOException {Runtime.getRuntime().exec(cmd);return \"y4tacker\";}';")
+
+    (is (= nil (#'h2/check-disallow-ddl-commands {:database (u/the-id (mt/db)) :native {:query nil}})))
 
     (is (= nil (#'h2/check-disallow-ddl-commands
                 {:database (u/the-id (mt/db))
@@ -184,14 +198,12 @@
                                            ["select 1"
                                             "update venues set name = 'bill'"
                                             "delete venues"])}})))
-    (let [trigger-creation-attempt
-          (str/join "\n" ["DROP TRIGGER IF EXISTS MY_SPECIAL_TRIG;"
-                          "CREATE OR REPLACE TRIGGER MY_SPECIAL_TRIG BEFORE SELECT ON INFORMATION_SCHEMA.Users AS '';"
-                          "SELECT * FROM INFORMATION_SCHEMA.Users;"])]
-      (is (thrown?
-           IllegalArgumentException
-           #"DDL commands are not allowed to be used with h2."
-           (#'h2/check-disallow-ddl-commands
-            {:database (u/the-id (mt/db))
-             :engine :h2
-             :native {:query trigger-creation-attempt}}))))))
+    (let [trigger-creation-attempt (str/join "\n" ["DROP TRIGGER IF EXISTS MY_SPECIAL_TRIG;"
+                                                   "CREATE OR REPLACE TRIGGER MY_SPECIAL_TRIG BEFORE SELECT ON INFORMATION_SCHEMA.Users AS '';"
+                                                   "SELECT * FROM INFORMATION_SCHEMA.Users;"])]
+      (is (thrown? clojure.lang.ExceptionInfo
+                   #"DDL commands are not allowed to be used with h2."
+                   (#'h2/check-disallow-ddl-commands
+                    {:database (u/the-id (mt/db))
+                     :engine :h2
+                     :native {:query trigger-creation-attempt}}))))))