1 files changed, 23 insertions, 18 deletions

diff --git a/flash_lock.c b/flash_lock.c
index 8cd1b82..164eaa4 100644
--- a/flash_lock.c
+++ b/flash_lock.c
@@ -23,14 +23,13 @@ int main(int argc, char *argv[])
 	int fd;
 	struct mtd_info_user mtdInfo;
 	struct erase_info_user mtdLockInfo;
-	int num_sectors;
-	int ofs;
+	int count;
 
 	/*
 	 * Parse command line options
 	 */
-	if (argc != 4) {
-		fprintf(stderr, "USAGE: %s <mtd device> <ofs in hex> <num of sectors in decimal or -1 for all sectors>\n", PROGRAM_NAME);
+	if (argc < 2) {
+		fprintf(stderr, "USAGE: %s <mtd device> <offset> <block count>\n", PROGRAM_NAME);
 		exit(1);
 	} else if (strncmp(argv[1], "/dev/mtd", 8) != 0) {
 		fprintf(stderr, "'%s' is not a MTD device.  Must specify mtd device: /dev/mtd?\n", argv[1]);
@@ -48,26 +47,32 @@ int main(int argc, char *argv[])
 		close(fd);
 		exit(1);
 	}
-	sscanf(argv[2], "%x", &ofs);
-	sscanf(argv[3], "%d", &num_sectors);
-	if (ofs > mtdInfo.size - mtdInfo.erasesize) {
-		fprintf(stderr, "%x is beyond device size %x\n", ofs,
-			(unsigned int)(mtdInfo.size - mtdInfo.erasesize));
+
+	if (argc > 2)
+		mtdLockInfo.start = strtol(argv[2], NULL, 0);
+	else
+		mtdLockInfo.start = 0;
+	if (mtdLockInfo.start > mtdInfo.size) {
+		fprintf(stderr, "%#x is beyond device size %#x\n",
+			mtdLockInfo.start, mtdInfo.size);
+		close(fd);
 		exit(1);
 	}
 
-	if (num_sectors == -1) {
-		num_sectors = mtdInfo.size / mtdInfo.erasesize;
+	if (argc > 3) {
+		count = strtol(argv[3], NULL, 0);
+		if (count == -1)
+			mtdLockInfo.length = mtdInfo.size - mtdInfo.erasesize;
+		else
+			mtdLockInfo.length = mtdInfo.erasesize * count;
 	} else {
-		if (num_sectors > mtdInfo.size / mtdInfo.erasesize) {
-			fprintf(stderr, "%d are too many sectors, device only has %d\n",
-				num_sectors, (int)(mtdInfo.size / mtdInfo.erasesize));
-			exit(1);
-		}
+		mtdLockInfo.length = mtdInfo.size - mtdInfo.erasesize;
+	}
+	if (mtdLockInfo.start + mtdLockInfo.length > mtdInfo.size) {
+		fprintf(stderr, "lock range is more than device supports\n");
+		exit(1);
 	}
 
-	mtdLockInfo.start = ofs;
-	mtdLockInfo.length = (num_sectors - 1) * mtdInfo.erasesize;
 	if (ioctl(fd, MEMLOCK, &mtdLockInfo)) {
 		fprintf(stderr, "Could not lock MTD device: %s\n", argv[1]);
 		close(fd);