BUG#21658: Crash when creating table with item in prepared statement that allocates memory in fix_fields

We need to use an arena to indicate we are preparing a statement when loading partition function and
parsing it as part of an open table.

5 files changed, 59 insertions, 27 deletions

diff --git a/mysql-test/r/partition.result b/mysql-test/r/partition.result
index 94891f4f91f..495aec35f54 100644
--- a/mysql-test/r/partition.result
+++ b/mysql-test/r/partition.result
@@ -1,4 +1,9 @@
 drop table if exists t1;
+create table t1 (s1 char(2) character set utf8)
+partition by list (case when s1 > 'cz' then 1 else 2 end)
+(partition p1 values in (1),
+partition p2 values in (2));
+drop table t1;
 create table t1 (a int)
 partition by key(a)
 partitions 0.2+e1;
@@ -801,11 +806,6 @@ t2	CREATE TABLE `t2` (
   PRIMARY KEY (`a`)
 ) ENGINE=MyISAM DEFAULT CHARSET=latin1 COMMENT='no comment' /*!50100 PARTITION BY KEY (a)  */
 drop table t2;
-create table t1 (s1 char(2) character set utf8)
-partition by list (case when s1 > 'cz' then 1 else 2 end)
-(partition p1 values in (1),
-partition p2 values in (2));
-drop table t1;
 create table t1 (f1 int) partition by hash (f1) as select 1;
 drop table t1;
 prepare stmt1 from 'create table t1 (s1 int) partition by hash (s1)';
diff --git a/mysql-test/t/partition.test b/mysql-test/t/partition.test
index 43e57221a42..75aceb4fafe 100644
--- a/mysql-test/t/partition.test
+++ b/mysql-test/t/partition.test
@@ -10,6 +10,15 @@ drop table if exists t1;
 --enable_warnings
 
 #
+# Bug#14367: Partitions: crash if utf8 column
+#
+create table t1 (s1 char(2) character set utf8)
+partition by list (case when s1 > 'cz' then 1 else 2 end)
+(partition p1 values in (1),
+ partition p2 values in (2));
+drop table t1;
+
+#
 # Bug 15890: Strange number of partitions accepted
 #
 -- error 1064
@@ -940,15 +949,6 @@ show create table t2;
 drop table t2;
 
 #
-# Bug#14367: Partitions: crash if utf8 column
-#
-create table t1 (s1 char(2) character set utf8)
-partition by list (case when s1 > 'cz' then 1 else 2 end)
-(partition p1 values in (1),
- partition p2 values in (2));
-drop table t1;
-
-#
 # Bug#15336 Partitions: crash if create table as select
 #
 create table t1 (f1 int) partition by hash (f1) as select 1;
