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