xtest: fix against uninited session

Client library crashes if we try to work with unitialized session.
This patch adds checks on values returned by TEE_OpenSession() to
avoid segfaults.

Also patch includes minor refactoring to make affected code more clear.

Signed-off-by: Volodymyr Babchuk <vlad.babchuk@gmail.com>
Reviewed-by: Etienne Carriere <etienne.carriere@linaro.org>
Reviewed-by: Igor Opaniuk <igor.opaniuk@linaro.org>
Tested-by: Jerome Forissier <jerome.forissier@linaro.org> (HiKey)
diff --git a/host/xtest/xtest_1000.c b/host/xtest/xtest_1000.c
index bd3ddb1..f46fa00 100644
--- a/host/xtest/xtest_1000.c
+++ b/host/xtest/xtest_1000.c
@@ -411,9 +411,10 @@
 	TEEC_Operation op = TEEC_OPERATION_INITIALIZER;
 	uint32_t ret_orig;
 
-	(void)ADBG_EXPECT_TEEC_SUCCESS(c,
+	if (!ADBG_EXPECT_TEEC_SUCCESS(c,
 		xtest_teec_open_session(&session, &os_test_ta_uuid, NULL,
-					&ret_orig));
+		                        &ret_orig)))
+		return;
 
 	op.params[0].value.a = n;
 	op.paramTypes = TEEC_PARAM_TYPES(TEEC_VALUE_INPUT, TEEC_NONE, TEEC_NONE,
@@ -424,11 +425,11 @@
 		TEEC_InvokeCommand(&session, TA_OS_TEST_CMD_BAD_MEM_ACCESS, &op,
 				   &ret_orig));
 
-	(void)ADBG_EXPECT_TEEC_RESULT(c, TEEC_ERROR_TARGET_DEAD,
-				      TEEC_InvokeCommand(&session,
-					TA_OS_TEST_CMD_BAD_MEM_ACCESS,
-					&op,
+	(void)ADBG_EXPECT_TEEC_RESULT(c,
+	        TEEC_ERROR_TARGET_DEAD,
+	        TEEC_InvokeCommand(&session, TA_OS_TEST_CMD_BAD_MEM_ACCESS, &op,
 					&ret_orig));
+
 	(void)ADBG_EXPECT_TEEC_ERROR_ORIGIN(c, TEEC_ORIGIN_TEE, ret_orig);
 
 	TEEC_CloseSession(&session);
@@ -482,9 +483,10 @@
 	TEEC_Session session = { 0 };
 	uint32_t ret_orig;
 
-	(void)ADBG_EXPECT_TEEC_SUCCESS(c,
+	if (!ADBG_EXPECT_TEEC_SUCCESS(c,
 		xtest_teec_open_session(&session, &os_test_ta_uuid, NULL,
-					&ret_orig));
+		                        &ret_orig)))
+		return;
 
 	(void)ADBG_EXPECT_TEEC_RESULT(c,
 		TEEC_ERROR_TARGET_DEAD,
@@ -668,15 +670,16 @@
 
 	Do_ADBG_BeginSubCase(c, "Invoke command");
 	{
-		(void)ADBG_EXPECT_TEEC_SUCCESS(c,
+		if (ADBG_EXPECT_TEEC_SUCCESS(c,
 			xtest_teec_open_session(&session, &os_test_ta_uuid,
-						NULL, &ret_orig));
+			                        NULL, &ret_orig))) {
 
-		(void)ADBG_EXPECT_TEEC_SUCCESS(c,
-			TEEC_InvokeCommand(&session, TA_OS_TEST_CMD_CLIENT,
-					   NULL, &ret_orig));
+			(void)ADBG_EXPECT_TEEC_SUCCESS(c,
+			    TEEC_InvokeCommand(&session, TA_OS_TEST_CMD_CLIENT,
+			                       NULL, &ret_orig));
+			TEEC_CloseSession(&session);
+		}
 
-		TEEC_CloseSession(&session);
 	}
 	Do_ADBG_EndSubCase(c, "Invoke command");
 
@@ -688,18 +691,18 @@
 		op.paramTypes = TEEC_PARAM_TYPES(
 			TEEC_VALUE_INPUT, TEEC_NONE, TEEC_NONE, TEEC_NONE);
 
-		(void)ADBG_EXPECT_TEEC_SUCCESS(c,
+		if (ADBG_EXPECT_TEEC_SUCCESS(c,
 			xtest_teec_open_session(&session,
 						&os_test_ta_uuid,
 						NULL,
-						&ret_orig));
+			                        &ret_orig))) {
 
-		(void)ADBG_EXPECT_TEEC_SUCCESS(c,
-			TEEC_InvokeCommand(&session,
-					   TA_OS_TEST_CMD_CLIENT_WITH_TIMEOUT,
-					   &op, &ret_orig));
-
-		TEEC_CloseSession(&session);
+			(void)ADBG_EXPECT_TEEC_SUCCESS(c,
+			  TEEC_InvokeCommand(&session,
+			                     TA_OS_TEST_CMD_CLIENT_WITH_TIMEOUT,
+			                     &op, &ret_orig));
+			TEEC_CloseSession(&session);
+		}
 	}
 	Do_ADBG_EndSubCase(c, "Invoke command with timeout");
 
