changes for gzipped userdata check

Signed-off-by: Abhishek Kumar <abhishek.mrt22@gmail.com>

Author: shwstppr

engine/userdata/cloud-init/src/main/java/org/apache/cloudstack/userdata/CloudInitUserDataProvider.java

@@ -16,27 +16,31 @@
 // under the License.
 package org.apache.cloudstack.userdata;
 
-import com.cloud.utils.component.AdapterBase;
-import com.cloud.utils.exception.CloudRuntimeException;
-import org.apache.commons.collections.CollectionUtils;
-import org.apache.commons.lang3.StringUtils;
-import org.apache.log4j.Logger;
-
-import javax.mail.BodyPart;
-import javax.mail.MessagingException;
-import javax.mail.Multipart;
-import javax.mail.Session;
-import javax.mail.internet.MimeBodyPart;
-import javax.mail.internet.MimeMessage;
-import javax.mail.internet.MimeMultipart;
 import java.io.ByteArrayInputStream;
 import java.io.ByteArrayOutputStream;
 import java.io.IOException;
+import java.nio.charset.StandardCharsets;
 import java.util.Arrays;
 import java.util.List;
 import java.util.Map;
 import java.util.Properties;
 import java.util.stream.Collectors;
+import java.util.zip.GZIPInputStream;
+
+import javax.mail.BodyPart;
+import javax.mail.MessagingException;
+import javax.mail.Multipart;
+import javax.mail.Session;
+import javax.mail.internet.MimeBodyPart;
+import javax.mail.internet.MimeMessage;
+import javax.mail.internet.MimeMultipart;
+
+import org.apache.commons.collections.CollectionUtils;
+import org.apache.commons.lang3.StringUtils;
+import org.apache.log4j.Logger;
+
+import com.cloud.utils.component.AdapterBase;
+import com.cloud.utils.exception.CloudRuntimeException;
 
 public class CloudInitUserDataProvider extends AdapterBase implements UserDataProvider {
 
@@ -65,7 +69,22 @@ public String getName() {
         return "cloud-init";
     }
 
+    protected boolean isGZipped(String userdata) {
+        if (StringUtils.isEmpty(userdata)) {
+            return false;
+        }
+        byte[] data = userdata.getBytes(StandardCharsets.ISO_8859_1);
+        if (data.length < 2) {
+            return false;
+        }
+        int magic = data[0] & 0xff | ((data[1] << 8) & 0xff00);
+        return magic == GZIPInputStream.GZIP_MAGIC;
+    }
+
     protected String extractUserDataHeader(String userdata) {
+        if (isGZipped(userdata)) {
+            throw new CloudRuntimeException("Gzipped user data can not be used together with other user data formats");
+        }
         List<String> lines = Arrays.stream(userdata.split("\n"))
                 .filter(x -> (x.startsWith("#") && !x.startsWith("##")) || (x.startsWith("Content-Type:")))
                 .collect(Collectors.toList());

engine/userdata/cloud-init/src/test/java/org/apache/cloudstack/userdata/CloudInitUserDataProviderTest.java

@@ -16,22 +16,28 @@
 // under the License.
 package org.apache.cloudstack.userdata;
 
-import com.cloud.utils.exception.CloudRuntimeException;
+import java.io.ByteArrayOutputStream;
+import java.io.IOException;
+import java.nio.charset.StandardCharsets;
+import java.util.zip.GZIPOutputStream;
+
 import org.junit.Assert;
 import org.junit.Test;
 
+import com.cloud.utils.exception.CloudRuntimeException;
+
 public class CloudInitUserDataProviderTest {
 
     private final CloudInitUserDataProvider provider = new CloudInitUserDataProvider();
+    private final static String CLOUD_CONFIG_USERDATA = "## template: jinja\n" +
+            "#cloud-config\n" +
+            "runcmd:\n" +
+            "   - echo 'TestVariable {{ ds.meta_data.variable1 }}' >> /tmp/variable\n" +
+            "   - echo 'Hostname {{ ds.meta_data.public_hostname }}' > /tmp/hostname";
 
     @Test
     public void testGetUserDataFormatType() {
-        String userdata = "## template: jinja\n" +
-                "#cloud-config\n" +
-                "runcmd:\n" +
-                "   - echo 'TestVariable {{ ds.meta_data.variable1 }}' >> /tmp/variable\n" +
-                "   - echo 'Hostname {{ ds.meta_data.public_hostname }}' > /tmp/hostname";
-        CloudInitUserDataProvider.FormatType type = provider.getUserDataFormatType(userdata);
+        CloudInitUserDataProvider.FormatType type = provider.getUserDataFormatType(CLOUD_CONFIG_USERDATA);
         Assert.assertEquals(CloudInitUserDataProvider.FormatType.CLOUD_CONFIG, type);
     }
 
@@ -94,4 +100,40 @@ public void testAppendUserDataInvalidUserData() {
                 "date > /provisioned";
         provider.appendUserData(templateData, vmData);
     }
+
+    @Test
+    public void testIsGzippedUserDataWithCloudConfigData() {
+        Assert.assertFalse(provider.isGZipped(CLOUD_CONFIG_USERDATA));
+    }
+
+    private String createGzipDataAsString() throws IOException {
+        byte[] input = CLOUD_CONFIG_USERDATA.getBytes(StandardCharsets.ISO_8859_1);
+
+        ByteArrayOutputStream arrayOutputStream = new ByteArrayOutputStream();
+        GZIPOutputStream outputStream = new GZIPOutputStream(arrayOutputStream);
+        outputStream.write(input,0, input.length);
+        outputStream.close();
+
+        return arrayOutputStream.toString(StandardCharsets.ISO_8859_1);
+    }
+
+    @Test
+    public void testIsGzippedUserDataWithValidGzipData() {
+        try {
+            String gzipped = createGzipDataAsString();
+            Assert.assertTrue(provider.isGZipped(gzipped));
+        } catch (IOException e) {
+            Assert.fail(e.getMessage());
+        }
+    }
+
+    @Test(expected = CloudRuntimeException.class)
+    public void testAppendUserDataWithGzippedData() {
+        try {
+            provider.appendUserData(CLOUD_CONFIG_USERDATA, createGzipDataAsString());
+            Assert.fail("Gzipped data shouldn't be appended with other data");
+        } catch (IOException e) {
+            Assert.fail("Exception encountered: " + e.getMessage());
+        }
+    }
 }