Fixes #2503

Author: jfarcand

modules/cpr/src/main/java/org/atmosphere/client/TrackMessageSizeInterceptor.java

@@ -26,6 +26,7 @@
 import org.atmosphere.cpr.AtmosphereResponse;
 import org.atmosphere.cpr.AtmosphereResponseImpl;
 import org.atmosphere.cpr.HeaderConfig;
+import org.atmosphere.interceptor.HeartbeatInterceptor;
 import org.atmosphere.interceptor.InvokationOrder;
 import org.atmosphere.util.IOUtils;
 import org.atmosphere.util.Utils;
@@ -43,7 +44,6 @@
 
 import static org.atmosphere.cpr.ApplicationConfig.EXCLUDED_CONTENT_TYPES;
 import static org.atmosphere.cpr.ApplicationConfig.MESSAGE_DELIMITER;
-
 /**
  * An {@link org.atmosphere.cpr.AtmosphereInterceptor} that add a message size and delimiter.
  * <p/>
@@ -66,6 +66,8 @@ public class TrackMessageSizeInterceptor extends AtmosphereInterceptorAdapter {
 
     private final Interceptor interceptor = new Interceptor();
 
+    private HeartbeatInterceptor heartbeatInterceptor;
+
     @Override
     public void configure(AtmosphereConfig config) {
         String s = config.getInitParameter(MESSAGE_DELIMITER);
@@ -76,6 +78,7 @@ public void configure(AtmosphereConfig config) {
         if (s != null) {
             excludedContentTypes.addAll(Arrays.asList(s.split(",")));
         }
+        heartbeatInterceptor = config.framework().interceptor(HeartbeatInterceptor.class);
     }
 
     /**
@@ -150,6 +153,8 @@ public byte[] transformPayload(AtmosphereResponse response, byte[] responseDraft
                             // This is likely padding written by PaddingAtmosphereInterceptor
                             return responseDraft;
                         }
+                    } else if (isMessageAlreadyEncoded(cb.toString())) {
+                        return responseDraft;
                     }
 
                     AtmosphereResource r = response.resource();
@@ -180,6 +185,15 @@ public byte[] transformPayload(AtmosphereResponse response, byte[] responseDraft
         }
     }
 
+    private boolean isMessageAlreadyEncoded(String message) {
+
+        if (heartbeatInterceptor != null &&
+                message.endsWith(endString + new String(heartbeatInterceptor.getPaddingBytes()))) {
+            return true;
+        }
+        return false;
+    }
+
     @Override
     public PRIORITY priority() {
         return InvokationOrder.BEFORE_DEFAULT;

modules/cpr/src/main/java/org/atmosphere/cpr/AtmosphereFramework.java

@@ -2691,12 +2691,14 @@ public LinkedList<AtmosphereInterceptor> interceptors() {
         return interceptors;
     }
 
-    /**
-     * Set the {@link AnnotationProcessor} class name.
-     *
-     * @param annotationProcessorClassName the {@link AnnotationProcessor} class name.
-     * @return this
-     */
+    public <T extends AtmosphereInterceptor> T interceptor(Class<T> c) {
+        for (AtmosphereInterceptor i : interceptors) {
+            if (c.isInstance(i)) {
+                return c.cast(i);
+            }
+        }
+        return null;
+    }
     public AtmosphereFramework annotationProcessorClassName(String annotationProcessorClassName) {
         this.annotationProcessorClassName = annotationProcessorClassName;
         return this;

modules/cpr/src/main/java/org/atmosphere/interceptor/TrackMessageSizeB64Interceptor.java

@@ -47,12 +47,15 @@ public class TrackMessageSizeB64Interceptor extends AtmosphereInterceptorAdapter
 
     private final Interceptor interceptor = new Interceptor();
 
+    private HeartbeatInterceptor heartbeatInterceptor;
+
     @Override
     public void configure(AtmosphereConfig config) {
         String s = config.getInitParameter(EXCLUDED_CONTENT_TYPES);
         if (s != null) {
             excludedContentTypes.addAll(Arrays.asList(s.split(",")));
         }
+        heartbeatInterceptor = config.framework().interceptor(HeartbeatInterceptor.class);
     }
 
     /**
@@ -99,11 +102,25 @@ public byte[] transformPayload(AtmosphereResponse response, byte[] responseDraft
                     || !excludedContentTypes.contains(response.getContentType().toLowerCase()))) {
                 response.setCharacterEncoding(OUT_ENCODING);
                 String s = Base64.getEncoder().encodeToString(responseDraft);
+
+                if (isMessageAlreadyEncoded(new String(data))) {
+                    logger.trace("Message already encoded {}", s);
+                    return responseDraft;
+                }
+
                 return (s.length() + DELIMITER + s).getBytes(OUT_ENCODING);
             } else {
                 return responseDraft;
             }
+        }
+    }
+
+    public boolean isMessageAlreadyEncoded(String message) {
 
+        if (heartbeatInterceptor != null &&
+                message.endsWith(DELIMITER + new String(heartbeatInterceptor.getPaddingBytes()))) {
+            return true;
         }
+        return false;
     }
 }

modules/cpr/src/test/java/org/atmosphere/cpr/TrackMessageSizeInterceptorTest.java

@@ -26,14 +26,12 @@
 import javax.servlet.ServletConfig;
 import javax.servlet.ServletContext;
 import javax.servlet.ServletException;
-
 import java.io.ByteArrayOutputStream;
 import java.io.IOException;
 import java.io.OutputStream;
 import java.util.Enumeration;
 import java.util.HashMap;
 import java.util.Map;
-import java.util.stream.IntStream;
 
 import static org.mockito.Mockito.mock;
 import static org.testng.Assert.assertEquals;
@@ -115,6 +113,16 @@ public void testTrackMessageSizeB64Disabled() throws Exception {
         testTrackMessageSize(false, new TrackMessageSizeB64Interceptor(), "yoCometyoWebSocket");
     }
 
+    @Test
+    public void testTrackMessageB64AlreadyEncoded() throws Exception {
+        testTrackMessageSize(true, new TrackMessageSizeB64Interceptor(), "1|X", "", "1|X");
+    }
+
+    @Test
+    public void testTrackMessageAlreadyEncoded() throws Exception {
+        testTrackMessageSize(true, new TrackMessageSizeInterceptor(), "1|X", "", "1|X");
+    }
+
     private void testTrackMessageSize(boolean enabled, AtmosphereInterceptor icp, String expected) throws Exception {
         testTrackMessageSize(enabled, icp, "yoComet", "yoWebSocket", expected);
     }