Do not read entire file into buffer when extracting

Only read FDT into buffer, and then in case of external image, use
lseek followed by sendfile if supported or otherwise read/write.

Remove unnecessary strdup of option arguments that would leak memory
in case options are repeated.

Fix confusing error message with --attribute when --image is used to
specify a non-existing image.

fdtextract.c

@@ -28,6 +28,7 @@
 #include <stdbool.h>
 #include <sys/types.h>
 #include <sys/stat.h>
+#include <sys/sendfile.h>
 #include <fcntl.h>
 #include <unistd.h>
 #include <errno.h>
@@ -37,8 +38,6 @@
 
 #include "util.h"
 
-#define FDT_MAGIC_SIZE  4
-#define MAX_VERSION 17
 #define MAX_PATH_LEN 100
 #define SHA_256_LEN 32
 
@@ -92,21 +91,56 @@ static int list_images(char *buf)
 	return 0;
 }
 
+static ssize_t copy_data(int out_fd, int in_fd, ssize_t size)
+{
+	ssize_t left = size;
+
+	if (left < 0) {
+		errno = EINVAL;
+		return -1;
+	}
+
+	while (left > 0) {
+		ssize_t written = sendfile(out_fd, in_fd, NULL, left);
+		if (written < 0 && errno == EINVAL)
+			break;
+		if (written < 0)
+			return -1;
+		if (written == 0)
+			return size - left;
+		left -= written;
+	}
+
+	while (left > 0) {
+		char buf[4096];
+		ssize_t count = sizeof(buf) < left ? sizeof(buf) : left;
+		ssize_t written;
+
+		count = read(in_fd, buf, count);
+		if (count < 0)
+			return -1;
+		if (count == 0)
+			break;
+
+		written = write(out_fd, buf, count);
+		if (written < 0)
+			return -1;
+		left -= written;
+		if (written < count)
+			break;
+	}
+
+	return size - left;
+}
+
 /* Extract an image embedded in the FIT. */
-static int extract_image(const char *file, char *name, char *out)
+static int extract_image(char *buf, char *name, int in_fd, int out_fd)
 {
 	char path[MAX_PATH_LEN] = {0};
-	int noffset, fd, count, data_size;
+	int noffset, count, data_size;
 	unsigned int data_offset;
 	const fdt32_t *val;
-	char *buf;
-	char *data;
-	size_t len;
-
-	buf = utilfdt_read(file, &len);
-	if (!buf) {
-		die("could not read: %s\n", file);
-	}
+	bool is_external = false;
 
 	snprintf(path, MAX_PATH_LEN, "/images/%s", name);
 	noffset = fdt_path_offset(buf, path);
@@ -115,9 +149,7 @@ static int extract_image(const char *file, char *name, char *out)
 		return -1;
 	}
 
-	bool is_external = false;
 	/* Get offset of image. Try both relative and absolute offset. */
-
 	if ((val = fdt_getprop(buf, noffset, "data-offset", NULL))) {
 		/* Relative offset */
 		data_offset = fdt32_to_cpu(*val);
@@ -136,32 +168,29 @@ static int extract_image(const char *file, char *name, char *out)
 			fprintf(stderr, "Error: Could not get image size: %s.\n", name);
 			return -1;
 		}
-		data = buf + data_offset;
 		data_size = fdt32_to_cpu(*val);
+		if (lseek(in_fd, data_offset, SEEK_SET) < 0) {
+			fprintf(stderr, "Error: Could not lseek to: %u\n", data_offset);
+			return errno;
+		}
+		count = copy_data(out_fd, in_fd, data_size);
 	} else {
-		data = fdt_getprop(buf, noffset, "data", &data_size);
+		const char *data = fdt_getprop(buf, noffset, "data", &data_size);
 		if (data_size < 0) {
 			fprintf(stderr, "Error: Could not get image data: %s.\n", name);
 			return -1;
 		}
+		count = write(out_fd, data, data_size);
 	}
 
