package/system/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  * Copyright (C) 2005-2009 Felix Fietkau <nbd@nbd.name>
   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 v2
   9  * as published by the Free Software Foundation.
  10  *
  11  * This program is distributed in the hope that it will be useful,
  12  * but WITHOUT ANY WARRANTY; without even the implied warranty of
  13  * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  14  * GNU General Public License for more details.
  15  *
  16  * You should have received a copy of the GNU General Public License
  17  * along with this program; if not, write to the Free Software
  18  * Foundation, Inc., 59 Temple Place - Suite 330, Boston, MA 02111-1307, USA.
  19  *
  20  *
  21  * The code is based on the linux-mtd examples.
  22  */
  23
  24 #define _GNU_SOURCE
  25 #include <byteswap.h>
  26 #include <endian.h>
  27 #include <limits.h>
  28 #include <unistd.h>
  29 #include <stdlib.h>
  30 #include <stdio.h>
  31 #include <stdint.h>
  32 #include <signal.h>
  33 #include <sys/ioctl.h>
  34 #include <sys/syscall.h>
  35 #include <fcntl.h>
  36 #include <errno.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/mtd-user.h>
  47 #include "crc32.h"
  48 #include "fis.h"
  49 #include "mtd.h"
  50
  51 #include <libubox/md5.h>
  52
  53 #define MAX_ARGS 8
  54 #define JFFS2_DEFAULT_DIR       "" /* directory name without /, empty means root dir */
  55
  56 #define TRX_MAGIC               0x48445230      /* "HDR0" */
  57 #define SEAMA_MAGIC             0x5ea3a417
  58 #define WRG_MAGIC               0x20040220
  59 #define WRGG03_MAGIC            0x20080321
  60
  61 #if !defined(__BYTE_ORDER)
  62 #error "Unknown byte order"
  63 #endif
  64
  65 #if __BYTE_ORDER == __BIG_ENDIAN
  66 #define cpu_to_be32(x)  (x)
  67 #define be32_to_cpu(x)  (x)
  68 #define le32_to_cpu(x)  bswap_32(x)
  69 #elif __BYTE_ORDER == __LITTLE_ENDIAN
  70 #define cpu_to_be32(x)  bswap_32(x)
  71 #define be32_to_cpu(x)  bswap_32(x)
  72 #define le32_to_cpu(x)  (x)
  73 #else
  74 #error "Unsupported endianness"
  75 #endif
  76
  77 enum mtd_image_format {
  78         MTD_IMAGE_FORMAT_UNKNOWN,
  79         MTD_IMAGE_FORMAT_TRX,
  80         MTD_IMAGE_FORMAT_SEAMA,
  81         MTD_IMAGE_FORMAT_WRG,
  82         MTD_IMAGE_FORMAT_WRGG03,
  83 };
  84
  85 static char *buf = NULL;
  86 static char *imagefile = NULL;
  87 static enum mtd_image_format imageformat = MTD_IMAGE_FORMAT_UNKNOWN;
  88 static char *jffs2file = NULL, *jffs2dir = JFFS2_DEFAULT_DIR;
  89 static char *tpl_uboot_args_part;
  90 static int buflen = 0;
  91 int quiet;
  92 int no_erase;
  93 int mtdsize = 0;
  94 int erasesize = 0;
  95 int jffs2_skip_bytes=0;
  96 int mtdtype = 0;
  97 uint32_t opt_trxmagic = TRX_MAGIC;
  98
  99 int mtd_open(const char *mtd, bool block)
 100 {
 101         FILE *fp;
 102         char dev[PATH_MAX];
 103         int i;
 104         int ret;
 105         int flags = O_RDWR | O_SYNC;
 106         char name[PATH_MAX];
 107
 108         snprintf(name, sizeof(name), "\"%s\"", mtd);
 109         if ((fp = fopen("/proc/mtd", "r"))) {
 110                 while (fgets(dev, sizeof(dev), fp)) {
 111                         if (sscanf(dev, "mtd%d:", &i) && strstr(dev, name)) {
 112                                 snprintf(dev, sizeof(dev), "/dev/mtd%s/%d", (block ? "block" : ""), i);
 113                                 if ((ret=open(dev, flags))<0) {
 114                                         snprintf(dev, sizeof(dev), "/dev/mtd%s%d", (block ? "block" : ""), i);
 115                                         ret=open(dev, flags);
 116                                 }
 117                                 fclose(fp);
 118                                 return ret;
 119                         }
 120                 }
 121                 fclose(fp);
 122         }
 123
 124         return open(mtd, flags);
 125 }
 126
 127 int mtd_check_open(const char *mtd)
 128 {
 129         struct mtd_info_user mtdInfo;
 130         int fd;
 131
 132         fd = mtd_open(mtd, false);
 133         if(fd < 0) {
 134                 fprintf(stderr, "Could not open mtd device: %s\n", mtd);
 135                 return -1;
 136         }
 137
 138         if(ioctl(fd, MEMGETINFO, &mtdInfo)) {
 139                 fprintf(stderr, "Could not get MTD device info from %s\n", mtd);
 140                 close(fd);
 141                 return -1;
 142         }
 143         mtdsize = mtdInfo.size;
 144         erasesize = mtdInfo.erasesize;
 145         mtdtype = mtdInfo.type;
 146
 147         return fd;
 148 }
 149
 150 int mtd_block_is_bad(int fd, int offset)
 151 {
 152         int r = 0;
 153         loff_t o = offset;
 154
 155         if (mtdtype == MTD_NANDFLASH)
 156         {
 157                 r = ioctl(fd, MEMGETBADBLOCK, &o);
 158                 if (r < 0)
 159                 {
 160                         fprintf(stderr, "Failed to get erase block status\n");
 161                         exit(1);
 162                 }
 163         }
 164         return r;
 165 }
 166
 167 int mtd_erase_block(int fd, int offset)
 168 {
 169         struct erase_info_user mtdEraseInfo;
 170
 171         mtdEraseInfo.start = offset;
 172         mtdEraseInfo.length = erasesize;
 173         ioctl(fd, MEMUNLOCK, &mtdEraseInfo);
 174         if (ioctl (fd, MEMERASE, &mtdEraseInfo) < 0)
 175                 return -1;
 176
 177         return 0;
 178 }
 179
 180 int mtd_write_buffer(int fd, const char *buf, int offset, int length)
 181 {
 182         lseek(fd, offset, SEEK_SET);
 183         write(fd, buf, length);
 184         return 0;
 185 }
 186
 187 static int
 188 image_check(int imagefd, const char *mtd)
 189 {
 190         uint32_t magic;
 191         int ret = 1;
 192         int bufread;
 193
 194         while (buflen < sizeof(magic)) {
 195                 bufread = read(imagefd, buf + buflen, sizeof(magic) - buflen);
 196                 if (bufread < 1)
 197                         break;
 198
 199                 buflen += bufread;
 200         }
 201
 202         if (buflen < sizeof(magic)) {
 203                 fprintf(stdout, "Could not get image magic\n");
 204                 return 0;
 205         }
 206
 207         magic = ((uint32_t *)buf)[0];
 208
 209         if (be32_to_cpu(magic) == opt_trxmagic)
 210                 imageformat = MTD_IMAGE_FORMAT_TRX;
 211         else if (be32_to_cpu(magic) == SEAMA_MAGIC)
 212                 imageformat = MTD_IMAGE_FORMAT_SEAMA;
 213         else if (le32_to_cpu(magic) == WRG_MAGIC)
 214                 imageformat = MTD_IMAGE_FORMAT_WRG;
 215         else if (le32_to_cpu(magic) == WRGG03_MAGIC)
 216                 imageformat = MTD_IMAGE_FORMAT_WRGG03;
 217
 218         switch (imageformat) {
 219         case MTD_IMAGE_FORMAT_TRX:
 220                 if (trx_check)
 221                         ret = trx_check(imagefd, mtd, buf, &buflen);
 222                 break;
 223         case MTD_IMAGE_FORMAT_SEAMA:
 224         case MTD_IMAGE_FORMAT_WRG:
 225         case MTD_IMAGE_FORMAT_WRGG03:
 226                 break;
 227         default:
 228 #ifdef target_brcm
 229                 if (!strcmp(mtd, "firmware"))
 230                         ret = 0;
 231 #endif
 232                 break;
 233         }
 234
 235         return ret;
 236 }
 237
 238 static int mtd_check(const char *mtd)
 239 {
 240         char *next = NULL;
 241         char *str = NULL;
 242         int fd;
 243
 244         if (strchr(mtd, ':')) {
 245                 str = strdup(mtd);
 246                 mtd = str;
 247         }
 248
 249         do {
 250                 next = strchr(mtd, ':');
 251                 if (next) {
 252                         *next = 0;
 253                         next++;
 254                 }
 255
 256                 fd = mtd_check_open(mtd);
 257                 if (fd < 0)
 258                         return 0;
 259
 260                 if (!buf)
 261                         buf = malloc(erasesize);
 262
 263                 close(fd);
 264                 mtd = next;
 265         } while (next);
 266
 267         if (str)
 268                 free(str);
 269
 270         return 1;
 271 }
 272
 273 static int
 274 mtd_unlock(const char *mtd)
 275 {
 276         struct erase_info_user mtdLockInfo;
 277         char *next = NULL;
 278         char *str = NULL;
 279         int fd;
 280
 281         if (strchr(mtd, ':')) {
 282                 str = strdup(mtd);
 283                 mtd = str;
 284         }
 285
 286         do {
 287                 next = strchr(mtd, ':');
 288                 if (next) {
 289                         *next = 0;
 290                         next++;
 291                 }
 292
 293                 fd = mtd_check_open(mtd);
 294                 if(fd < 0) {
 295                         fprintf(stderr, "Could not open mtd device: %s\n", mtd);
 296                         exit(1);
 297                 }
 298
 299                 if (quiet < 2)
 300                         fprintf(stderr, "Unlocking %s ...\n", mtd);
 301
 302                 mtdLockInfo.start = 0;
 303                 mtdLockInfo.length = mtdsize;
 304                 ioctl(fd, MEMUNLOCK, &mtdLockInfo);
 305                 close(fd);
 306                 mtd = next;
 307         } while (next);
 308
 309         if (str)
 310                 free(str);
 311
 312         return 0;
 313 }
 314
 315 static int
 316 mtd_erase(const char *mtd)
 317 {
 318         int fd;
 319         struct erase_info_user mtdEraseInfo;
 320
 321         if (quiet < 2)
 322                 fprintf(stderr, "Erasing %s ...\n", mtd);
 323
 324         fd = mtd_check_open(mtd);
 325         if(fd < 0) {
 326                 fprintf(stderr, "Could not open mtd device: %s\n", mtd);
 327                 exit(1);
 328         }
 329
 330         mtdEraseInfo.length = erasesize;
 331
 332         for (mtdEraseInfo.start = 0;
 333                  mtdEraseInfo.start < mtdsize;
 334                  mtdEraseInfo.start += erasesize) {
 335                 if (mtd_block_is_bad(fd, mtdEraseInfo.start)) {
 336                         if (!quiet)
 337                                 fprintf(stderr, "\nSkipping bad block at 0x%x   ", mtdEraseInfo.start);
 338                 } else {
 339                         ioctl(fd, MEMUNLOCK, &mtdEraseInfo);
 340                         if(ioctl(fd, MEMERASE, &mtdEraseInfo))
 341                                 fprintf(stderr, "Failed to erase block on %s at 0x%x\n", mtd, mtdEraseInfo.start);
 342                 }
 343         }
 344
 345         close(fd);
 346         return 0;
 347
 348 }
 349
 350 static int
 351 mtd_dump(const char *mtd, int part_offset, int size)
 352 {
 353         int ret = 0, offset = 0;
 354         int fd;
 355         char *buf;
 356
 357         if (quiet < 2)
 358                 fprintf(stderr, "Dumping %s ...\n", mtd);
 359
 360         fd = mtd_check_open(mtd);
 361         if(fd < 0) {
 362                 fprintf(stderr, "Could not open mtd device: %s\n", mtd);
 363                 return -1;
 364         }
 365
 366         if (!size)
 367                 size = mtdsize;
 368
 369         if (part_offset)
 370                 lseek(fd, part_offset, SEEK_SET);
 371
 372         buf = malloc(erasesize);
 373         if (!buf)
 374                 return -1;
 375
 376         do {
 377                 int len = (size > erasesize) ? (erasesize) : (size);
 378                 int rlen = read(fd, buf, len);
 379
 380                 if (rlen < 0) {
 381                         if (errno == EINTR)
 382                                 continue;
 383                         ret = -1;
 384                         goto out;
 385                 }
 386                 if (!rlen || rlen != len)
 387                         break;
 388                 if (mtd_block_is_bad(fd, offset)) {
 389                         fprintf(stderr, "skipping bad block at 0x%08x\n", offset);
 390                 } else {
 391                         size -= rlen;
 392                         write(1, buf, rlen);
 393                 }
 394                 offset += rlen;
 395         } while (size > 0);
 396
 397 out:
 398         close(fd);
 399         return ret;
 400 }
 401
 402 static int
 403 mtd_verify(const char *mtd, char *file)
 404 {
 405         uint32_t f_md5[4], m_md5[4];
 406         struct stat s;
 407         md5_ctx_t ctx;
 408         int ret = 0;
 409         int fd;
 410
 411         if (quiet < 2)
 412                 fprintf(stderr, "Verifying %s against %s ...\n", mtd, file);
 413
 414         if (stat(file, &s) || md5sum(file, f_md5) < 0) {
 415                 fprintf(stderr, "Failed to hash %s\n", file);
 416                 return -1;
 417         }
 418
 419         fd = mtd_check_open(mtd);
 420         if(fd < 0) {
 421                 fprintf(stderr, "Could not open mtd device: %s\n", mtd);
 422                 return -1;
 423         }
 424
 425         md5_begin(&ctx);
 426         do {
 427                 char buf[256];
 428                 int len = (s.st_size > sizeof(buf)) ? (sizeof(buf)) : (s.st_size);
 429                 int rlen = read(fd, buf, len);
 430
 431                 if (rlen < 0) {
 432                         if (errno == EINTR)
 433                                 continue;
 434                         ret = -1;
 435                         goto out;
 436                 }
 437                 if (!rlen)
 438                         break;
 439                 md5_hash(buf, rlen, &ctx);
 440                 s.st_size -= rlen;
 441         } while (s.st_size > 0);
 442
 443         md5_end(m_md5, &ctx);
 444
 445         fprintf(stderr, "%08x%08x%08x%08x - %s\n", m_md5[0], m_md5[1], m_md5[2], m_md5[3], mtd);
 446         fprintf(stderr, "%08x%08x%08x%08x - %s\n", f_md5[0], f_md5[1], f_md5[2], f_md5[3], file);
 447
 448         ret = memcmp(f_md5, m_md5, sizeof(m_md5));
 449         if (!ret)
 450                 fprintf(stderr, "Success\n");
 451         else
 452                 fprintf(stderr, "Failed\n");
 453
 454 out:
 455         close(fd);
 456         return ret;
 457 }
 458
 459 static void
 460 indicate_writing(const char *mtd)
 461 {
 462         if (quiet < 2)
 463                 fprintf(stderr, "\nWriting from %s to %s ... ", imagefile, mtd);
 464
 465         if (!quiet)
 466                 fprintf(stderr, " [ ]");
 467 }
 468
 469 static int
 470 mtd_write(int imagefd, const char *mtd, char *fis_layout, size_t part_offset)
 471 {
 472         char *next = NULL;
 473         char *str = NULL;
 474         int fd, result;
 475         ssize_t r, w, e;
 476         ssize_t skip = 0;
 477         uint32_t offset = 0;
 478         int buflen_raw = 0;
 479         int jffs2_replaced = 0;
 480         int skip_bad_blocks = 0;
 481
 482 #ifdef FIS_SUPPORT
 483         static struct fis_part new_parts[MAX_ARGS];
 484         static struct fis_part old_parts[MAX_ARGS];
 485         struct fis_part *cur_part = NULL;
 486         int n_new = 0, n_old = 0;
 487
 488         if (fis_layout) {
 489                 const char *tmp = mtd;
 490                 char *word, *brkt;
 491                 int ret;
 492
 493                 memset(&old_parts, 0, sizeof(old_parts));
 494                 memset(&new_parts, 0, sizeof(new_parts));
 495                 if (!part_offset)
 496                         cur_part = new_parts;
 497
 498                 do {
 499                         next = strchr(tmp, ':');
 500                         if (!next)
 501                                 next = (char *) tmp + strlen(tmp);
 502
 503                         memcpy(old_parts[n_old].name, tmp, next - tmp);
 504
 505                         n_old++;
 506                         tmp = next + 1;
 507                 } while(*next);
 508
 509                 for (word = strtok_r(fis_layout, ",", &brkt);
 510                      word;
 511                          word = strtok_r(NULL, ",", &brkt)) {
 512
 513                         tmp = strtok(word, ":");
 514                         strncpy((char *) new_parts[n_new].name, tmp, sizeof(new_parts[n_new].name) - 1);
 515
 516                         tmp = strtok(NULL, ":");
 517                         if (!tmp)
 518                                 goto next;
 519
 520                         new_parts[n_new].size = strtoul(tmp, NULL, 0);
 521
 522                         tmp = strtok(NULL, ":");
 523                         if (!tmp)
 524                                 goto next;
 525
 526                         new_parts[n_new].loadaddr = strtoul(tmp, NULL, 16);
 527 next:
 528                         n_new++;
 529                 }
 530                 ret = fis_validate(old_parts, n_old, new_parts, n_new);
 531                 if (ret < 0) {
 532                         fprintf(stderr, "Failed to validate the new FIS partition table\n");
 533                         exit(1);
 534                 }
 535                 if (ret == 0)
 536                         fis_layout = NULL;
 537         }
 538 #endif
 539
 540         if (strchr(mtd, ':')) {
 541                 str = strdup(mtd);
 542                 mtd = str;
 543         }
 544
 545         r = 0;
 546
 547 resume:
 548         next = strchr(mtd, ':');
 549         if (next) {
 550                 *next = 0;
 551                 next++;
 552         }
 553
 554         fd = mtd_check_open(mtd);
 555         if(fd < 0) {
 556                 fprintf(stderr, "Could not open mtd device: %s\n", mtd);
 557                 exit(1);
 558         }
 559         if (part_offset > 0) {
 560                 fprintf(stderr, "Seeking on mtd device '%s' to: %zu\n", mtd, part_offset);
 561                 lseek(fd, part_offset, SEEK_SET);
 562         }
 563
 564         /* Write TP-Link recovery flag */
 565         if (tpl_uboot_args_part && mtd_tpl_recoverflag_write) {
 566                 if (quiet < 2)
 567                         fprintf(stderr, "Writing recovery flag to %s\n", tpl_uboot_args_part);
 568                 result = mtd_tpl_recoverflag_write(tpl_uboot_args_part, true);
 569                 if (result < 0) {
 570                         fprintf(stderr, "Could not write TP-Link recovery flag to %s: %i", mtd, result);
 571                         exit(1);
 572                 }
 573         }
 574
 575         indicate_writing(mtd);
 576
 577         w = e = 0;
 578         for (;;) {
 579                 /* buffer may contain data already (from trx check or last mtd partition write attempt) */
 580                 while (buflen < erasesize) {
 581                         r = read(imagefd, buf + buflen, erasesize - buflen);
 582                         if (r < 0) {
 583                                 if ((errno == EINTR) || (errno == EAGAIN))
 584                                         continue;
 585                                 else {
 586                                         perror("read");
 587                                         break;
 588                                 }
 589                         }
 590
 591                         if (r == 0)
 592                                 break;
 593
 594                         buflen += r;
 595                 }
 596
 597                 if (buflen_raw == 0)
 598                         buflen_raw = buflen;
 599
 600                 if (buflen == 0)
 601                         break;
 602
 603                 if (buflen < erasesize) {
 604                         /* Pad block to eraseblock size */
 605                         memset(&buf[buflen], 0xff, erasesize - buflen);
 606                         buflen = erasesize;
 607                 }
 608
 609                 if (skip > 0) {
 610                         skip -= buflen;
 611                         buflen_raw = 0;
 612                         buflen = 0;
 613                         if (skip <= 0)
 614                                 indicate_writing(mtd);
 615
 616                         continue;
 617                 }
 618
 619                 if (jffs2file && w >= jffs2_skip_bytes) {
 620                         if (memcmp(buf, JFFS2_EOF, sizeof(JFFS2_EOF) - 1) == 0) {
 621                                 if (!quiet)
 622                                         fprintf(stderr, "\b\b\b   ");
 623                                 if (quiet < 2)
 624                                         fprintf(stderr, "\nAppending jffs2 data from %s to %s..\n.", jffs2file, mtd);
 625                                 /* got an EOF marker - this is the place to add some jffs2 data */
 626                                 skip = mtd_replace_jffs2(mtd, fd, e, jffs2file);
 627                                 jffs2_replaced = 1;
 628
 629                                 /* don't add it again */
 630                                 jffs2file = NULL;
 631
 632                                 w += skip;
 633                                 e += skip;
 634                                 skip -= buflen;
 635                                 buflen_raw = 0;
 636                                 buflen = 0;
 637                                 offset = 0;
 638                                 continue;
 639                         }
 640                         /* no EOF marker, make sure we figure out the last inode number
 641                          * before appending some data */
 642                         mtd_parse_jffs2data(buf, jffs2dir);
 643                 }
 644
 645                 /* need to erase the next block before writing data to it */
 646                 if(!no_erase)
 647                 {
 648                         while (w + buflen > e - skip_bad_blocks) {
 649                                 if (!quiet)
 650                                         fprintf(stderr, "\b\b\b[e]");
 651
 652                                 if (mtd_block_is_bad(fd, e)) {
 653                                         if (!quiet)
 654                                                 fprintf(stderr, "\nSkipping bad block at 0x%08zx   ", e);
 655
 656                                         skip_bad_blocks += erasesize;
 657                                         e += erasesize;
 658
 659                                         // Move the file pointer along over the bad block.
 660                                         lseek(fd, erasesize, SEEK_CUR);
 661                                         continue;
 662                                 }
 663
 664                                 if (mtd_erase_block(fd, e + part_offset) < 0) {
 665                                         if (next) {
 666                                                 if (w < e) {
 667                                                         write(fd, buf + offset, e - w);
 668                                                         offset = e - w;
 669                                                 }
 670                                                 w = 0;
 671                                                 e = 0;
 672                                                 close(fd);
 673                                                 mtd = next;
 674                                                 fprintf(stderr, "\b\b\b   \n");
 675                                                 goto resume;
 676                                         } else {
 677                                                 fprintf(stderr, "Failed to erase block\n");
 678                                                 exit(1);
 679                                         }
 680                                 }
 681
 682                                 /* erase the chunk */
 683                                 e += erasesize;
 684                         }
 685                 }
 686
 687                 if (!quiet)
 688                         fprintf(stderr, "\b\b\b[w]");
 689
 690                 if ((result = write(fd, buf + offset, buflen)) < buflen) {
 691                         if (result < 0) {
 692                                 fprintf(stderr, "Error writing image.\n");
 693                                 exit(1);
 694                         } else {
 695                                 fprintf(stderr, "Insufficient space.\n");
 696                                 exit(1);
 697                         }
 698                 }
 699                 w += buflen;
 700
 701 #ifdef FIS_SUPPORT
 702                 if (cur_part && cur_part->size
 703                 && cur_part < &new_parts[MAX_ARGS - 1]
 704                 && cur_part->length + buflen_raw > cur_part->size)
 705                         cur_part++;
 706                 if (cur_part) {
 707                         cur_part->length += buflen_raw;
 708                         cur_part->crc = crc32(cur_part->crc, buf, buflen_raw);
 709                 }
 710 #endif
 711                 buflen_raw = 0;
 712                 buflen = 0;
 713                 offset = 0;
 714         }
 715
 716         if (jffs2_replaced) {
 717                 switch (imageformat) {
 718                 case MTD_IMAGE_FORMAT_TRX:
 719                         if (trx_fixup)
 720                                 trx_fixup(fd, mtd);
 721                         break;
 722                 case MTD_IMAGE_FORMAT_SEAMA:
 723                         if (mtd_fixseama)
 724                                 mtd_fixseama(mtd, 0, 0);
 725                         break;
 726                 case MTD_IMAGE_FORMAT_WRG:
 727                         if (mtd_fixwrg)
 728                                 mtd_fixwrg(mtd, 0, 0);
 729                         break;
 730                 case MTD_IMAGE_FORMAT_WRGG03:
 731                         if (mtd_fixwrgg)
 732                                 mtd_fixwrgg(mtd, 0, 0);
 733                         break;
 734                 default:
 735                         break;
 736                 }
 737         }
 738
 739         if (!quiet)
 740                 fprintf(stderr, "\b\b\b\b    ");
 741
 742         if (quiet < 2)
 743                 fprintf(stderr, "\n");
 744
 745 #ifdef FIS_SUPPORT
 746         if (fis_layout) {
 747                 if (fis_remap(old_parts, n_old, new_parts, n_new) < 0)
 748                         fprintf(stderr, "Failed to update the FIS partition table\n");
 749         }
 750 #endif
 751
 752         close(fd);
 753
 754         /* Clear TP-Link recovery flag */
 755         if (tpl_uboot_args_part && mtd_tpl_recoverflag_write) {
 756                 if (quiet < 2)
 757                         fprintf(stderr, "Removing recovery flag from %s\n", tpl_uboot_args_part);
 758                 result = mtd_tpl_recoverflag_write(tpl_uboot_args_part, false);
 759                 if (result < 0) {
 760                         fprintf(stderr, "Could not clear TP-Link recovery flag to %s: %i", mtd, result);
 761                         exit(1);
 762                 }
 763         }
 764
 765         return 0;
 766 }
 767
 768 static void usage(void)
 769 {
 770         fprintf(stderr, "Usage: mtd [<options> ...] <command> [<arguments> ...] <device>[:<device>...]\n\n"
 771         "The device is in the format of mtdX (eg: mtd4) or its label.\n"
 772         "mtd recognizes these commands:\n"
 773         "        unlock                  unlock the device\n"
 774         "        refresh                 refresh mtd partition\n"
 775         "        erase                   erase all data on device\n"
 776         "        verify <imagefile>|-    verify <imagefile> (use - for stdin) to device\n"
 777         "        write <imagefile>|-     write <imagefile> (use - for stdin) to device\n"
 778         "        jffs2write <file>       append <file> to the jffs2 partition on the device\n");
 779         if (mtd_resetbc) {
 780             fprintf(stderr,
 781         "        resetbc <device>        reset the uboot boot counter\n");
 782         }
 783         if (mtd_fixtrx) {
 784             fprintf(stderr,
 785         "        fixtrx                  fix the checksum in a trx header on first boot\n");
 786         }
 787         if (mtd_fixseama) {
 788             fprintf(stderr,
 789         "        fixseama                fix the checksum in a seama header on first boot\n");
 790         }
 791         if (mtd_fixwrg) {
 792             fprintf(stderr,
 793         "        fixwrg                  fix the checksum in a wrg header on first boot\n");
 794         }
 795         if (mtd_fixwrgg) {
 796             fprintf(stderr,
 797         "        fixwrgg                 fix the checksum in a wrgg header on first boot\n");
 798         }
 799         fprintf(stderr,
 800         "Following options are available:\n"
 801         "        -q                      quiet mode (once: no [w] on writing,\n"
 802         "                                           twice: no status messages)\n"
 803         "        -n                      write without first erasing the blocks\n"
 804         "        -r                      reboot after successful command\n"
 805         "        -f                      force write without trx checks\n"
 806         "        -e <device>             erase <device> before executing the command\n"
 807         "        -d <name>               directory for jffs2write, defaults to \"tmp\"\n"
 808         "        -j <name>               integrate <file> into jffs2 data when writing an image\n"
 809         "        -s <number>             skip the first n bytes when appending data to the jffs2 partiton, defaults to \"0\"\n"
 810         "        -p <number>             write beginning at partition offset\n"
 811         "        -l <length>             the length of data that we want to dump\n");
 812         if (mtd_fixtrx) {
 813             fprintf(stderr,
 814         "        -M <magic>              magic number of the image header in the partition (for fixtrx)\n"
 815         "        -o offset               offset of the image header in the partition(for fixtrx)\n");
 816         }
 817         if (mtd_fixtrx || mtd_fixseama || mtd_fixwrg || mtd_fixwrgg) {
 818                 fprintf(stderr,
 819         "        -c datasize             amount of data to be used for checksum calculation (for fixtrx / fixseama / fixwrg / fixwrgg)\n");
 820         }
 821         if (mtd_tpl_recoverflag_write) {
 822                 fprintf(stderr,
 823         "        -t <partition>          write TP-Link recovery-flag to <partition> (for write)\n");
 824         }
 825         fprintf(stderr,
 826 #ifdef FIS_SUPPORT
 827         "        -F <part>[:<size>[:<entrypoint>]][,<part>...]\n"
 828         "                                alter the fis partition table to create new partitions replacing\n"
 829         "                                the partitions provided as argument to the write command\n"
 830         "                                (only valid together with the write command)\n"
 831 #endif
 832         "\n"
 833         "Example: To write linux.trx to mtd4 labeled as linux and reboot afterwards\n"
 834         "         mtd -r write linux.trx linux\n\n");
 835         exit(1);
 836 }
 837
 838 static void do_reboot(void)
 839 {
 840         fprintf(stderr, "Rebooting ...\n");
 841         fflush(stderr);
 842
 843         /* try regular reboot method first */
 844         system("/sbin/reboot");
 845         sleep(2);
 846
 847         /* if we're still alive at this point, force the kernel to reboot */
 848         syscall(SYS_reboot,LINUX_REBOOT_MAGIC1,LINUX_REBOOT_MAGIC2,LINUX_REBOOT_CMD_RESTART,NULL);
 849 }
 850
 851 int main (int argc, char **argv)
 852 {
 853         int ch, i, boot, imagefd = 0, force, unlocked;
 854         char *erase[MAX_ARGS], *device = NULL;
 855         char *fis_layout = NULL;
 856         size_t offset = 0, data_size = 0, part_offset = 0, dump_len = 0;
 857         enum {
 858                 CMD_ERASE,
 859                 CMD_WRITE,
 860                 CMD_UNLOCK,
 861                 CMD_JFFS2WRITE,
 862                 CMD_FIXTRX,
 863                 CMD_FIXSEAMA,
 864                 CMD_FIXWRG,
 865                 CMD_FIXWRGG,
 866                 CMD_VERIFY,
 867                 CMD_DUMP,
 868                 CMD_RESETBC,
 869         } cmd = -1;
 870
 871         erase[0] = NULL;
 872         boot = 0;
 873         force = 0;
 874         buflen = 0;
 875         quiet = 0;
 876         no_erase = 0;
 877
 878         while ((ch = getopt(argc, argv,
 879 #ifdef FIS_SUPPORT
 880                         "F:"
 881 #endif
 882                         "frnqe:d:s:j:p:o:c:t:l:M:")) != -1)
 883                 switch (ch) {
 884                         case 'f':
 885                                 force = 1;
 886                                 break;
 887                         case 'r':
 888                                 boot = 1;
 889                                 break;
 890                         case 'n':
 891                                 no_erase = 1;
 892                                 break;
 893                         case 'j':
 894                                 jffs2file = optarg;
 895                                 break;
 896                         case 's':
 897                                 errno = 0;
 898                                 jffs2_skip_bytes = strtoul(optarg, 0, 0);
 899                                 if (errno) {
 900                                                 fprintf(stderr, "-s: illegal numeric string\n");
 901                                                 usage();
 902                                 }
 903                                 break;
 904                         case 'q':
 905                                 quiet++;
 906                                 break;
 907                         case 'e':
 908                                 i = 0;
 909                                 while ((erase[i] != NULL) && ((i + 1) < MAX_ARGS))
 910                                         i++;
 911
 912                                 erase[i++] = optarg;
 913                                 erase[i] = NULL;
 914                                 break;
 915                         case 'd':
 916                                 jffs2dir = optarg;
 917                                 break;
 918                         case 'p':
 919                                 errno = 0;
 920                                 part_offset = strtoul(optarg, 0, 0);
 921                                 if (errno) {
 922                                         fprintf(stderr, "-p: illegal numeric string\n");
 923                                         usage();
 924                                 }
 925                                 break;
 926                         case 'l':
 927                                 errno = 0;
 928                                 dump_len = strtoul(optarg, 0, 0);
 929                                 if (errno) {
 930                                         fprintf(stderr, "-l: illegal numeric string\n");
 931                                         usage();
 932                                 }
 933                                 break;
 934                         case 'M':
 935                                 errno = 0;
 936                                 opt_trxmagic = strtoul(optarg, 0, 0);
 937                                 if (errno) {
 938                                         fprintf(stderr, "-M: illegal numeric string\n");
 939                                         usage();
 940                                 }
 941                                 break;
 942                         case 'o':
 943                                 errno = 0;
 944                                 offset = strtoul(optarg, 0, 0);
 945                                 if (errno) {
 946                                         fprintf(stderr, "-o: illegal numeric string\n");
 947                                         usage();
 948                                 }
 949                                 break;
 950                         case 'c':
 951                                 errno = 0;
 952                                 data_size = strtoul(optarg, 0, 0);
 953                                 if (errno) {
 954                                         fprintf(stderr, "-c: illegal numeric string\n");
 955                                         usage();
 956                                 }
 957                                 break;
 958                         case 't':
 959                                 tpl_uboot_args_part = optarg;
 960                                 break;
 961 #ifdef FIS_SUPPORT
 962                         case 'F':
 963                                 fis_layout = optarg;
 964                                 break;
 965 #endif
 966                         case '?':
 967                         default:
 968                                 usage();
 969                 }
 970         argc -= optind;
 971         argv += optind;
 972
 973         if (argc < 2)
 974                 usage();
 975
 976         if ((strcmp(argv[0], "unlock") == 0) && (argc == 2)) {
 977                 cmd = CMD_UNLOCK;
 978                 device = argv[1];
 979         } else if ((strcmp(argv[0], "erase") == 0) && (argc == 2)) {
 980                 cmd = CMD_ERASE;
 981                 device = argv[1];
 982         } else if (((strcmp(argv[0], "resetbc") == 0) && (argc == 2)) && mtd_resetbc) {
 983                 cmd = CMD_RESETBC;
 984                 device = argv[1];
 985         } else if (((strcmp(argv[0], "fixtrx") == 0) && (argc == 2)) && mtd_fixtrx) {
 986                 cmd = CMD_FIXTRX;
 987                 device = argv[1];
 988         } else if (((strcmp(argv[0], "fixseama") == 0) && (argc == 2)) && mtd_fixseama) {
 989                 cmd = CMD_FIXSEAMA;
 990                 device = argv[1];
 991         } else if (((strcmp(argv[0], "fixwrg") == 0) && (argc == 2)) && mtd_fixwrg) {
 992                 cmd = CMD_FIXWRG;
 993                 device = argv[1];
 994         } else if (((strcmp(argv[0], "fixwrgg") == 0) && (argc == 2)) && mtd_fixwrgg) {
 995                 cmd = CMD_FIXWRGG;
 996                 device = argv[1];
 997         } else if ((strcmp(argv[0], "verify") == 0) && (argc == 3)) {
 998                 cmd = CMD_VERIFY;
 999                 imagefile = argv[1];
1000                 device = argv[2];
1001         } else if ((strcmp(argv[0], "dump") == 0) && (argc == 2)) {
1002                 cmd = CMD_DUMP;
1003                 device = argv[1];
1004         } else if ((strcmp(argv[0], "write") == 0) && (argc == 3)) {
1005                 cmd = CMD_WRITE;
1006                 device = argv[2];
1007
1008                 if (strcmp(argv[1], "-") == 0) {
1009                         imagefile = "<stdin>";
1010                         imagefd = 0;
1011                 } else {
1012                         imagefile = argv[1];
1013                         if ((imagefd = open(argv[1], O_RDONLY)) < 0) {
1014                                 fprintf(stderr, "Couldn't open image file: %s!\n", imagefile);
1015                                 exit(1);
1016                         }
1017                 }
1018
1019                 if (!mtd_check(device)) {
1020                         fprintf(stderr, "Can't open device for writing!\n");
1021                         exit(1);
1022                 }
1023                 /* check trx file before erasing or writing anything */
1024                 if (!image_check(imagefd, device) && !force) {
1025                         fprintf(stderr, "Image check failed.\n");
1026                         exit(1);
1027                 }
1028         } else if ((strcmp(argv[0], "jffs2write") == 0) && (argc == 3)) {
1029                 cmd = CMD_JFFS2WRITE;
1030                 device = argv[2];
1031
1032                 imagefile = argv[1];
1033                 if (!mtd_check(device)) {
1034                         fprintf(stderr, "Can't open device for writing!\n");
1035                         exit(1);
1036                 }
1037         } else {
1038                 usage();
1039         }
1040
1041         sync();
1042
1043         i = 0;
1044         unlocked = 0;
1045         while (erase[i] != NULL) {
1046                 mtd_unlock(erase[i]);
1047                 mtd_erase(erase[i]);
1048                 if (strcmp(erase[i], device) == 0)
1049                         unlocked = 1;
1050                 i++;
1051         }
1052
1053         switch (cmd) {
1054                 case CMD_UNLOCK:
1055                         if (!unlocked)
1056                                 mtd_unlock(device);
1057                         break;
1058                 case CMD_VERIFY:
1059                         mtd_verify(device, imagefile);
1060                         break;
1061                 case CMD_DUMP:
1062                         mtd_dump(device, offset, dump_len);
1063                         break;
1064                 case CMD_ERASE:
1065                         if (!unlocked)
1066                                 mtd_unlock(device);
1067                         mtd_erase(device);
1068                         break;
1069                 case CMD_WRITE:
1070                         if (!unlocked)
1071                                 mtd_unlock(device);
1072                         mtd_write(imagefd, device, fis_layout, part_offset);
1073                         break;
1074                 case CMD_JFFS2WRITE:
1075                         if (!unlocked)
1076                                 mtd_unlock(device);
1077                         mtd_write_jffs2(device, imagefile, jffs2dir);
1078                         break;
1079                 case CMD_FIXTRX:
1080                         if (mtd_fixtrx) {
1081                                 mtd_fixtrx(device, offset, data_size);
1082                         }
1083                         break;
1084                 case CMD_RESETBC:
1085                         if (mtd_resetbc) {
1086                                 mtd_resetbc(device);
1087                         }
1088                         break;
1089                 case CMD_FIXSEAMA:
1090                         if (mtd_fixseama)
1091                                 mtd_fixseama(device, 0, data_size);
1092                         break;
1093                 case CMD_FIXWRG:
1094                         if (mtd_fixwrg)
1095                                 mtd_fixwrg(device, 0, data_size);
1096                         break;
1097                 case CMD_FIXWRGG:
1098                         if (mtd_fixwrgg)
1099                                 mtd_fixwrgg(device, 0, data_size);
1100                         break;
1101         }
1102
1103         sync();
1104
1105         if (boot)
1106                 do_reboot();
1107
1108         return 0;
1109 }