@@ -765,109 +768,70 @@
 }
 #endif
 
-static void xtest_tee_test_1009(ADBG_Case_t *c)
+static void xtest_tee_test_1009_subcase(ADBG_Case_t *c, const char *subcase,
+                                        uint32_t timeout, bool cancel)
 {
 	TEEC_Session session = { 0 };
 	uint32_t ret_orig;
-
-	Do_ADBG_BeginSubCase(c, "TEE Wait 0.1s");
-	{
-		TEEC_Operation op = TEEC_OPERATION_INITIALIZER;
-
-		(void)ADBG_EXPECT_TEEC_SUCCESS(c,
-			xtest_teec_open_session(&session, &os_test_ta_uuid,
-						NULL, &ret_orig));
-
-		(void)ADBG_EXPECT_TEEC_ERROR_ORIGIN(c, TEEC_ORIGIN_TRUSTED_APP,
-						    ret_orig);
-
-		op.params[0].value.a = 100;
-		op.paramTypes = TEEC_PARAM_TYPES(TEEC_VALUE_INPUT, TEEC_NONE,
-						 TEEC_NONE, TEEC_NONE);
-
-		(void)ADBG_EXPECT_TEEC_SUCCESS(c,
-			TEEC_InvokeCommand(&session, TA_OS_TEST_CMD_WAIT, &op,
-					   &ret_orig));
-		TEEC_CloseSession(&session);
-	}
-	Do_ADBG_EndSubCase(c, "TEE Wait 0.1s");
-
-	Do_ADBG_BeginSubCase(c, "TEE Wait 0.5s");
-	{
-		TEEC_Operation op = TEEC_OPERATION_INITIALIZER;
-
-		(void)ADBG_EXPECT_TEEC_SUCCESS(c,
-			xtest_teec_open_session(&session, &os_test_ta_uuid,
-						NULL, &ret_orig));
-
-		(void)ADBG_EXPECT_TEEC_ERROR_ORIGIN(c, TEEC_ORIGIN_TRUSTED_APP,
-						    ret_orig);
-
-		op.params[0].value.a = 500;
-		op.paramTypes = TEEC_PARAM_TYPES(TEEC_VALUE_INPUT, TEEC_NONE,
-						 TEEC_NONE, TEEC_NONE);
-
-		(void)ADBG_EXPECT_TEEC_SUCCESS(c,
-			TEEC_InvokeCommand(&session, TA_OS_TEST_CMD_WAIT, &op,
-					   &ret_orig));
-
-		(void)ADBG_EXPECT_TEEC_ERROR_ORIGIN(c, TEEC_ORIGIN_TRUSTED_APP,
-						    ret_orig);
-		TEEC_CloseSession(&session);
-	}
-	Do_ADBG_EndSubCase(c, "TEE Wait 0.5s");
-
 #ifdef USER_SPACE
-	Do_ADBG_BeginSubCase(c, "TEE Wait 2s cancel");
-	{
-		pthread_t thr;
-		TEEC_Operation op = TEEC_OPERATION_INITIALIZER;
-
-		(void)ADBG_EXPECT_TEEC_SUCCESS(c,
-			xtest_teec_open_session(&session, &os_test_ta_uuid,
-						NULL, &ret_orig));
-
-		(void)ADBG_EXPECT_TEEC_ERROR_ORIGIN(c, TEEC_ORIGIN_TRUSTED_APP,
-						    ret_orig);
-
-		op.params[0].value.a = 2000;
-		op.paramTypes = TEEC_PARAM_TYPES(TEEC_VALUE_INPUT, TEEC_NONE,
-						 TEEC_NONE, TEEC_NONE);
-
-		(void)ADBG_EXPECT(c, 0,
-			pthread_create(&thr, NULL, cancellation_thread, &op));
-
-		(void)ADBG_EXPECT_TEEC_RESULT(c, TEEC_ERROR_CANCEL,
-			TEEC_InvokeCommand(&session, TA_OS_TEST_CMD_WAIT, &op,
-					   &ret_orig));
-
-		(void)ADBG_EXPECT_TEEC_ERROR_ORIGIN(c, TEEC_ORIGIN_TRUSTED_APP,
-						    ret_orig);
-		(void)ADBG_EXPECT(c, 0, pthread_join(thr, NULL));
-		TEEC_CloseSession(&session);
-	}
-	Do_ADBG_EndSubCase(c, "TEE Wait 2s cancel");
+	pthread_t thr;
 #endif
 
-	Do_ADBG_BeginSubCase(c, "TEE Wait 2s");
+	Do_ADBG_BeginSubCase(c, "%s", subcase);
 	{
 		TEEC_Operation op = TEEC_OPERATION_INITIALIZER;
 
-		(void)ADBG_EXPECT_TEEC_SUCCESS(c,
-			xtest_teec_open_session(&session, &os_test_ta_uuid,
-						NULL, &ret_orig));
+		if (ADBG_EXPECT_TEEC_SUCCESS(c,
+		         xtest_teec_open_session(&session, &os_test_ta_uuid,
+		                                 NULL, &ret_orig))) {
 
-		op.params[0].value.a = 2000;
-		op.paramTypes = TEEC_PARAM_TYPES(TEEC_VALUE_INPUT, TEEC_NONE,
-						 TEEC_NONE, TEEC_NONE);
+			(void)ADBG_EXPECT_TEEC_ERROR_ORIGIN(c,
+			           TEEC_ORIGIN_TRUSTED_APP,
+			           ret_orig);
 
-		(void)ADBG_EXPECT_TEEC_SUCCESS(c,
-			TEEC_InvokeCommand(&session, TA_OS_TEST_CMD_WAIT, &op,
-					   &ret_orig));
+			op.params[0].value.a = timeout;
+			op.paramTypes = TEEC_PARAM_TYPES(TEEC_VALUE_INPUT,
+			                                 TEEC_NONE,
+			                                 TEEC_NONE, TEEC_NONE);
+#ifdef USER_SPACE
+			if (cancel) {
+				(void)ADBG_EXPECT(c, 0,
+				      pthread_create(&thr, NULL,
+				                     cancellation_thread, &op));
 
-		TEEC_CloseSession(&session);
+				(void)ADBG_EXPECT_TEEC_RESULT(c,
+				           TEEC_ERROR_CANCEL,
+			                   TEEC_InvokeCommand(&session,
+			                                 TA_OS_TEST_CMD_WAIT,
+			                                  &op,
+			                                  &ret_orig));
+			} else
+#endif
+
+			(void)ADBG_EXPECT_TEEC_SUCCESS(c,
+			           TEEC_InvokeCommand(&session,
+			                              TA_OS_TEST_CMD_WAIT,
+			                              &op,
+			                              &ret_orig));
+#ifdef USER_SPACE
+			if (cancel)
+				(void)ADBG_EXPECT(c, 0, pthread_join(thr, NULL));
+#endif
+
+			TEEC_CloseSession(&session);
+		}
 	}
-	Do_ADBG_EndSubCase(c, "TEE Wait 2s");
+	Do_ADBG_EndSubCase(c, "%s", subcase);
+}
+
+static void xtest_tee_test_1009(ADBG_Case_t *c)
+{
+	xtest_tee_test_1009_subcase(c, "TEE Wait 0.1s", 100, false);
+	xtest_tee_test_1009_subcase(c, "TEE Wait 0.5s", 500, false);
+#ifdef USER_SPACE
+	xtest_tee_test_1009_subcase(c, "TEE Wait 2s cancel", 2000, true);
+#endif
+	xtest_tee_test_1009_subcase(c, "TEE Wait 2s", 2000, false);
 }
 
 static void xtest_tee_test_1010(ADBG_Case_t *c)
