avoid leaking HMA on install error

1 files changed, 8 insertions, 6 deletions

diff --git a/sfmain.c b/sfmain.c
index 710a97c..e2f0b58 100644
--- a/sfmain.c
+++ b/sfmain.c
@@ -426,10 +426,7 @@ static __declspec(aborts) int install_driver(LPTSRDATA data, bool high)
 	// If we reallocated ourselves to UMB,
 	// it's time to free our initial conventional memory allocation
 	if (high) {
-		// We are about to free() our own code segment.
-		// Nothing should try to allocate memory between this and the TSR call
-		// below, since it could overwrite our code...
-		dos_free(_psp);
+		finish_reallocation(_psp, FP_SEG(data));
 	}
 
 	_dos_keep(EXIT_SUCCESS, get_paragraphs(resident_size));
@@ -570,9 +567,14 @@ int main(int argc, const char *argv[])
 			deallocate_environment(_psp);
 		}
 		err = configure_driver(data);
-		if (err) return EXIT_FAILURE;
+		if (err) {
+			if (high) cancel_reallocation(FP_SEG(data));
+			return EXIT_FAILURE;
+		}
 		err = automount(data);
-		if (err) return EXIT_FAILURE;
+		if (err) {
+			// Automount errors are not fatal
+		}
 		return install_driver(data, high);
 	} else if (stricmp(argv[argi], "uninstall") == 0) {
 		if (!data) return driver_not_found();