Merge pull request #3570 from IntegraSources/mtp-changes

mtp: fixes in add_string

2 files changed, 43 insertions, 19 deletions

diff --git a/src/class/mtp/mtp.h b/src/class/mtp/mtp.h
index 236cf98e0..7b22837cd 100644
--- a/src/class/mtp/mtp.h
+++ b/src/class/mtp/mtp.h
@@ -798,45 +798,65 @@ TU_ATTR_ALWAYS_INLINE static inline uint32_t mtp_container_add_array(mtp_contain
 }
 
 TU_ATTR_ALWAYS_INLINE static inline uint32_t mtp_container_add_string(mtp_container_info_t* p_container, uint16_t* utf16) {
-  uint8_t count = 0;
+  uint32_t count = 0;
   while (utf16[count] != 0u) {
     count++;
   }
-  const uint32_t added_len = 1u + (uint32_t) count * 2u;
-  TU_ASSERT(p_container->header->len + added_len < CFG_TUD_MTP_EP_BUFSIZE, 0);
+  // MTP strings store length in a single uint8_t, including trailing null.
+  TU_ASSERT(count < UINT8_MAX, 0);
+  count++;
+
   uint8_t* buf = p_container->payload + p_container->header->len - sizeof(mtp_container_header_t);
 
-  *buf++ = count;
+  if (count == 1) {
+    // empty string (size only): single zero byte
+    TU_ASSERT(p_container->header->len + 1 < CFG_TUD_MTP_EP_BUFSIZE, 0);
+    *buf = 0;
+    p_container->header->len++;
+    return 1u;
+  }
+
+  const uint32_t added_len = 1u + count * 2u;
+  TU_ASSERT(p_container->header->len + added_len < CFG_TUD_MTP_EP_BUFSIZE, 0);
+
+  *buf++ = (uint8_t) count;
   p_container->header->len++;
 
-  memcpy(buf, utf16, 2u * (uint32_t) count);
+  memcpy(buf, utf16, 2u * count);
   p_container->header->len += 2u * count;
 
   return added_len;
 }
 
 TU_ATTR_ALWAYS_INLINE static inline uint32_t mtp_container_add_cstring(mtp_container_info_t* p_container, const char* str) {
-  const uint8_t len = (uint8_t) (strlen(str) + 1); // include null
-  TU_ASSERT(p_container->header->len + 1 + 2 * len < CFG_TUD_MTP_EP_BUFSIZE, 0);
+  const size_t cstr_len = strlen(str);
+  // MTP strings store length in a single uint8_t, including trailing null.
+  TU_ASSERT(cstr_len < UINT8_MAX, 0);
+
+  const uint32_t count = (uint32_t) cstr_len + 1u; // include null
   uint8_t* buf = p_container->payload + p_container->header->len - sizeof(mtp_container_header_t);
 
-  if (len == 1) {
-    // empty string (null only): single zero byte
+  if (count == 1u) {
+    // empty string (size only): single zero byte
+    TU_ASSERT(p_container->header->len + 1 < CFG_TUD_MTP_EP_BUFSIZE, 0);
     *buf = 0;
     p_container->header->len++;
     return 1u;
-  } else {
-    *buf++ = len;
-    p_container->header->len++;
+  }
 
-    for (uint8_t i = 0; i < len; i++) {
-      buf[0] = str[i];
-      buf[1] = 0;
-      buf += 2;
-      p_container->header->len += 2;
-    }
-    return 1u + 2u * len;
+  const uint32_t added_len = 1u + 2u * count;
+  TU_ASSERT(p_container->header->len + added_len < CFG_TUD_MTP_EP_BUFSIZE, 0);
+
+  *buf++ = (uint8_t) count;
+  p_container->header->len++;
+
+  for (uint32_t i = 0; i < count; i++) {
+    *buf++ = (uint8_t) str[i];
+    *buf++ = 0;
   }
+  p_container->header->len += 2u * count;
+
+  return added_len;
 }
 
 TU_ATTR_ALWAYS_INLINE static inline uint32_t mtp_container_add_uint8(mtp_container_info_t* p_container, uint8_t data) {
diff --git a/src/class/mtp/mtp_device.h b/src/class/mtp/mtp_device.h
index 6cce7efbb..f2c5cef7f 100644
--- a/src/class/mtp/mtp_device.h
+++ b/src/class/mtp/mtp_device.h
@@ -78,6 +78,10 @@ typedef struct {
     mtp_auint16_t(_capture_count) capture_formats; \
     mtp_auint16_t(_playback_count) playback_formats; \
     /* string fields will be added using append function */ \
+    /* mtp_string_t() Manufacturer */ \
+    /* mtp_string_t() Model */ \
+    /* mtp_string_t() Device Version */ \
+    /* mtp_string_t() Serial Number */ \
   }
 
 typedef MTP_DEVICE_INFO_STRUCT( //-V2586 [MISRA-C-18.7] Flexible array members should not be declared