-	/* If output file is provided, write to that file. If not, write
-	   to stdout. */
-	if (out) {
-		fd = open(out, O_WRONLY | O_CREAT | O_TRUNC, 0666);
-		if (fd < 0) {
-			fprintf(stderr, "Error opening output file %s.\n", out);
-			return errno;
-		}
-	} else {
-		fd = STDOUT_FILENO;
+	if (count < 0) {
+		fprintf(stderr, "Error: I/O error while copying image data.\n");
+		return errno;
 	}
 
-	count = write(fd, data, data_size);
 	if (count < data_size) {
-		fprintf(stderr, "Error writing output file %s.\n", out);
-		return errno;
+		fprintf(stderr, "Error: Image data was truncated.\n");
+		return -1;
 	}
 
 	return 0;
@@ -215,18 +244,14 @@ static int get_attribute(char *buf, char *name, char *imagename)
 {
 	int noffset;
 	const char *val = NULL;
-	char path[MAX_PATH_LEN] = {0};
+	char path[MAX_PATH_LEN] = "/";
 
-	if (imagename == NULL) {
-		/* Get path of root node. */
-		noffset = fdt_path_offset(buf, "/");
-	} else {
-		/* Get path of image node. */
+	if (imagename)
 		snprintf(path, MAX_PATH_LEN, "/images/%s", imagename);
-		noffset = fdt_path_offset(buf, path);
-	}
+
+	noffset = fdt_path_offset(buf, path);
 	if (noffset < 0) {
-		fprintf(stderr, "Error: could not find image hash: %s.\n", name);
+		fprintf(stderr, "Error: invalid FDT path: %s\n", path);
 		return -1;
 	}
 
@@ -241,15 +266,10 @@ static int get_attribute(char *buf, char *name, char *imagename)
 	return 0;
 }
 
-char *read_header(const char *file)
+char *read_header(int fd)
 {
 	char *buf;
-	int fd, total_size;
-	size_t len;
-
-	fd = open(file, O_RDONLY);
-	if (fd < 0)
-		return NULL;
+	ssize_t len, total_size;
 
 	/* Read minimal static struct */
 	buf = malloc(FDT_V1_SIZE);
@@ -257,18 +277,27 @@ char *read_header(const char *file)
 		return NULL;
 
 	len = read(fd, buf, FDT_V1_SIZE);
-	if (len < FDT_V1_SIZE)
+	if (len < FDT_V1_SIZE) {
+		free(buf);
 		return NULL;
+	}
 
 	/* Read rest of header */
 	total_size = fdt_totalsize(buf);
+	if (total_size < FDT_V1_SIZE) {
+		free(buf);
+		return NULL;
+	}
+
 	buf = realloc(buf, total_size);
 	if (!buf)
 		return NULL;
 
 	len = read(fd, buf + FDT_V1_SIZE, total_size - FDT_V1_SIZE);
-	if (len < total_size - FDT_V1_SIZE)
+	if (len < total_size - FDT_V1_SIZE) {
+		free(buf);
 		return NULL;
+	}
 
 	return buf;
 }
@@ -280,8 +309,13 @@ int main(int argc, char *argv[])
 	char *buf, *name = NULL, *out = NULL, *imagename = NULL;
 	bool list = false, extract = false,
 		hash = false, attribute = false;
+	int in_fd, out_fd = STDOUT_FILENO;
 
 	while ((opt = util_getopt_long()) != EOF) {
+		if ((opt == 'l' || opt == 'e' || opt == 's' || opt == 'a') &&
+		    (list || extract || hash || attribute))
+			usage("only one of --list/--extract/--hash/--attribute allowed");
+
 		switch (opt) {
 		case_USAGE_COMMON_FLAGS
 
@@ -290,21 +324,21 @@ int main(int argc, char *argv[])
 			break;
 		case 'e':
 			extract = true;
-			name = strdup(optarg);
+			name = optarg;
 			break;
 		case 'o':
-			out = strdup(optarg);
+			out = optarg;
 			break;
 		case 's':
 			hash = true;
-			name = strdup(optarg);
+			name = optarg;
 			break;
 		case 'a':
 			attribute = true;
-			name = strdup(optarg);
+			name = optarg;
 			break;
 		case 'i':
-			imagename = strdup(optarg);
+			imagename = optarg;
 			break;
 		}
 	}
@@ -312,7 +346,12 @@ int main(int argc, char *argv[])
 		usage("missing input filename");
 	file = argv[optind];
 
-	buf = read_header(file);
+	in_fd = open(file, O_RDONLY);
+	if (in_fd < 0) {
+		die("could not open: %s\n", file);
+	}
+
+	buf = read_header(in_fd);
 	if (!buf) {
 		die("could not read header from: %s\n", file);
 	}
@@ -325,6 +364,14 @@ int main(int argc, char *argv[])
 		die("--image should be used with --attribute\n");
 	}
 
+	if (out != NULL && extract == false) {
+		die("--out should be used with --extract\n");
+	}
+
+	if (out && (out_fd = open(out, O_WRONLY | O_CREAT | O_TRUNC, 0666)) < 0) {
+		die("could not open output file: %s\n", out);
+	}
+
 	/* Pass the pointer to the header to read attributes. */
 	if (list)
 		ret = list_images(buf);
@@ -337,17 +384,13 @@ int main(int argc, char *argv[])
 
 	/* Let extract image read the entire file. */
 	if (extract)
-		ret = extract_image(file, name, out);
+		ret = extract_image(buf, name, in_fd, out_fd);
 
-	/* Free up memory */
-	if(buf != NULL)
-		free(buf);
-	if(name != NULL)
-		free(name);
-	if(out != NULL)
-		free(out);
-	if(imagename != NULL)
-		free(imagename);
+	if (out_fd != STDOUT_FILENO)
+		close(out_fd);
+
+	free(buf);
+	close(in_fd);
 
 	return ret;
 }