Merge pull request #3117 from harawata/3115-secure-invocation

Prevent `Invocation` from invoking arbitrary method

Author: harawata

src/main/java/org/apache/ibatis/plugin/Invocation.java

@@ -1,5 +1,5 @@
 /*
- *    Copyright 2009-2022 the original author or authors.
+ *    Copyright 2009-2024 the original author or authors.
  *
  *    Licensed under the Apache License, Version 2.0 (the "License");
  *    you may not use this file except in compliance with the License.
@@ -17,17 +17,29 @@
 
 import java.lang.reflect.InvocationTargetException;
 import java.lang.reflect.Method;
+import java.util.Arrays;
+import java.util.List;
+
+import org.apache.ibatis.executor.Executor;
+import org.apache.ibatis.executor.parameter.ParameterHandler;
+import org.apache.ibatis.executor.resultset.ResultSetHandler;
+import org.apache.ibatis.executor.statement.StatementHandler;
 
 /**
  * @author Clinton Begin
  */
 public class Invocation {
 
+  private static final List<Class<?>> targetClasses = Arrays.asList(Executor.class, ParameterHandler.class,
+      ResultSetHandler.class, StatementHandler.class);
   private final Object target;
   private final Method method;
   private final Object[] args;
 
   public Invocation(Object target, Method method, Object[] args) {
+    if (!targetClasses.contains(method.getDeclaringClass())) {
+      throw new IllegalArgumentException("Method '" + method + "' is not supported as a plugin target.");
+    }
     this.target = target;
     this.method = method;
     this.args = args;

src/test/java/org/apache/ibatis/plugin/Mapper.java

@@ -0,0 +1,26 @@
+/*
+ *    Copyright 2009-2024 the original author or authors.
+ *
+ *    Licensed under the Apache License, Version 2.0 (the "License");
+ *    you may not use this file except in compliance with the License.
+ *    You may obtain a copy of the License at
+ *
+ *       https://www.apache.org/licenses/LICENSE-2.0
+ *
+ *    Unless required by applicable law or agreed to in writing, software
+ *    distributed under the License is distributed on an "AS IS" BASIS,
+ *    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ *    See the License for the specific language governing permissions and
+ *    limitations under the License.
+ */
+
+package org.apache.ibatis.plugin;
+
+import org.apache.ibatis.annotations.Select;
+
+public interface Mapper {
+
+  @Select("select name from users where id = #{id}")
+  String selectNameById(Integer id);
+
+}

src/test/java/org/apache/ibatis/plugin/PluginTest.java

@@ -1,5 +1,5 @@
 /*
- *    Copyright 2009-2023 the original author or authors.
+ *    Copyright 2009-2024 the original author or authors.
  *
  *    Licensed under the Apache License, Version 2.0 (the "License");
  *    you may not use this file except in compliance with the License.
@@ -16,27 +16,87 @@
 package org.apache.ibatis.plugin;
 
 import static org.junit.jupiter.api.Assertions.assertEquals;
-import static org.junit.jupiter.api.Assertions.assertNotEquals;
+import static org.junit.jupiter.api.Assertions.fail;
 
+import java.io.Reader;
+import java.sql.Connection;
 import java.util.HashMap;
 import java.util.Map;
 
+import org.apache.ibatis.BaseDataTest;
+import org.apache.ibatis.executor.statement.StatementHandler;
+import org.apache.ibatis.io.Resources;
+import org.apache.ibatis.session.SqlSession;
+import org.apache.ibatis.session.SqlSessionFactory;
+import org.apache.ibatis.session.SqlSessionFactoryBuilder;
+import org.junit.jupiter.api.BeforeAll;
 import org.junit.jupiter.api.Test;
 
 class PluginTest {
 
+  private static SqlSessionFactory sqlSessionFactory;
+
+  @BeforeAll
+  static void setUp() throws Exception {
+    try (Reader reader = Resources.getResourceAsReader("org/apache/ibatis/plugin/mybatis-config.xml")) {
+      sqlSessionFactory = new SqlSessionFactoryBuilder().build(reader);
+    }
+    BaseDataTest.runScript(sqlSessionFactory.getConfiguration().getEnvironment().getDataSource(),
+        "org/apache/ibatis/plugin/CreateDB.sql");
+  }
+
   @Test
-  void mapPluginShouldInterceptGet() {
-    Map map = new HashMap();
-    map = (Map) new AlwaysMapPlugin().plugin(map);
-    assertEquals("Always", map.get("Anything"));
+  void shouldPluginSwitchSchema() {
+    try (SqlSession sqlSession = sqlSessionFactory.openSession()) {
+      Mapper mapper = sqlSession.getMapper(Mapper.class);
+      assertEquals("Public user 1", mapper.selectNameById(1));
+    }
+
+    SchemaHolder.set("MYSCHEMA");
+
+    try (SqlSession sqlSession = sqlSessionFactory.openSession()) {
+      Mapper mapper = sqlSession.getMapper(Mapper.class);
+      assertEquals("Private user 1", mapper.selectNameById(1));
+    }
+  }
+
+  static class SchemaHolder {
+    private static ThreadLocal<String> value = ThreadLocal.withInitial(() -> "PUBLIC");
+
+    public static void set(String tenantName) {
+      value.set(tenantName);
+    }
+
+    public static String get() {
+      return value.get();
+    }
+  }
+
+  @Intercepts(@Signature(type = StatementHandler.class, method = "prepare", args = { Connection.class, Integer.class }))
+  public static class SwitchCatalogInterceptor implements Interceptor {
+    @Override
+    public Object intercept(Invocation invocation) throws Throwable {
+      Object[] args = invocation.getArgs();
+      Connection con = (Connection) args[0];
+      con.setSchema(SchemaHolder.get());
+      return invocation.proceed();
+    }
   }
 
   @Test
-  void shouldNotInterceptToString() {
-    Map map = new HashMap();
-    map = (Map) new AlwaysMapPlugin().plugin(map);
-    assertNotEquals("Always", map.toString());
+  void shouldPluginNotInvokeArbitraryMethod() {
+    Map<?, ?> map = new HashMap<>();
+    map = (Map<?, ?>) new AlwaysMapPlugin().plugin(map);
+    try {
+      map.get("Anything");
+      fail("Exected IllegalArgumentException, but no exception was thrown.");
+    } catch (IllegalArgumentException e) {
+      assertEquals(
+          "Method 'public abstract java.lang.Object java.util.Map.get(java.lang.Object)' is not supported as a plugin target.",
+          e.getMessage());
+    } catch (Exception e) {
+      fail("Exected IllegalArgumentException, but was " + e.getClass(), e);
+    }
   }
 
   @Intercepts({ @Signature(type = Map.class, method = "get", args = { Object.class }) })
@@ -45,7 +105,5 @@ public static class AlwaysMapPlugin implements Interceptor {
     public Object intercept(Invocation invocation) {
       return "Always";
     }
-
   }
-
 }

src/test/resources/org/apache/ibatis/plugin/CreateDB.sql

@@ -0,0 +1,35 @@
+--
+--    Copyright 2009-2024 the original author or authors.
+--
+--    Licensed under the Apache License, Version 2.0 (the "License");
+--    you may not use this file except in compliance with the License.
+--    You may obtain a copy of the License at
+--
+--       https://www.apache.org/licenses/LICENSE-2.0
+--
+--    Unless required by applicable law or agreed to in writing, software
+--    distributed under the License is distributed on an "AS IS" BASIS,
+--    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+--    See the License for the specific language governing permissions and
+--    limitations under the License.
+--
+
+drop schema public if exists; -- empty public remains
+
+create table public.users (
+  id int,
+  name varchar(20)
+);
+
+insert into public.users (id, name) values (1, 'Public user 1');
+
+drop schema myschema if exists;
+create schema myschema;
+
+create table myschema.users (
+  id int,
+  name varchar(20)
+);
+
+insert into myschema.users (id, name) values (1, 'Private user 1');
+

src/test/resources/org/apache/ibatis/plugin/mybatis-config.xml

@@ -0,0 +1,47 @@
+<?xml version="1.0" encoding="UTF-8" ?>
+<!--
+
+       Copyright 2009-2024 the original author or authors.
+
+       Licensed under the Apache License, Version 2.0 (the "License");
+       you may not use this file except in compliance with the License.
+       You may obtain a copy of the License at
+
+          https://www.apache.org/licenses/LICENSE-2.0
+
+       Unless required by applicable law or agreed to in writing, software
+       distributed under the License is distributed on an "AS IS" BASIS,
+       WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+       See the License for the specific language governing permissions and
+       limitations under the License.
+
+-->
+<!DOCTYPE configuration
+    PUBLIC "-//mybatis.org//DTD Config 3.0//EN"
+    "https://mybatis.org/dtd/mybatis-3-config.dtd">
+
+<configuration>
+
+  <plugins>
+    <plugin
+      interceptor="org.apache.ibatis.plugin.PluginTest$SwitchCatalogInterceptor" />
+  </plugins>
+
+  <environments default="development">
+    <environment id="development">
+      <transactionManager type="JDBC">
+        <property name="" value="" />
+      </transactionManager>
+      <dataSource type="UNPOOLED">
+        <property name="driver" value="org.hsqldb.jdbcDriver" />
+        <property name="url" value="jdbc:hsqldb:mem:plugintest" />
+        <property name="username" value="sa" />
+      </dataSource>
+    </environment>
+  </environments>
+
+  <mappers>
+    <mapper class="org.apache.ibatis.plugin.Mapper" />
+  </mappers>
+
+</configuration>