mtd: improve support for reading input from a pipe. fixes sysupgrade
[openwrt.git] / package / mtd / src / mtd.c
1 /*
2  * mtd - simple memory technology device manipulation tool
3  *
4  * Copyright (C) 2005 Waldemar Brodkorb <wbx@dass-it.de>,
5  *                        Felix Fietkau <nbd@openwrt.org>
6  *
7  * This program is free software; you can redistribute it and/or
8  * modify it under the terms of the GNU General Public License
9  * as published by the Free Software Foundation; either version 2
10  * of the License, or (at your option) any later version.
11  *
12  * This program is distributed in the hope that it will be useful,
13  * but WITHOUT ANY WARRANTY; without even the implied warranty of
14  * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
15  * GNU General Public License for more details.
16  *
17  * You should have received a copy of the GNU General Public License
18  * along with this program; if not, write to the Free Software
19  * Foundation, Inc., 59 Temple Place - Suite 330, Boston, MA 02111-1307, USA.
20  *
21  * $Id$
22  *
23  * The code is based on the linux-mtd examples.
24  */
25
26 #include <limits.h>
27 #include <unistd.h>
28 #include <stdlib.h>
29 #include <stdio.h>
30 #include <stdint.h>
31 #include <signal.h>
32 #include <sys/ioctl.h>
33 #include <sys/syscall.h>
34 #include <fcntl.h>
35 #include <errno.h>
36 #include <error.h>
37 #include <time.h>
38 #include <string.h>
39 #include <sys/ioctl.h>
40 #include <sys/types.h>
41 #include <sys/param.h>
42 #include <sys/mount.h>
43 #include <sys/stat.h>
44 #include <sys/reboot.h>
45 #include <linux/reboot.h>
46 #include "mtd-api.h"
47 #include "mtd.h"
48
49 #define MAX_ARGS 8
50 #define JFFS2_DEFAULT_DIR       "" /* directory name without /, empty means root dir */
51
52 struct trx_header {
53         uint32_t magic;         /* "HDR0" */
54         uint32_t len;           /* Length of file including header */
55         uint32_t crc32;         /* 32-bit CRC from flag_version to end of file */
56         uint32_t flag_version;  /* 0:15 flags, 16:31 version */
57         uint32_t offsets[3];    /* Offsets of partitions from start of header */
58 };
59
60 static char *buf = NULL;
61 static char *imagefile = NULL;
62 static char *jffs2file = NULL, *jffs2dir = JFFS2_DEFAULT_DIR;
63 static int buflen = 0;
64 int quiet;
65 int mtdsize = 0;
66 int erasesize = 0;
67
68 int mtd_open(const char *mtd, bool block)
69 {
70         FILE *fp;
71         char dev[PATH_MAX];
72         int i;
73         int ret;
74         int flags = O_RDWR | O_SYNC;
75
76         if ((fp = fopen("/proc/mtd", "r"))) {
77                 while (fgets(dev, sizeof(dev), fp)) {
78                         if (sscanf(dev, "mtd%d:", &i) && strstr(dev, mtd)) {
79                                 snprintf(dev, sizeof(dev), "/dev/mtd%s/%d", (block ? "block" : ""), i);
80                                 if ((ret=open(dev, flags))<0) {
81                                         snprintf(dev, sizeof(dev), "/dev/mtd%s%d", (block ? "block" : ""), i);
82                                         ret=open(dev, flags);
83                                 }
84                                 fclose(fp);
85                                 return ret;
86                         }
87                 }
88                 fclose(fp);
89         }
90
91         return open(mtd, flags);
92 }
93
94 int mtd_check_open(const char *mtd)
95 {
96         struct mtd_info_user mtdInfo;
97         int fd;
98
99         fd = mtd_open(mtd, false);
100         if(fd < 0) {
101                 fprintf(stderr, "Could not open mtd device: %s\n", mtd);
102                 return 0;
103         }
104
105         if(ioctl(fd, MEMGETINFO, &mtdInfo)) {
106                 fprintf(stderr, "Could not get MTD device info from %s\n", mtd);
107                 close(fd);
108                 return 0;
109         }
110         mtdsize = mtdInfo.size;
111         erasesize = mtdInfo.erasesize;
112
113         return fd;
114 }
115
116 int mtd_erase_block(int fd, int offset)
117 {
118         struct erase_info_user mtdEraseInfo;
119
120         mtdEraseInfo.start = offset;
121         mtdEraseInfo.length = erasesize;
122         ioctl(fd, MEMUNLOCK, &mtdEraseInfo);
123         if (ioctl (fd, MEMERASE, &mtdEraseInfo) < 0) {
124                 fprintf(stderr, "Erasing mtd failed.\n");
125                 exit(1);
126         }
127         return 0;
128 }
129
130 int mtd_write_buffer(int fd, const char *buf, int offset, int length)
131 {
132         lseek(fd, offset, SEEK_SET);
133         write(fd, buf, length);
134         return 0;
135 }
136
137
138 static int
139 image_check(int imagefd, const char *mtd)
140 {
141 #ifdef target_brcm
142         return trx_check(imagefd, mtd, buf, &buflen);
143 #endif
144 }
145
146 static int mtd_check(const char *mtd)
147 {
148         int fd;
149
150         fd = mtd_check_open(mtd);
151         if (!fd)
152                 return 0;
153
154         if (!buf)
155                 buf = malloc(erasesize);
156
157         close(fd);
158         return 1;
159 }
160
161 static int
162 mtd_unlock(const char *mtd)
163 {
164         int fd;
165         struct erase_info_user mtdLockInfo;
166
167         fd = mtd_check_open(mtd);
168         if(fd <= 0) {
169                 fprintf(stderr, "Could not open mtd device: %s\n", mtd);
170                 exit(1);
171         }
172
173         if (quiet < 2) 
174                 fprintf(stderr, "Unlocking %s ...\n", mtd);
175
176         mtdLockInfo.start = 0;
177         mtdLockInfo.length = mtdsize;
178         if(ioctl(fd, MEMUNLOCK, &mtdLockInfo)) {
179                 close(fd);
180                 return 0;
181         }
182                 
183         close(fd);
184         return 0;
185 }
186
187 static int
188 mtd_erase(const char *mtd)
189 {
190         int fd;
191         struct erase_info_user mtdEraseInfo;
192
193         if (quiet < 2)
194                 fprintf(stderr, "Erasing %s ...\n", mtd);
195
196         fd = mtd_check_open(mtd);
197         if(fd <= 0) {
198                 fprintf(stderr, "Could not open mtd device: %s\n", mtd);
199                 exit(1);
200         }
201
202         mtdEraseInfo.length = erasesize;
203
204         for (mtdEraseInfo.start = 0;
205                  mtdEraseInfo.start < mtdsize;
206                  mtdEraseInfo.start += erasesize) {
207                 
208                 ioctl(fd, MEMUNLOCK, &mtdEraseInfo);
209                 if(ioctl(fd, MEMERASE, &mtdEraseInfo))
210                         fprintf(stderr, "Failed to erase block on %s at 0x%x\n", mtd, mtdEraseInfo.start);
211         }               
212
213         close(fd);
214         return 0;
215
216 }
217
218 static int
219 mtd_refresh(const char *mtd)
220 {
221         int fd;
222
223         if (quiet < 2)
224                 fprintf(stderr, "Refreshing mtd partition %s ... ", mtd);
225
226         fd = mtd_check_open(mtd);
227         if(fd <= 0) {
228                 fprintf(stderr, "Could not open mtd device: %s\n", mtd);
229                 exit(1);
230         }
231
232         if (ioctl(fd, MTDREFRESH, NULL)) {
233                 fprintf(stderr, "Failed to refresh the MTD device\n");
234                 close(fd);
235                 exit(1);
236         }
237         close(fd);
238
239         if (quiet < 2)
240                 fprintf(stderr, "\n");
241
242         return 0;
243 }
244
245 static int
246 mtd_write(int imagefd, const char *mtd)
247 {
248         int fd, result;
249         ssize_t r, w, e;
250
251         fd = mtd_check_open(mtd);
252         if(fd < 0) {
253                 fprintf(stderr, "Could not open mtd device: %s\n", mtd);
254                 exit(1);
255         }
256                 
257         if (quiet < 2)
258                 fprintf(stderr, "Writing from %s to %s ... ", imagefile, mtd);
259
260         r = w = e = 0;
261         if (!quiet)
262                 fprintf(stderr, " [ ]");
263
264         for (;;) {
265                 /* buffer may contain data already (from trx check) */
266                 do {
267                         r = read(imagefd, buf + buflen, erasesize - buflen);
268                         if (r < 0) {
269                                 if ((errno == EINTR) || (errno == EAGAIN))
270                                         continue;
271                                 else {
272                                         perror("read");
273                                         break;
274                                 }
275                         }
276
277                         if (r == 0) {
278                                 fprintf(stderr, "No more data left\n");
279                                 break;
280                         }
281
282                         buflen += r;
283                 } while (buflen < erasesize);
284
285                 if (buflen == 0)
286                         break;
287
288                 if (jffs2file) {
289                         if (memcmp(buf, JFFS2_EOF, sizeof(JFFS2_EOF)) == 0) {
290                                 if (!quiet)
291                                         fprintf(stderr, "\b\b\b   ");
292                                 if (quiet < 2)
293                                         fprintf(stderr, "\nAppending jffs2 data to from %s to %s...", jffs2file, mtd);
294                                 /* got an EOF marker - this is the place to add some jffs2 data */
295                                 mtd_replace_jffs2(mtd, fd, e, jffs2file);
296                                 goto done;
297                         }
298                         /* no EOF marker, make sure we figure out the last inode number
299                          * before appending some data */
300                         mtd_parse_jffs2data(buf, jffs2dir);
301                 }
302
303                 /* need to erase the next block before writing data to it */
304                 while (w + buflen > e) {
305                         if (!quiet)
306                                 fprintf(stderr, "\b\b\b[e]");
307
308                         mtd_erase_block(fd, e);
309
310                         /* erase the chunk */
311                         e += erasesize;
312                 }
313                 
314                 if (!quiet)
315                         fprintf(stderr, "\b\b\b[w]");
316                 
317                 if ((result = write(fd, buf, buflen)) < buflen) {
318                         if (result < 0) {
319                                 fprintf(stderr, "Error writing image.\n");
320                                 exit(1);
321                         } else {
322                                 fprintf(stderr, "Insufficient space.\n");
323                                 exit(1);
324                         }
325                 }
326                 w += buflen;
327
328                 buflen = 0;
329         }
330         if (!quiet)
331                 fprintf(stderr, "\b\b\b\b");
332
333 done:
334         if (quiet < 2)
335                 fprintf(stderr, "\n");
336
337         close(fd);
338         return 0;
339 }
340
341 static void usage(void)
342 {
343         fprintf(stderr, "Usage: mtd [<options> ...] <command> [<arguments> ...] <device>\n\n"
344         "The device is in the format of mtdX (eg: mtd4) or its label.\n"
345         "mtd recognizes these commands:\n"
346         "        unlock                  unlock the device\n"
347         "        refresh                 refresh mtd partition\n"
348         "        erase                   erase all data on device\n"
349         "        write <imagefile>|-     write <imagefile> (use - for stdin) to device\n"
350         "        jffs2write <file>       append <file> to the jffs2 partition on the device\n"
351         "Following options are available:\n"
352         "        -q                      quiet mode (once: no [w] on writing,\n"
353         "                                           twice: no status messages)\n"
354         "        -r                      reboot after successful command\n"
355         "        -f                      force write without trx checks\n"
356         "        -e <device>             erase <device> before executing the command\n"
357         "        -d <name>               directory for jffs2write, defaults to \"tmp\"\n"
358         "        -j <name>               integrate <file> into jffs2 data when writing an image\n"
359         "\n"
360         "Example: To write linux.trx to mtd4 labeled as linux and reboot afterwards\n"
361         "         mtd -r write linux.trx linux\n\n");
362         exit(1);
363 }
364
365 static void do_reboot(void)
366 {
367         fprintf(stderr, "Rebooting ...\n");
368         fflush(stderr);
369
370         /* try regular reboot method first */
371         system("/sbin/reboot");
372         sleep(2);
373
374         /* if we're still alive at this point, force the kernel to reboot */
375         syscall(SYS_reboot,LINUX_REBOOT_MAGIC1,LINUX_REBOOT_MAGIC2,LINUX_REBOOT_CMD_RESTART,NULL);
376 }
377
378 int main (int argc, char **argv)
379 {
380         int ch, i, boot, imagefd = 0, force, unlocked;
381         char *erase[MAX_ARGS], *device = NULL;
382         enum {
383                 CMD_ERASE,
384                 CMD_WRITE,
385                 CMD_UNLOCK,
386                 CMD_REFRESH,
387                 CMD_JFFS2WRITE
388         } cmd = -1;
389         
390         erase[0] = NULL;
391         boot = 0;
392         force = 0;
393         buflen = 0;
394         quiet = 0;
395
396         while ((ch = getopt(argc, argv, "frqe:d:j:")) != -1)
397                 switch (ch) {
398                         case 'f':
399                                 force = 1;
400                                 break;
401                         case 'r':
402                                 boot = 1;
403                                 break;
404                         case 'j':
405                                 jffs2file = optarg;
406                                 break;
407                         case 'q':
408                                 quiet++;
409                                 break;
410                         case 'e':
411                                 i = 0;
412                                 while ((erase[i] != NULL) && ((i + 1) < MAX_ARGS))
413                                         i++;
414                                         
415                                 erase[i++] = optarg;
416                                 erase[i] = NULL;
417                                 break;
418                         case 'd':
419                                 jffs2dir = optarg;
420                                 break;
421                         case '?':
422                         default:
423                                 usage();
424                 }
425         argc -= optind;
426         argv += optind;
427         
428         if (argc < 2)
429                 usage();
430
431         if ((strcmp(argv[0], "unlock") == 0) && (argc == 2)) {
432                 cmd = CMD_UNLOCK;
433                 device = argv[1];
434         } else if ((strcmp(argv[0], "refresh") == 0) && (argc == 2)) {
435                 cmd = CMD_REFRESH;
436                 device = argv[1];
437         } else if ((strcmp(argv[0], "erase") == 0) && (argc == 2)) {
438                 cmd = CMD_ERASE;
439                 device = argv[1];
440         } else if ((strcmp(argv[0], "write") == 0) && (argc == 3)) {
441                 cmd = CMD_WRITE;
442                 device = argv[2];
443         
444                 if (strcmp(argv[1], "-") == 0) {
445                         imagefile = "<stdin>";
446                         imagefd = 0;
447                 } else {
448                         imagefile = argv[1];
449                         if ((imagefd = open(argv[1], O_RDONLY)) < 0) {
450                                 fprintf(stderr, "Couldn't open image file: %s!\n", imagefile);
451                                 exit(1);
452                         }
453                 }
454         
455                 if (!mtd_check(device)) {
456                         fprintf(stderr, "Can't open device for writing!\n");
457                         exit(1);
458                 }
459                 /* check trx file before erasing or writing anything */
460                 if (!image_check(imagefd, device) && !force) {
461                         fprintf(stderr, "Image check failed.\n");
462                         exit(1);
463                 }
464         } else if ((strcmp(argv[0], "jffs2write") == 0) && (argc == 3)) {
465                 cmd = CMD_JFFS2WRITE;
466                 device = argv[2];
467         
468                 imagefile = argv[1];
469                 if (!mtd_check(device)) {
470                         fprintf(stderr, "Can't open device for writing!\n");
471                         exit(1);
472                 }
473         } else {
474                 usage();
475         }
476
477         sync();
478         
479         i = 0;
480         unlocked = 0;
481         while (erase[i] != NULL) {
482                 mtd_unlock(erase[i]);
483                 mtd_erase(erase[i]);
484                 if (strcmp(erase[i], device) == 0)
485                         unlocked = 1;
486                 i++;
487         }
488         
489                 
490         switch (cmd) {
491                 case CMD_UNLOCK:
492                         if (!unlocked)
493                                 mtd_unlock(device);
494                         break;
495                 case CMD_ERASE:
496                         if (!unlocked)
497                                 mtd_unlock(device);
498                         mtd_erase(device);
499                         break;
500                 case CMD_WRITE:
501                         if (!unlocked)
502                                 mtd_unlock(device);
503                         mtd_write(imagefd, device);
504                         break;
505                 case CMD_JFFS2WRITE:
506                         if (!unlocked)
507                                 mtd_unlock(device);
508                         mtd_write_jffs2(device, imagefile, jffs2dir);
509                         break;
510                 case CMD_REFRESH:
511                         mtd_refresh(device);
512                         break;
513         }
514
515         sync();
516         
517         if (boot)
518                 do_reboot();
519
520         return 0;
521 }