[UTC][nntrainer][Non-ACR] Assgin NULL pointer to handle 35/284935/1
authorhyeonseok lee <hs89.lee@samsung.com>
Wed, 30 Nov 2022 07:10:02 +0000 (16:10 +0900)
committerhyeonseok lee <hs89.lee@samsung.com>
Thu, 1 Dec 2022 08:42:33 +0000 (17:42 +0900)
 - For now the layer/optimizer/dataset handle which is used by model is cleanup twice.
   One by in model_destroy and the other one by in model_cleanup.
   This caused heap use after free issue.
 - To prohibit heap use after free assign NULL pointer to handle if it is used by model.
   By checking the handle in cleanup function it will protect from heap use after free.

Change-Id: Iac1b9b200c7eff8363b21f1ec0a109eb03c86b0e
Signed-off-by: hyeonseok lee <hs89.lee@samsung.com>
src/utc/nntrainer/utc-nntrainer-model.c

index 1eea41f163aff1336149bbd7ea9dac4630b16405..0d0c78703e0bc1ae70dc8723a3885d10d9a8ed7b 100644 (file)
@@ -65,11 +65,20 @@ void utc_nntrainer_model_startup(void)
  */
 void utc_nntrainer_model_cleanup(void)
 {
-  ml_train_layer_destroy(layers[0]);
-  ml_train_layer_destroy(layers[1]);
+  if (layers[0]) {
+    ml_train_layer_destroy(layers[0]);
+  }
+  if (layers[1]) {
+    ml_train_layer_destroy(layers[1]);
+  }
 
-  ml_train_optimizer_destroy(adam);
-  ml_train_dataset_destroy(dataset);
+  if (adam) {
+    ml_train_optimizer_destroy(adam);
+  }
+
+  if (dataset) {
+    ml_train_dataset_destroy(dataset);
+  }
   return;
 }
 
@@ -219,6 +228,7 @@ int utc_nntrainer_model_add_layer_p(void)
   }
 
   ml_train_model_destroy(model);
+  layers[0] = NULL;
 
   return 0;
 }
@@ -245,6 +255,7 @@ int utc_nntrainer_model_add_layer_n1(void)
   }
 
   ml_train_model_destroy(model);
+  layers[1] = NULL;
 
   return 0;
 }
@@ -272,6 +283,7 @@ int utc_nntrainer_model_add_layer_n2(void)
 
   ml_train_model_destroy(model);
   ml_train_model_destroy(other_model);
+  layers[1] = NULL;
 
   return 0;
 }
@@ -295,6 +307,7 @@ int utc_nntrainer_model_set_optimizer_p(void)
   }
 
   ml_train_model_destroy(model);
+  adam = NULL;
 
   return 0;
 }
@@ -324,6 +337,7 @@ int utc_nntrainer_model_set_optimizer_n1(void)
 
   ml_train_model_destroy(model);
   ml_train_model_destroy(other_model);
+  adam = NULL;
 
   return 0;
 }
@@ -372,6 +386,7 @@ int utc_nntrainer_model_set_dataset_p(void)
   }
 
   ml_train_model_destroy(model);
+  dataset = NULL;
 
   return 0;
 }
@@ -401,6 +416,7 @@ int utc_nntrainer_model_set_dataset_n1(void)
 
   ml_train_model_destroy(model);
   ml_train_model_destroy(other_model);
+  dataset = NULL;
 
   return 0;
 }
@@ -455,6 +471,7 @@ int utc_nntrainer_model_get_summary_p(void)
     free(sum);
 
   ml_train_model_destroy(model);
+  layers[0] = NULL;
 
   return 0;
 }
@@ -576,7 +593,6 @@ int utc_nntrainer_model_run_p1(void)
 int utc_nntrainer_model_run_p2(void)
 {
   ml_train_dataset_h dataset;
-  ml_train_optimizer_h optimizer;
 
   ml_train_model_construct(&model);
   ml_train_dataset_create_with_generator(&dataset, constant_generator_cb, NULL, NULL);
@@ -584,12 +600,7 @@ int utc_nntrainer_model_run_p2(void)
   ml_train_dataset_set_property(dataset, "buffer_size=9", NULL);
   ml_train_model_set_dataset(model, dataset);
 
-  ml_train_optimizer_create(&optimizer, ML_TRAIN_OPTIMIZER_TYPE_ADAM);
-  ml_train_optimizer_set_property(
-      optimizer, "learning_rate=0.0001", "decay_rate=0.96", "decay_steps=1000",
-      "beta1=0.9", "beta2=0.9999", "epsilon=1e-7", NULL);
-
-  ml_train_model_set_optimizer(model, optimizer);
+  ml_train_model_set_optimizer(model, adam);
 
   ml_train_model_add_layer(model, layers[0]);
   ml_train_model_add_layer(model, layers[1]);
@@ -607,6 +618,9 @@ int utc_nntrainer_model_run_p2(void)
   }
 
   ml_train_model_destroy(model);
+  layers[0] = NULL;
+  layers[1] = NULL;
+  adam = NULL;
 
   return 0;
 }
@@ -929,6 +943,8 @@ int utc_nntrainer_model_get_layer_p1(void)
     assert_eq(status, ML_ERROR_NONE);
   }
 
+  layers[0] = NULL;
+
   return 0;
 }
 
@@ -982,6 +998,8 @@ int utc_nntrainer_model_get_layer_n1(void)
     assert_eq(status, ML_ERROR_NONE);
   }
 
+  layers[0] = NULL;
+
   return 0;
 }