Merge pull request #1279 from microsoft/littleaj/fix-object-deserialization

Object deserialization improvements

Author: littleaj

ApplicationInsightsInternalLogger/src/main/java/com/microsoft/applicationinsights/internal/logger/InternalLogger.java

@@ -204,6 +204,11 @@ public void stop() {
             loggingLevel = LoggingLevel.OFF;
         }
     }
+    /* VisibleForTesting */
+    void reset() {
+        INSTANCE.stop();
+        INSTANCE.initialized = false;
+    }
 
     public boolean isTraceEnabled() {
         return loggingLevel.getValue() <= LoggingLevel.TRACE.getValue();

core/src/main/java/com/microsoft/applicationinsights/internal/channel/common/TransmissionFileSystemOutput.java

@@ -23,14 +23,15 @@
 
 import java.io.File;
 import java.io.InputStream;
+import java.io.InvalidClassException;
 import java.io.ObjectInput;
 import java.io.FileInputStream;
 import java.io.BufferedInputStream;
 import java.io.ObjectInputStream;
 import java.io.FileNotFoundException;
 import java.io.IOException;
 import java.io.FileOutputStream;
-import java.io.OutputStream;
+import java.io.ObjectStreamClass;
 import java.io.BufferedOutputStream;
 import java.io.ObjectOutput;
 import java.io.ObjectOutputStream;
@@ -244,7 +245,7 @@ private Optional<Transmission> loadTransmission(File file) {
         if (file == null) {
             return Optional.absent();
         }
-        try (ObjectInput input = new ObjectInputStream(new BufferedInputStream(new FileInputStream(file)))) {
+        try (ObjectInput input = new SafeObjectInputStream(new BufferedInputStream(new FileInputStream(file)))) {
             transmission = (Transmission)input.readObject();
         } catch (FileNotFoundException e) {
             InternalLogger.INSTANCE.error("Failed to load transmission, file not found, exception: %s", e.toString());
@@ -257,6 +258,21 @@ private Optional<Transmission> loadTransmission(File file) {
         return Optional.fromNullable(transmission);
     }
 
+    private final static class SafeObjectInputStream extends ObjectInputStream {
+
+        public SafeObjectInputStream(InputStream in) throws IOException {
+            super(in);
+        }
+
+        protected Class<?> resolveClass(ObjectStreamClass desc) throws IOException, ClassNotFoundException {
+            if (!desc.getName().equals(Transmission.class.getName()) && !desc.getName().equals(byte[].class.getName())) {
+                throw new InvalidClassException("Cannot deserialize "+desc.getName());
+            } else {
+                return super.resolveClass(desc);
+            }
+        }
+    }
+
     private boolean renameToPermanentName(File tempTransmissionFile) {
         File transmissionFile = new File(folder, FilenameUtils.getBaseName(tempTransmissionFile.getName()) + TRANSMISSION_FILE_EXTENSION);
         try {

core/src/test/java/com/microsoft/applicationinsights/internal/channel/common/ActiveTransmissionLoaderTest.java

@@ -23,10 +23,13 @@
 
 import java.io.File;
 import java.io.IOException;
+import java.util.HashMap;
 import java.util.concurrent.TimeUnit;
 
 import com.microsoft.applicationinsights.internal.channel.TransmissionDispatcher;
-import org.junit.Test;
+import com.microsoft.applicationinsights.internal.logger.InternalLogger;
+import com.microsoft.applicationinsights.internal.logger.LoggerTestHelper;
+import org.junit.*;
 import org.mockito.Mockito;
 
 import org.apache.commons.io.FileUtils;
@@ -38,6 +41,16 @@
 public class ActiveTransmissionLoaderTest {
     private final static String TEMP_TEST_FOLDER = "TransmissionTests";
 
+    @BeforeClass
+    public static void initLogger() {
+        InternalLogger.INSTANCE.initialize("CONSOLE", new HashMap<String, String>());
+    }
+
+    @AfterClass
+    public static void tearDownLogger() {
+        LoggerTestHelper.resetInternalLogger();
+    }
+
     @Test(expected = NullPointerException.class)
     public void testNullFileSystem() throws Exception {
         new ActiveTransmissionLoader(null, Mockito.mock(TransmissionDispatcher.class), mockStateFetcher(), 1);

core/src/test/java/com/microsoft/applicationinsights/internal/logger/LoggerTestHelper.java

@@ -0,0 +1,7 @@
+package com.microsoft.applicationinsights.internal.logger;
+
+public class LoggerTestHelper {
+    public static void resetInternalLogger() {
+        InternalLogger.INSTANCE.reset();
+    }
+}

gradle/common-java.gradle

@@ -110,7 +110,7 @@ if (hasProperty("JDKToUse")) {
         options.with {
             fork = true
             forkOptions.executable = javaExecutables.javac
-            forkOptions.javaHome = JDKToUse
+            forkOptions.javaHome = file(JDKToUse)
         }
     }
     tasks.withType(Javadoc) {