Modify error handling codes to free resources

In the error handling codes, allocated resources were not freed
correctly.
To fix this, resource free codes like fclose and free are added.
In addition, memory allocation failure checking code is added right
after the malloc.
Also, 'return 0;' is added to the main function of postlinker_64.c since
it did not have it.

Change-Id: I9ea6dc96267c7760596b9260ef9f0d27f1a79b19
Signed-off-by: SangYoun Kwak <sy.kwak@samsung.com>

2 files changed, 198 insertions, 49 deletions

diff --git a/script/postlinker.c b/script/postlinker.c
index b7a5ee2..a25da2f 100755
--- a/script/postlinker.c
+++ b/script/postlinker.c
@@ -84,9 +84,11 @@ typedef struct {
 int PostLinker(char *fileName)
 {
 	Elf32_Ehdr elfHeader;
-	Elf32_Phdr *pSegmentEntries;
-	Elf32_Dyn *DynEntries;
-	Elf32_Sym *pMapSymbolTable;
+	Elf32_Phdr *pSegmentEntries = NULL;
+	Elf32_Dyn *DynEntries = NULL;
+	Elf32_Sym *pMapSymbolTable = NULL;
+
+	int exit_status = EXIT_SUCCESS;
 	int totalSymbolCnt;
 	int pos;
 
@@ -100,50 +102,76 @@ int PostLinker(char *fileName)
 
 	if (fp == NULL) {
 		printf("error[%s] : file open\n", fileName);
-		exit(EXIT_FAILURE);
+		exit_status = EXIT_FAILURE;
+		goto finalize;
 	}
 
 	/* Get ELF Header Information */
 	if (!fread((void*)&elfHeader, sizeof(Elf32_Ehdr), 1, fp)) {
 		printf("error[%s] : file read <header>\n", fileName);
-		exit(EXIT_FAILURE);
+		exit_status = EXIT_FAILURE;
+		goto finalize;
 	}
 
 	if (*(int*)(elfHeader.e_ident) != ELF_FORMAT) {
 		printf("error[%s] : no ELF format\n", fileName);
-		exit(EXIT_FAILURE);
+		exit_status = EXIT_FAILURE;
+		goto finalize;
 	}
-	pSegmentEntries = NULL;
-	DynEntries = NULL;
 
 	/* Memory Allocation for Segment and Section entries */
 	pSegmentEntries = (Elf32_Phdr *)malloc(elfHeader.e_phentsize * elfHeader.e_phnum);
+	if (pSegmentEntries == NULL) {
+		printf("error[%s] : memory allocation <pSegmentEntries>\n", fileName);
+		exit_status = EXIT_FAILURE;
+		goto finalize;
+	}
 
 	/* Get Segment Header Information using ELF Header */
-	fseek(fp, elfHeader.e_phoff, SEEK_SET);
-	if (!fread(pSegmentEntries, elfHeader.e_phentsize * elfHeader.e_phnum, 1, fp))
+	if (fseek(fp, elfHeader.e_phoff, SEEK_SET) != 0) {
+		printf("error[%s] : fseek <pSegmentEntries>\n", fileName);
+		exit_status = EXIT_FAILURE;
+		goto finalize;
+	}
+
+	if (!fread(pSegmentEntries, elfHeader.e_phentsize * elfHeader.e_phnum, 1, fp)) {
 		printf("error[%s] : file read <pSegmentEntries>\n", fileName);
+		exit_status = EXIT_FAILURE;
+		goto finalize;
+	}
 
 	/* Get Dynamic segment Information using Segment Header */
 	for (int i = 0; i < elfHeader.e_phnum; ++i) {
 		if (pSegmentEntries[i].p_type == PT_DYNAMIC) {
-			fseek(fp, pSegmentEntries[i].p_offset, SEEK_SET);
-			DynEntries = (Elf32_Dyn*)malloc(pSegmentEntries[i].p_memsz);
+			if (fseek(fp, pSegmentEntries[i].p_offset, SEEK_SET) != 0) {
+				printf("error[%s] : fseek <DynEntries>\n", fileName);
+				exit_status = EXIT_FAILURE;
+				goto finalize;
+			}
 
+			DynEntries = (Elf32_Dyn*)malloc(pSegmentEntries[i].p_memsz);
 			if (DynEntries == NULL) {
 				printf("error[%s] : memory allocation <DynEntries>\n", fileName);
-				exit(EXIT_FAILURE);
+				exit_status = EXIT_FAILURE;
+				goto finalize;
 			}
 
 			if (!fread(DynEntries, pSegmentEntries[i].p_memsz, 1, fp)) {
 				printf("error[%s] : file read <DynEntries>\n", fileName);
-				exit(EXIT_FAILURE);
+				exit_status = EXIT_FAILURE;
+				goto finalize;
 			}
 			pos = i;
 			break;
 		}
 	}
 
+	if (DynEntries == NULL) {
+		printf("error[%s] : no such p_type <PT_DYNAMIC>\n", fileName);
+		exit_status = EXIT_FAILURE;
+		goto finalize;
+	}
+
 	for (int i = 0; DynEntries[i].d_tag != DT_NULL; ++i) {
 		switch (DynEntries[i].d_tag) {
 		case DT_NEEDED:
@@ -166,14 +194,36 @@ int PostLinker(char *fileName)
 		}
 	}
 
-	fseek(fp, hashOff + 4, SEEK_SET);
-	if (!fread(&totalSymbolCnt, sizeof(Elf32_Word), 1, fp))
+	if (fseek(fp, hashOff + 4, SEEK_SET) != 0) {
+		printf("error[%s] : fseek <symEntryCnt>\n", fileName);
+		exit_status = EXIT_FAILURE;
+		goto finalize;
+	}
+
+	if (!fread(&totalSymbolCnt, sizeof(Elf32_Word), 1, fp)) {
 		printf("error[%s] : file read <symEntryCnt>\n", fileName);
+		exit_status = EXIT_FAILURE;
+		goto finalize;
+	}
 
 	pMapSymbolTable = (Elf32_Sym*)malloc(totalSymbolCnt * sizeof(Elf32_Sym));
-	fseek(fp, symbolOff, SEEK_SET);
-	if (!fread(pMapSymbolTable, symEntrySize * totalSymbolCnt, 1, fp))
+	if (pMapSymbolTable == NULL) {
+		printf("error[%s] : memory allocation <pMapSymbolTable>\n", fileName);
+		exit_status = EXIT_FAILURE;
+		goto finalize;
+	}
+
+	if (fseek(fp, symbolOff, SEEK_SET) != 0) {
+		printf("error[%s] : fseek <pMapSymbolTable>\n", fileName);
+		exit_status = EXIT_FAILURE;
+		goto finalize;
+	}
+
+	if (!fread(pMapSymbolTable, symEntrySize * totalSymbolCnt, 1, fp)) {
 		printf("error[%s] : file read <pMapSymbolTable>\n", fileName);
+		exit_status = EXIT_FAILURE;
+		goto finalize;
+	}
 
 	for (int i = 1; i < totalSymbolCnt; ++i) {
 		if (pMapSymbolTable[i].st_shndx == SHN_UNDEF &&
@@ -192,23 +242,46 @@ int PostLinker(char *fileName)
 
 	if (!done) {
 		printf("Undefined symbols are already removed : %s\n", fileName);
-		return 1;
+		exit_status = EXIT_SUCCESS;
+		goto finalize;
+	}
+
+	if (fseek(fp, pSegmentEntries[pos].p_offset, SEEK_SET) != 0) {
+		printf("error[%s] : fseek <DynEntries>\n", fileName);
+		exit_status = EXIT_FAILURE;
+		goto finalize;
 	}
 
-	fseek(fp, pSegmentEntries[pos].p_offset, SEEK_SET);
 	if (!fwrite(DynEntries, pSegmentEntries[pos].p_memsz, 1, fp)) {
 		printf("error[%s] : file write <DynEntries>\n", fileName);
-		exit(EXIT_FAILURE);
+		exit_status = EXIT_FAILURE;
+		goto finalize;
 	}
-	fseek(fp, symbolOff, SEEK_SET);
+
+	if (fseek(fp, symbolOff, SEEK_SET) != 0) {
+		printf("error[%s] : fseek <pMapSymbolTable>\n", fileName);
+		exit_status = EXIT_FAILURE;
+		goto finalize;
+	}
+
 	if (!fwrite(pMapSymbolTable, symEntrySize * totalSymbolCnt, 1, fp)) {
 		printf("error[%s] : file write <pMapSymbolTable>\n", fileName);
-		exit(EXIT_FAILURE);
+		exit_status = EXIT_FAILURE;
+		goto finalize;
 	}
 
-	fclose(fp);
 	printf("Success : %s\n", fileName);
 
+finalize:
+	free(pMapSymbolTable);
+	free(DynEntries);
+	free(pSegmentEntries);
+	if (fp != NULL)
+		fclose(fp);
+
+	if (exit_status == EXIT_FAILURE)
+		exit(EXIT_FAILURE);
+
 	return 1;
 }
 
diff --git a/script/postlinker_64.c b/script/postlinker_64.c
index 994cae7..0764fde 100755
--- a/script/postlinker_64.c
+++ b/script/postlinker_64.c
@@ -10,9 +10,11 @@
 int PostLinker64(char *fileName)
 {
 	Elf64_Ehdr elfHeader;
-	Elf64_Phdr *pSegmentEntries;
-	Elf64_Dyn *DynEntries;
-	Elf64_Sym *pMapSymbolTable;
+	Elf64_Phdr *pSegmentEntries = NULL;
+	Elf64_Dyn *DynEntries = NULL;
+	Elf64_Sym *pMapSymbolTable = NULL;
+
+	int exit_status = EXIT_SUCCESS;
 	int totalSymbolCnt;
 	int pos;
 
@@ -26,54 +28,81 @@ int PostLinker64(char *fileName)
 
 	if (fp == NULL) {
 		printf("error[%s] : file open\n", fileName);
-		exit(EXIT_FAILURE);
+		exit_status = EXIT_FAILURE;
+		goto finalize;
 	}
 
 	/* Get ELF Header Information */
 	if (!fread((void*)&elfHeader, sizeof(Elf64_Ehdr), 1, fp)) {
 		printf("error[%s] : file read <header>\n", fileName);
-		exit(EXIT_FAILURE);
+		exit_status = EXIT_FAILURE;
+		goto finalize;
 	}
 
 	if (*(int*)(elfHeader.e_ident) != ELF_FORMAT) {
 		printf("error[%s] : no ELF format\n", fileName);
-		exit(EXIT_FAILURE);
+		exit_status = EXIT_FAILURE;
+		goto finalize;
 	}
 
 	if (elfHeader.e_ident[EI_CLASS] != 2) {
 		printf("error[%s] : EI_CLASS is not 64-bit format\n", fileName);
-		exit(EXIT_FAILURE);
+		exit_status = EXIT_FAILURE;
+		goto finalize;
 	}
 
-	pSegmentEntries = NULL;
-	DynEntries = NULL;
-
 	/* Memory Allocation for Segment and Section entries */
 	pSegmentEntries = (Elf64_Phdr *)malloc(elfHeader.e_phentsize * elfHeader.e_phnum);
+	if (pSegmentEntries == NULL) {
+		printf("error[%s] : memory allocation <pSegmentEntries>\n", fileName);
+		exit_status = EXIT_FAILURE;
+		goto finalize;
+	}
 
 	/* Get Segment Header Information using ELF Header */
-	fseek(fp, elfHeader.e_phoff, SEEK_SET);
-	if (!fread(pSegmentEntries, elfHeader.e_phentsize * elfHeader.e_phnum, 1, fp))
+	if (fseek(fp, elfHeader.e_phoff, SEEK_SET) != 0) {
+		printf("error[%s] : fseek <pSegmentEntries>\n", fileName);
+		exit_status = EXIT_FAILURE;
+		goto finalize;
+	}
+
+	if (!fread(pSegmentEntries, elfHeader.e_phentsize * elfHeader.e_phnum, 1, fp)) {
 		printf("error[%s] : file read <pSegmentEntries>\n", fileName);
+		exit_status = EXIT_FAILURE;
+		goto finalize;
+	}
 
 	/* Get Dynamic segment Information using Segment Header */
 	for (int i = 0; i < elfHeader.e_phnum; ++i) {
 		if (pSegmentEntries[i].p_type == PT_DYNAMIC) {
-			fseek(fp, pSegmentEntries[i].p_offset, SEEK_SET);
+			if (fseek(fp, pSegmentEntries[i].p_offset, SEEK_SET) != 0) {
+				printf("error[%s] : fseek <DynEntries>\n", fileName);
+				exit_status = EXIT_FAILURE;
+				goto finalize;
+			}
+
 			DynEntries = (Elf64_Dyn*)malloc(pSegmentEntries[i].p_memsz);
 			if (DynEntries == NULL) {
 				printf("error[%s] : memory allocation <DynEntries>\n", fileName);
-				exit(EXIT_FAILURE);
+				exit_status = EXIT_FAILURE;
+				goto finalize;
 			}
 			if (!fread(DynEntries, pSegmentEntries[i].p_memsz, 1, fp)) {
 				printf("error[%s] : file read <DynEntries>\n", fileName);
-				exit(EXIT_FAILURE);
+				exit_status = EXIT_FAILURE;
+				goto finalize;
 			}
 			pos = i;
 			break;
 		}
 	}
 
+	if (DynEntries == NULL) {
+		printf("error[%s] : no such p_type <PT_DYNAMIC>\n", fileName);
+		exit_status = EXIT_FAILURE;
+		goto finalize;
+	}
+
 	for (int i = 0; DynEntries[i].d_tag != DT_NULL; ++i) {
 		switch (DynEntries[i].d_tag) {
 		case DT_NEEDED:
@@ -96,14 +125,36 @@ int PostLinker64(char *fileName)
 		}
 	}
 
-	fseek(fp, hashOff + 4, SEEK_SET);
-	if (!fread(&totalSymbolCnt, sizeof(Elf64_Word), 1, fp))
+	if (fseek(fp, hashOff + 4, SEEK_SET) != 0) {
+		printf("error[%s] : fseek <symEntryCnt>\n", fileName);
+		exit_status = EXIT_FAILURE;
+		goto finalize;
+	}
+
+	if (!fread(&totalSymbolCnt, sizeof(Elf64_Word), 1, fp)) {
 		printf("error[%s] : file read <symEntryCnt>\n", fileName);
+		exit_status = EXIT_FAILURE;
+		goto finalize;
+	}
 
 	pMapSymbolTable = (Elf64_Sym*)malloc(totalSymbolCnt * sizeof(Elf64_Sym));
-	fseek(fp, symbolOff, SEEK_SET);
-	if (!fread(pMapSymbolTable, symEntrySize * totalSymbolCnt, 1, fp))
+	if (pMapSymbolTable == NULL) {
+		printf("error[%s] : memory allocation <pMapSymbolTable>\n", fileName);
+		exit_status = EXIT_FAILURE;
+		goto finalize;
+	}
+
+	if (fseek(fp, symbolOff, SEEK_SET) != 0) {
+		printf("error[%s] : fseek <pMapSymbolTable>\n", fileName);
+		exit_status = EXIT_FAILURE;
+		goto finalize;
+	}
+
+	if (!fread(pMapSymbolTable, symEntrySize * totalSymbolCnt, 1, fp)) {
 		printf("error[%s] : file read <pMapSymbolTable>\n", fileName);
+		exit_status = EXIT_FAILURE;
+		goto finalize;
+	}
 
 	for (int i = 1; i < totalSymbolCnt; ++i) {
 		if (pMapSymbolTable[i].st_shndx == SHN_UNDEF &&
@@ -122,23 +173,46 @@ int PostLinker64(char *fileName)
 
 	if (!done) {
 		printf("Undefined symbols are already removed : %s\n", fileName);
-		return 1;
+		exit_status = EXIT_SUCCESS;
+		goto finalize;
+	}
+
+	if (fseek(fp, pSegmentEntries[pos].p_offset, SEEK_SET) != 0) {
+		printf("error[%s] : fseek <DynEntries>\n", fileName);
+		exit_status = EXIT_FAILURE;
+		goto finalize;
 	}
 
-	fseek(fp, pSegmentEntries[pos].p_offset, SEEK_SET);
 	if (!fwrite(DynEntries, pSegmentEntries[pos].p_memsz, 1, fp)) {
 		printf("error[%s] : file write <DynEntries>\n", fileName);
-		exit(EXIT_FAILURE);
+		exit_status = EXIT_FAILURE;
+		goto finalize;
+	}
+
+	if (fseek(fp, symbolOff, SEEK_SET) != 0) {
+		printf("error[%s] : fseek <pMapSymbolTable>\n", fileName);
+		exit_status = EXIT_FAILURE;
+		goto finalize;
 	}
 
-	fseek(fp, symbolOff, SEEK_SET);
 	if (!fwrite(pMapSymbolTable, symEntrySize * totalSymbolCnt, 1, fp)) {
 		printf("error[%s] : file write <pMapSymbolTable>\n", fileName);
-		exit(EXIT_FAILURE);
+		exit_status = EXIT_FAILURE;
+		goto finalize;
 	}
 
-	fclose(fp);
 	printf("Success : %s\n", fileName);
+
+finalize:
+	free(pMapSymbolTable);
+	free(DynEntries);
+	free(pSegmentEntries);
+	if (fp != NULL)
+		fclose(fp);
+
+	if (exit_status == EXIT_FAILURE)
+		exit(EXIT_FAILURE);
+
 	return 1;
 }
 
@@ -146,4 +220,6 @@ int main(int argc, char **argv)
 {
 	for (int i = 1; i < argc; ++i)
 		PostLinker64(argv[i]);
+
+	return 0;
 }