@@ -929,9 +893,10 @@
 		uint8_t out[32] = { 0 };
 		int i;
 
-		(void)ADBG_EXPECT_TEEC_SUCCESS(c,
+		if (!ADBG_EXPECT_TEEC_SUCCESS(c,
 			xtest_teec_open_session(&session1, &uuid, NULL,
-						&ret_orig));
+			                        &ret_orig)))
+			return;
 
 		op.params[0].value.a = 0;
 		op.params[1].tmpref.buffer = (void *)in;
@@ -945,9 +910,10 @@
 					   &ret_orig));
 
 		for (i = 1; i < 1000; i++) {
-			(void)ADBG_EXPECT_TEEC_SUCCESS(c,
+			if (!ADBG_EXPECT_TEEC_SUCCESS(c,
 				xtest_teec_open_session(&session2, &uuid, NULL,
-							&ret_orig));
+				                        &ret_orig)))
+				continue;
 
 			op.params[0].value.a = 0;
 			op.params[1].tmpref.buffer = out;
diff --git a/host/xtest/xtest_20000.c b/host/xtest/xtest_20000.c
index 3efb4af..d7671c4 100644
--- a/host/xtest/xtest_20000.c
+++ b/host/xtest/xtest_20000.c
@@ -424,14 +424,17 @@
 	uint32_t obj_id;
 	uint32_t nb;
 	size_t n;
