Set current working URI to be that of the executable image when

executing any image, not just a script.  (This will enable pxelinux to
use relative URIs, should it wish to.)

2 files changed, 13 insertions, 12 deletions

diff --git a/src/core/image.c b/src/core/image.c
index 63c2502bd..440a68c9f 100644
--- a/src/core/image.c
+++ b/src/core/image.c
@@ -237,6 +237,7 @@ int image_autoload ( struct image *image ) {
  * @ret rc		Return status code
  */
 int image_exec ( struct image *image ) {
+	struct uri *old_cwuri;
 	int rc;
 
 	/* Image must be loaded first */
@@ -252,15 +253,23 @@ int image_exec ( struct image *image ) {
 	if ( ! image->type->exec )
 		return -ENOEXEC;
 
+	/* Switch current working directory to be that of the image itself */
+	old_cwuri = uri_get ( cwuri );
+	churi ( image->uri );
+
 	/* Try executing the image */
 	if ( ( rc = image->type->exec ( image ) ) != 0 ) {
 		DBGC ( image, "IMAGE %p could not execute: %s\n",
 		       image, strerror ( rc ) );
-		return rc;
+		goto done;
 	}
 
-	/* Well, some formats might return... */
-	return 0;
+ done:
+	/* Reset current working directory */
+	churi ( old_cwuri );
+	uri_put ( old_cwuri );
+
+	return rc;
 }
 
 /**
diff --git a/src/image/script.c b/src/image/script.c
index 2f159c97e..c8821522b 100644
--- a/src/image/script.c
+++ b/src/image/script.c
@@ -27,7 +27,6 @@
 #include <stdlib.h>
 #include <errno.h>
 #include <gpxe/image.h>
-#include <gpxe/uri.h>
 
 struct image_type script_image_type __image_type ( PROBE_NORMAL );
 
@@ -38,7 +37,6 @@ struct image_type script_image_type __image_type ( PROBE_NORMAL );
  * @ret rc		Return status code
  */
 static int script_exec ( struct image *image ) {
-	struct uri *old_cwuri;
 	char cmdbuf[256];
 	size_t offset = 0;
 	size_t remaining;
@@ -53,10 +51,6 @@ static int script_exec ( struct image *image ) {
 	image_get ( image );
 	unregister_image ( image );
 
-	/* Switch current working directory to be that of the script itself */
-	old_cwuri = uri_get ( cwuri );
-	churi ( image->uri );
-
 	while ( offset < image->len ) {
 	
 		/* Read up to cmdbuf bytes from script into buffer */
@@ -93,9 +87,7 @@ static int script_exec ( struct image *image ) {
 
 	rc = 0;
  done:
-	/* Reset current working directory, re-register image and return */
-	churi ( old_cwuri );
-	uri_put ( old_cwuri );
+	/* Re-register image and return */
 	register_image ( image );
 	image_put ( image );
 	return rc;