diff --git a/sql/item.cc b/sql/item.cc
index 6e661afbd91..0401646a41f 100644
--- a/sql/item.cc
+++ b/sql/item.cc
@@ -1492,7 +1492,7 @@ bool agg_item_charsets(DTCollation &coll, const char *fname,
       been created in prepare. In this case register the change for
       rollback.
     */
-    if (arena)
+    if (thd->is_stmt_prepare())
       *arg= conv;
     else
       thd->change_item_tree(arg, conv);
diff --git a/sql/sql_partition.cc b/sql/sql_partition.cc
index 208f9a0837a..f5a30994f7d 100644
--- a/sql/sql_partition.cc
+++ b/sql/sql_partition.cc
@@ -868,9 +868,13 @@ int check_signed_flag(partition_info *part_info)
 bool fix_fields_part_func(THD *thd, Item* func_expr, TABLE *table,
                           bool is_sub_part, bool is_field_to_be_setup)
 {
+  MEM_ROOT new_mem_root;
+  Query_arena partition_arena(&new_mem_root, Query_arena::INITIALIZED);
+  Query_arena backup_arena;
   partition_info *part_info= table->part_info;
   uint dir_length, home_dir_length;
   bool result= TRUE;
+  bool is_prepare;
   TABLE_LIST tables;
   TABLE_LIST *save_table_list, *save_first_table, *save_last_table;
   int error;
@@ -917,7 +921,25 @@ bool fix_fields_part_func(THD *thd, Item* func_expr, TABLE *table,
   func_expr->walk(&Item::change_context_processor, 0, (byte*) context);
   save_where= thd->where;
   thd->where= "partition function";
+  /*
+    In execution we must avoid the use of thd->change_item_tree since
+    we might release memory before statement is completed. We do this
+    by temporarily setting the stmt_arena->mem_root to be the mem_root
+    of the table object, this also ensures that any memory allocated
+    during fix_fields will not be released at end of execution of this
+    statement. Thus the item tree will remain valid also in subsequent
+    executions of this table object. We do however not at the moment
+    support allocations during execution of val_int so any item class
+    that does this during val_int must be disallowed as partition
+    function.
+    SEE Bug #21658
+  */
+  /*
+    This is a tricky call to prepare for since it can have a large number
+    of interesting side effects, both desirable and undesirable.
+  */
   error= func_expr->fix_fields(thd, (Item**)0);
+
   context->table_list= save_table_list;
   context->first_name_resolution_table= save_first_table;
   context->last_name_resolution_table= save_last_table;
@@ -1422,7 +1444,6 @@ bool fix_partition_func(THD *thd, TABLE *table,
       DBUG_RETURN(TRUE);
     }
   }
-  thd->free_list= part_info->item_free_list;
   if (part_info->is_sub_partitioned())
   {
     DBUG_ASSERT(part_info->subpart_type == HASH_PARTITION);
@@ -1530,7 +1551,6 @@ bool fix_partition_func(THD *thd, TABLE *table,
   set_up_range_analysis_info(part_info);
   result= FALSE;
 end:
-  thd->free_list= thd_free_list;
   thd->mark_used_columns= save_mark_used_columns;
   DBUG_PRINT("info", ("thd->mark_used_columns: %d", thd->mark_used_columns));
   DBUG_RETURN(result);
@@ -3368,7 +3388,6 @@ bool mysql_unpack_partition(THD *thd, const uchar *part_buf,
                             TABLE* table, bool is_create_table_ind,
                             handlerton *default_db_type)
 {
-  Item *thd_free_list= thd->free_list;
   bool result= TRUE;
   partition_info *part_info;
   CHARSET_INFO *old_character_set_client= thd->variables.character_set_client;
@@ -3396,7 +3415,6 @@ bool mysql_unpack_partition(THD *thd, const uchar *part_buf,
     Thus we move away the current list temporarily and start a new list that
     we then save in the partition info structure.
   */
-  thd->free_list= NULL;
   lex.part_info= new partition_info();/* Indicates MYSQLparse from this place */
   if (!lex.part_info)
   {
@@ -3409,6 +3427,7 @@ bool mysql_unpack_partition(THD *thd, const uchar *part_buf,
   if (MYSQLparse((void*)thd) || thd->is_fatal_error)
   {
     free_items(thd->free_list);
+    thd->free_list= NULL;
     goto end;
   }
   /*
@@ -3477,7 +3496,6 @@ bool mysql_unpack_partition(THD *thd, const uchar *part_buf,
   if (!part_info->default_engine_type)
     part_info->default_engine_type= default_db_type;
   DBUG_ASSERT(part_info->default_engine_type == default_db_type);
-  part_info->item_free_list= thd->free_list;
 
   {
   /*
@@ -3500,7 +3518,7 @@ bool mysql_unpack_partition(THD *thd, const uchar *part_buf,
     {
       mem_alloc_error(part_func_len);
       free_items(thd->free_list);
-      part_info->item_free_list= 0;
+      thd->free_list= NULL;
       goto end;
     }
     if (part_func_len)
@@ -3515,7 +3533,6 @@ bool mysql_unpack_partition(THD *thd, const uchar *part_buf,
   result= FALSE;
 end:
   lex_end(thd->lex);
-  thd->free_list= thd_free_list;
   thd->lex= old_lex;
   thd->variables.character_set_client= old_character_set_client;
   DBUG_RETURN(result);
diff --git a/sql/table.cc b/sql/table.cc
index dec4901807d..bafd254ea51 100644
--- a/sql/table.cc
+++ b/sql/table.cc
@@ -1475,11 +1475,23 @@ int open_table_from_share(THD *thd, TABLE_SHARE *share, const char *alias,
 #ifdef WITH_PARTITION_STORAGE_ENGINE
   if (share->partition_info_len)
   {
-    MEM_ROOT **root_ptr, *old_root;
+  /*
+    In this execution we must avoid calling thd->change_item_tree since
+    we might release memory before statement is completed. We do this
+    by changing to a new statement arena. As part of this arena we also
+    set the memory root to be the memory root of the table since we
+    call the parser and fix_fields which both can allocate memory for
+    item objects. We keep the arena to ensure that we can release the
+    free_list when closing the table object.
+    SEE Bug #21658
+  */
+
+    Query_arena *backup_stmt_arena_ptr= thd->stmt_arena;
+    Query_arena backup_arena;
+    Query_arena part_func_arena(&outparam->mem_root, Query_arena::INITIALIZED);
+    thd->set_n_backup_active_arena(&part_func_arena, &backup_arena);
+    thd->stmt_arena= &part_func_arena;
     bool tmp;
-    root_ptr= my_pthread_getspecific_ptr(MEM_ROOT**, THR_MALLOC);
-    old_root= *root_ptr;
-    *root_ptr= &outparam->mem_root;
 
     tmp= mysql_unpack_partition(thd, share->partition_info,
                                 share->partition_info_len,
@@ -1491,7 +1503,10 @@ int open_table_from_share(THD *thd, TABLE_SHARE *share, const char *alias,
     DBUG_PRINT("info", ("autopartitioned: %u", share->auto_partitioned));
     if (!tmp)
       tmp= fix_partition_func(thd, outparam, is_create_table);
-    *root_ptr= old_root;
+    if (!tmp)
+      outparam->part_info->item_free_list= part_func_arena.free_list;
+    thd->stmt_arena= backup_stmt_arena_ptr;
+    thd->restore_active_arena(&part_func_arena, &backup_arena);
     if (tmp)
     {
       if (is_create_table)