+	char *filedata = NULL;
 
 
-	ADBG_EXPECT(c, TEE_SUCCESS,
-		    TEEC_InitializeContext(_device, &ctx));
+	if (!ADBG_EXPECT(c, TEE_SUCCESS,
+	            TEEC_InitializeContext(_device, &ctx)))
+		return;
 
-	ADBG_EXPECT(c, TEE_SUCCESS,
+	if (!ADBG_EXPECT(c, TEE_SUCCESS,
 		    TEEC_OpenSession(&ctx, &sess, &uuid,
-			TEEC_LOGIN_PUBLIC, NULL, NULL, &error));
+		        TEEC_LOGIN_PUBLIC, NULL, NULL, &error)))
+		return;
 
 	for (n = 0; n < ARRAY_SIZE(xtest_enc_fs_cases); n++) {
 
@@ -442,7 +445,6 @@
 
 		char buffer[tv->data_len];
 		char filename[20];
-		char *filedata = NULL;
 		size_t p;
 		uint8_t data_byte = 0;
 
@@ -461,13 +463,15 @@
 		Do_ADBG_BeginSubCase(c, "| filename: %s , data size: %d byte(s)",
 				     filename, tv->data_len);
 
-		ADBG_EXPECT(c, TEE_SUCCESS,
+		if (ADBG_EXPECT(c, TEE_SUCCESS,
 			    obj_create(&sess, filename, ARRAY_SIZE(filename),
 				TEE_DATA_FLAG_ACCESS_WRITE, 0, filedata,
-				tv->data_len, &obj_id));
+			        tv->data_len, &obj_id)))
 
-		ADBG_EXPECT(c, TEE_SUCCESS,
-			    obj_close(&sess, obj_id));
+			ADBG_EXPECT(c, TEE_SUCCESS,
+			            obj_close(&sess, obj_id));
+		else
+			goto exit;
 
 		if (file_type == META_FILE)
 			ADBG_EXPECT_COMPARE_UNSIGNED(c, 0, !=,
@@ -571,10 +575,12 @@
 		}
 
 		free(filedata);
+		filedata = NULL;
 		Do_ADBG_EndSubCase(c, NULL);
 	};
 
 exit:
+	free(filedata);
 	TEEC_CloseSession(&sess);
 	TEEC_FinalizeContext(&ctx);
 }