From 7994ac0b1924cf0cc3e3e99ea6f11842f56bc4c8 Mon Sep 17 00:00:00 2001
From: Joshua Harlow <harlowja@yahoo-inc.com>
Date: Sun, 26 Aug 2012 20:52:41 -0700
Subject: Fix tests running and add in a check on the content type before we
 look into the payload as well as make the skip test a function that the
 datasource module can also use.

---
 cloudinit/sources/__init__.py | 17 +++++++++++------
 cloudinit/user_data.py        | 33 +++++++++++++++++++++++----------
 2 files changed, 34 insertions(+), 16 deletions(-)

(limited to 'cloudinit')

diff --git a/cloudinit/sources/__init__.py b/cloudinit/sources/__init__.py
index a3f4af1e..d49b67b2 100644
--- a/cloudinit/sources/__init__.py
+++ b/cloudinit/sources/__init__.py
@@ -81,8 +81,7 @@ class DataSource(object):
         # headers, if not just skip this....
         launch_idxs = 0
         for part in processed_ud.walk():
-            # multipart/* are just containers
-            if part.get_content_maintype() == 'multipart':
+            if ud.is_skippable(part):
                 continue
             launch_idx_h = part.get('Launch-Index', None)
             if launch_idx_h is not None:
@@ -94,17 +93,23 @@ class DataSource(object):
         # that have some other garbage that we don't know what to do with
         accumulating_msg = MIMEMultipart()
         tot_attached = 0
+        tot_processed = 0
         for part in processed_ud.walk():
-            # multipart/* are just containers
-            if part.get_content_maintype() == 'multipart':
+            if ud.is_skippable(part):
                 continue
             try:
+                tot_processed += 1
                 launch_idx_h = part.get('Launch-Index', None)
                 if launch_idx_h is None or int(launch_idx_h) == int(idx):
                     accumulating_msg.attach(part)
                     tot_attached += 1
-            except:
-                # If any int conversion fails (or other error), keep the part
+                else:
+                    LOG.debug(("Discarding multipart message %s, "
+                               "launch-index provided destined for %s "
+                               "and not %s"),
+                               tot_processed, launch_idx_h, idx)
+            except (TypeError, ValueError):
+                # If any int conversion fails keep the message
                 accumulating_msg.attach(part)
                 tot_attached += 1
         accumulating_msg[ud.ATTACHMENT_FIELD] = str(tot_attached)
diff --git a/cloudinit/user_data.py b/cloudinit/user_data.py
index 244e9223..d104d237 100644
--- a/cloudinit/user_data.py
+++ b/cloudinit/user_data.py
@@ -52,6 +52,9 @@ ARCHIVE_UNDEF_TYPE = "text/cloud-config"
 # Msg header used to track attachments
 ATTACHMENT_FIELD = 'Number-Attachments'
 
+# Only the following content types can have there launch index examined
+CAN_HAVE_LAUNCH_INDEX = ["text/cloud-config", "text/cloud-config-archive"]
+
 
 class UserDataProcessor(object):
     def __init__(self, paths):
@@ -64,8 +67,7 @@ class UserDataProcessor(object):
 
     def _process_msg(self, base_msg, append_msg):
         for part in base_msg.walk():
-            # multipart/* are just containers
-            if part.get_content_maintype() == 'multipart':
+            if is_skippable(part):
                 continue
 
             ctype = None
@@ -99,13 +101,16 @@ class UserDataProcessor(object):
     def _attach_launch_index(self, msg):
         header_idx = msg.get('Launch-Index', None)
         payload_idx = None
-        try:
-            payload = util.load_yaml(msg.get_payload(decode=True))
-            if payload:
-                payload_idx = payload.get('launch-index')
-        except:
-            pass
-        # Header overrides contents...
+        if msg.get_content_type() in CAN_HAVE_LAUNCH_INDEX:
+            try:
+                # See if it has a launch-index field
+                # that might affect the final header
+                payload = util.load_yaml(msg.get_payload(decode=True))
+                if payload:
+                    payload_idx = payload.get('launch-index')
+            except:
+                pass
+        # Header overrides contents, for now (?) or the other way around?
         if header_idx is not None:
             payload_idx = header_idx
         # Nothing found in payload, use header (if anything there)
@@ -114,7 +119,7 @@ class UserDataProcessor(object):
         if payload_idx is not None:
             try:
                 msg.add_header('Launch-Index', str(int(payload_idx)))
-            except:
+            except (ValueError, TypeError):
                 pass
 
     def _get_include_once_filename(self, entry):
@@ -241,6 +246,14 @@ class UserDataProcessor(object):
         self._multi_part_count(outer_msg, part_count + 1)
 
 
+def is_skippable(part):
+    # multipart/* are just containers
+    part_maintype = part.get_content_maintype() or ''
+    if part_maintype.lower() == 'multipart':
+        return True
+    return False
+
+
 # Coverts a raw string into a mime message
 def convert_string(raw_data, headers=None):
     if not raw_data:
-- 
cgit v1.2.3