Create macro
  for_each_opt_hdr_v1
and functions
  opt_hdr_v1_size(),
  opt_hdr_v1_valid_size(),
  opt_hdr_v1_ext(),
  opt_hdr_v1_first() and
  opt_hdr_v1_next()
to simplify iteration over version 1 optional headers.

This prevents ugly code repetition and makes it nicer to read.

Signed-off-by: Marek Behún <marek.be...@nic.cz>
---
 tools/kwbimage.c | 90 +++++++++++++-----------------------------------
 tools/kwbimage.h | 58 +++++++++++++++++++++++++++++++
 2 files changed, 81 insertions(+), 67 deletions(-)

diff --git a/tools/kwbimage.c b/tools/kwbimage.c
index ed11c835a5..6678ef15bd 100644
--- a/tools/kwbimage.c
+++ b/tools/kwbimage.c
@@ -1620,34 +1620,20 @@ static void kwbimage_set_header(void *ptr, struct stat 
*sbuf, int ifd,
 static void kwbimage_print_header(const void *ptr)
 {
        struct main_hdr_v0 *mhdr = (struct main_hdr_v0 *)ptr;
+       struct opt_hdr_v1 *ohdr;
 
        printf("Image Type:   MVEBU Boot from %s Image\n",
               image_boot_mode_name(mhdr->blockid));
        printf("Image version:%d\n", image_version((void *)ptr));
-       if (image_version((void *)ptr) == 1) {
-               struct main_hdr_v1 *mhdr = (struct main_hdr_v1 *)ptr;
 
-               if (mhdr->ext & 0x1) {
-                       struct opt_hdr_v1 *ohdr = (struct opt_hdr_v1 *)
-                                                 ((uint8_t *)ptr +
-                                                  sizeof(*mhdr));
-
-                       while (1) {
-                               uint32_t ohdr_size;
-
-                               ohdr_size = (ohdr->headersz_msb << 16) |
-                                           le16_to_cpu(ohdr->headersz_lsb);
-                               if (ohdr->headertype == OPT_HDR_V1_BINARY_TYPE) 
{
-                                       printf("BIN Hdr Size: ");
-                                       genimg_print_size(ohdr_size - 12 - 4 * 
ohdr->data[0]);
-                               }
-                               if (!(*((uint8_t *)ohdr + ohdr_size - 4) & 0x1))
-                                       break;
-                               ohdr = (struct opt_hdr_v1 *)((uint8_t *)ohdr +
-                                                            ohdr_size);
-                       }
+       for_each_opt_hdr_v1 (ohdr, mhdr) {
+               if (ohdr->headertype == OPT_HDR_V1_BINARY_TYPE) {
+                       printf("BIN Hdr Size: ");
+                       genimg_print_size(opt_hdr_v1_size(ohdr) - 12 -
+                                         4 * ohdr->data[0]);
                }
        }
+
        printf("Data Size:    ");
        genimg_print_size(mhdr->blocksize - sizeof(uint32_t));
        printf("Load Address: %08x\n", mhdr->destaddr);
@@ -1694,33 +1680,15 @@ static int kwbimage_verify_header(unsigned char *ptr, 
int image_size,
                }
        } else if (image_version((void *)ptr) == 1) {
                struct main_hdr_v1 *mhdr = (struct main_hdr_v1 *)ptr;
+               const uint8_t *mhdr_end;
+               struct opt_hdr_v1 *ohdr;
                uint32_t offset;
                uint32_t size;
 
-               if (mhdr->ext & 0x1) {
-                       uint32_t ohdr_size;
-                       struct opt_hdr_v1 *ohdr = (struct opt_hdr_v1 *)
-                                                 (ptr + sizeof(*mhdr));
-
-                       while (1) {
-                               if ((uint8_t *)ohdr + sizeof(*ohdr) >
-                                   (uint8_t *)mhdr + header_size)
-                                       return -FDT_ERR_BADSTRUCTURE;
-
-                               ohdr_size = (ohdr->headersz_msb << 16) |
-                                           le16_to_cpu(ohdr->headersz_lsb);
-
-                               if (ohdr_size < 8 ||
-                                   (uint8_t *)ohdr + ohdr_size >
-                                   (uint8_t *)mhdr + header_size)
-                                       return -FDT_ERR_BADSTRUCTURE;
-
-                               if (!(*((uint8_t *)ohdr + ohdr_size - 4) & 0x1))
-                                       break;
-                               ohdr = (struct opt_hdr_v1 *)((uint8_t *)ohdr +
-                                                            ohdr_size);
-                       }
-               }
+               mhdr_end = (uint8_t *)mhdr + header_size;
+               for_each_opt_hdr_v1 (ohdr, ptr)
+                       if (!opt_hdr_v1_valid_size(ohdr, mhdr_end))
+                               return -FDT_ERR_BADSTRUCTURE;
 
                offset = le32_to_cpu(mhdr->srcaddr);
 
@@ -1867,36 +1835,24 @@ static int kwbimage_extract_subimage(void *ptr, struct 
image_tool_params *params
 {
        struct main_hdr_v1 *mhdr = (struct main_hdr_v1 *)ptr;
        size_t header_size = kwbimage_header_size(ptr);
+       struct opt_hdr_v1 *ohdr;
        int idx = params->pflag;
        int cur_idx = 0;
        uint32_t offset;
        ulong image;
        ulong size;
 
-       if (image_version((void *)ptr) == 1 && (mhdr->ext & 0x1)) {
-               struct opt_hdr_v1 *ohdr = (struct opt_hdr_v1 *)
-                                         ((uint8_t *)ptr +
-                                          sizeof(*mhdr));
+       for_each_opt_hdr_v1 (ohdr, ptr) {
+               if (ohdr->headertype != OPT_HDR_V1_BINARY_TYPE)
+                       continue;
 
-               while (1) {
-                       uint32_t ohdr_size = (ohdr->headersz_msb << 16) |
-                                            le16_to_cpu(ohdr->headersz_lsb);
-
-                       if (ohdr->headertype == OPT_HDR_V1_BINARY_TYPE) {
-                               if (idx == cur_idx) {
-                                       image = (ulong)&ohdr->data[4 +
-                                                4 * ohdr->data[0]];
-                                       size = ohdr_size - 12 -
-                                              4 * ohdr->data[0];
-                                       goto extract;
-                               }
-                               ++cur_idx;
-                       }
-                       if (!(*((uint8_t *)ohdr + ohdr_size - 4) & 0x1))
-                               break;
-                       ohdr = (struct opt_hdr_v1 *)((uint8_t *)ohdr +
-                                                    ohdr_size);
+               if (idx == cur_idx) {
+                       image = (ulong)&ohdr->data[4 + 4 * ohdr->data[0]];
+                       size = opt_hdr_v1_size(ohdr) - 12 - 4 * ohdr->data[0];
+                       goto extract;
                }
+
+               ++cur_idx;
        }
 
        if (idx != cur_idx) {
diff --git a/tools/kwbimage.h b/tools/kwbimage.h
index 074bdfbe46..9a949e03c0 100644
--- a/tools/kwbimage.h
+++ b/tools/kwbimage.h
@@ -235,4 +235,62 @@ static inline unsigned int image_version(const void 
*header)
        return ptr[8];
 }
 
+static inline uint32_t opt_hdr_v1_size(const struct opt_hdr_v1 *ohdr)
+{
+       return (ohdr->headersz_msb << 16) | le16_to_cpu(ohdr->headersz_lsb);
+}
+
+static inline int opt_hdr_v1_valid_size(const struct opt_hdr_v1 *ohdr,
+                                       const void *mhdr_end)
+{
+       uint32_t ohdr_size;
+
+       if ((void *)(ohdr + 1) > mhdr_end)
+               return 0;
+
+       ohdr_size = opt_hdr_v1_size(ohdr);
+       if (ohdr_size < 8 || (void *)((uint8_t *)ohdr + ohdr_size) > mhdr_end)
+               return 0;
+
+       return 1;
+}
+
+static inline struct opt_hdr_v1 *opt_hdr_v1_first(void *img) {
+       struct main_hdr_v1 *mhdr;
+
+       if (image_version(img) != 1)
+               return NULL;
+
+       mhdr = img;
+       if (mhdr->ext & 0x1)
+               return (struct opt_hdr_v1 *)(mhdr + 1);
+       else
+               return NULL;
+}
+
+static inline uint8_t *opt_hdr_v1_ext(struct opt_hdr_v1 *cur)
+{
+       uint32_t size = opt_hdr_v1_size(cur);
+
+       return (uint8_t *)cur + size - 4;
+}
+
+static inline struct opt_hdr_v1 *_opt_hdr_v1_next(struct opt_hdr_v1 *cur)
+{
+       return (struct opt_hdr_v1 *)((uint8_t *)cur + opt_hdr_v1_size(cur));
+}
+
+static inline struct opt_hdr_v1 *opt_hdr_v1_next(struct opt_hdr_v1 *cur)
+{
+       if (*opt_hdr_v1_ext(cur) & 0x1)
+               return _opt_hdr_v1_next(cur);
+       else
+               return NULL;
+}
+
+#define for_each_opt_hdr_v1(ohdr, img)         \
+       for ((ohdr) = opt_hdr_v1_first((img));  \
+            (ohdr) != NULL;                    \
+            (ohdr) = opt_hdr_v1_next((ohdr)))
+
 #endif /* _KWBIMAGE_H_ */
-- 
2.31.1

Reply via email to