Allowing multiple @OnWebSocketFrame annotations, but preventing duplicate declarations on the same frameType

jetty-websocket/websocket-core/src/main/java/org/eclipse/jetty/websocket/annotations/EventMethod.java

@@ -103,4 +103,14 @@ public class EventMethod
             LOG.warn("Cannot call method {} on {} with {}",method,pojo,args,e);
         }
     }
+
+    protected Method getMethod()
+    {
+        return method;
+    }
+
+    protected Class<?>[] getParamTypes()
+    {
+        return this.paramTypes;
+    }
 }

jetty-websocket/websocket-core/src/main/java/org/eclipse/jetty/websocket/annotations/EventMethods.java

@@ -1,5 +1,12 @@
 package org.eclipse.jetty.websocket.annotations;
 
+import java.util.HashMap;
+import java.util.Map;
+
+import org.eclipse.jetty.util.StringUtil;
+import org.eclipse.jetty.websocket.api.InvalidWebSocketException;
+import org.eclipse.jetty.websocket.frames.BaseFrame;
+
 /**
  * A representation of the methods available to call for a particular class.
  * <p>
@@ -9,12 +16,14 @@ public class EventMethods
 {
     private Class<?> pojoClass;
     private boolean isAnnotated = false;
-    public EventMethod onConnect = EventMethod.NOOP;
+    public EventMethod onConnect = null;
-    public EventMethod onClose = EventMethod.NOOP;
+    public EventMethod onClose = null;
-    public EventMethod onBinary = EventMethod.NOOP;
+    public EventMethod onBinary = null;
-    public EventMethod onText = EventMethod.NOOP;
+    public EventMethod onText = null;
-    public EventMethod onFrame = EventMethod.NOOP;
+    public EventMethod onException = null;
-    public EventMethod onException = EventMethod.NOOP;
+
+    // special case, multiple methods allowed
+    private Map<Class<? extends BaseFrame>, EventMethod> onFrames = new HashMap<Class<? extends BaseFrame>, EventMethod>();
 
     public EventMethods(Class<?> pojoClass, boolean annotated)
     {
@@ -22,6 +31,30 @@ public class EventMethods
         this.isAnnotated = annotated;
     }
 
+    public void addOnFrame(EventMethod eventMethod)
+    {
+        Class<?> paramTypes[] = eventMethod.getParamTypes();
+        @SuppressWarnings("unchecked")
+        Class<? extends BaseFrame> frameType = (Class<? extends BaseFrame>)((paramTypes.length == 1)?paramTypes[0]:paramTypes[1]);
+
+        if (onFrames.containsKey(frameType))
+        {
+            // Attempt to add duplicate frame type (a no-no)
+            StringBuilder err = new StringBuilder();
+            err.append("Duplicate Frame Type declaration on ");
+            err.append(eventMethod.getMethod());
+            err.append(StringUtil.__LINE_SEPARATOR);
+
+            EventMethod dup = onFrames.get(frameType);
+            err.append("Type ").append(frameType.getSimpleName()).append(" previously declared at ");
+            err.append(dup.getMethod());
+
+            throw new InvalidWebSocketException(err.toString());
+        }
+
+        onFrames.put(frameType,eventMethod);
+    }
+
     @Override
     public boolean equals(Object obj)
     {
@@ -52,6 +85,16 @@ public class EventMethods
         return true;
     }
 
+    public EventMethod getOnFrame(Class<? extends BaseFrame> frameType)
+    {
+        return onFrames.get(frameType);
+    }
+
+    public Map<Class<? extends BaseFrame>, EventMethod> getOnFrames()
+    {
+        return onFrames;
+    }
+
     public Class<?> getPojoClass()
     {
         return pojoClass;

jetty-websocket/websocket-core/src/main/java/org/eclipse/jetty/websocket/annotations/EventMethodsCache.java

@@ -255,9 +255,8 @@ public class EventMethodsCache
 
             if (method.getAnnotation(OnWebSocketFrame.class) != null)
             {
-                assertUnset(events.onFrame,OnWebSocketFrame.class,pojo,method);
                 assertValidParams(pojo,method,OnWebSocketFrame.class,validFrameParams);
-                events.onFrame = new EventMethod(pojo,method);
+                events.addOnFrame(new EventMethod(pojo,method));
                 continue;
             }
 

jetty-websocket/websocket-core/src/test/java/org/eclipse/jetty/websocket/annotations/BadDuplicateFrameSocket.java

@@ -0,0 +1,35 @@
+package org.eclipse.jetty.websocket.annotations;
+
+import org.eclipse.jetty.websocket.frames.BaseFrame;
+import org.eclipse.jetty.websocket.frames.TextFrame;
+
+@WebSocket
+public class BadDuplicateFrameSocket
+{
+    /**
+     * The most basic frame type
+     */
+    @OnWebSocketFrame
+    public void frameMe(BaseFrame frame)
+    {
+        /* ignore */
+    }
+
+    /**
+     * Should allow for a more specific frame type as well.
+     */
+    @OnWebSocketFrame
+    public void messageMe(TextFrame frame)
+    {
+        /* ignore */
+    }
+
+    /**
+     * This is a duplicate frame type (should throw an exception attempting to use)
+     */
+    @OnWebSocketFrame
+    public void textMe(TextFrame frame)
+    {
+        /* ignore */
+    }
+}

jetty-websocket/websocket-core/src/test/java/org/eclipse/jetty/websocket/annotations/EventMethodsCacheTest.java

@@ -2,6 +2,9 @@ package org.eclipse.jetty.websocket.annotations;
 
 import static org.hamcrest.Matchers.*;
 
+import org.eclipse.jetty.websocket.api.InvalidWebSocketException;
+import org.eclipse.jetty.websocket.frames.BaseFrame;
+import org.eclipse.jetty.websocket.frames.TextFrame;
 import org.junit.Assert;
 import org.junit.Test;
 
@@ -9,12 +12,31 @@ public class EventMethodsCacheTest
 {
     private void assertHasEventMethod(String message, EventMethod actual)
     {
-        Assert.assertNotSame(message + "Event method should have been discovered",actual,EventMethod.NOOP);
+        Assert.assertThat(message + "Event method should have been discovered",actual,notNullValue());
     }
 
     private void assertNoEventMethod(String message, EventMethod actual)
     {
-        Assert.assertEquals(message + "Event method should have been NOOP",actual,EventMethod.NOOP);
+        Assert.assertThat(message + "Event method should have been NOOP",actual,nullValue());
+    }
+
+    /**
+     * Test Case for bad declaration (duplicate frame type methods)
+     */
+    @Test
+    public void testDiscoverBadDuplicateFrameSocket()
+    {
+        EventMethodsCache cache = new EventMethodsCache();
+        try
+        {
+            // Should toss exception
+            cache.getMethods(BadDuplicateFrameSocket.class);
+        }
+        catch (InvalidWebSocketException e)
+        {
+            // Validate that we have clear error message to the developer
+            Assert.assertThat(e.getMessage(),containsString("Duplicate Frame Type"));
+        }
     }
 
     /**
@@ -32,8 +54,9 @@ public class EventMethodsCacheTest
         assertHasEventMethod("MyEchoSocket.onClose",methods.onClose);
         assertHasEventMethod("MyEchoSocket.onConnect",methods.onConnect);
         assertNoEventMethod("MyEchoSocket.onException",methods.onException);
-        assertNoEventMethod("MyEchoSocket.onFrame",methods.onFrame);
         assertHasEventMethod("MyEchoSocket.onText",methods.onText);
+
+        Assert.assertThat("MyEchoSocket.getOnFrames()",methods.getOnFrames().size(),is(0));
     }
 
     /**
@@ -51,8 +74,9 @@ public class EventMethodsCacheTest
         assertNoEventMethod("MyStatelessEchoSocket.onClose",methods.onClose);
         assertNoEventMethod("MyStatelessEchoSocket.onConnect",methods.onConnect);
         assertNoEventMethod("MyStatelessEchoSocket.onException",methods.onException);
-        assertNoEventMethod("MyStatelessEchoSocket.onFrame",methods.onFrame);
         assertHasEventMethod("MyStatelessEchoSocket.onText",methods.onText);
+
+        Assert.assertThat("MyEchoSocket.getOnFrames()",methods.getOnFrames().size(),is(0));
     }
 
     /**
@@ -70,8 +94,9 @@ public class EventMethodsCacheTest
         assertNoEventMethod("NoopSocket.onClose",methods.onClose);
         assertNoEventMethod("NoopSocket.onConnect", methods.onConnect);
         assertNoEventMethod("NoopSocket.onException",methods.onException);
-        assertNoEventMethod("NoopSocket.onFrame",methods.onFrame);
         assertNoEventMethod("NoopSocket.onText",methods.onText);
+
+        Assert.assertThat("MyEchoSocket.getOnFrames()",methods.getOnFrames().size(),is(0));
     }
 
     /**
@@ -89,7 +114,10 @@ public class EventMethodsCacheTest
         assertNoEventMethod("MyEchoSocket.onClose",methods.onClose);
         assertNoEventMethod("MyEchoSocket.onConnect",methods.onConnect);
         assertNoEventMethod("MyEchoSocket.onException",methods.onException);
-        assertHasEventMethod("MyEchoSocket.onFrame",methods.onFrame);
         assertNoEventMethod("MyEchoSocket.onText",methods.onText);
+
+        Assert.assertThat("MyEchoSocket.getOnFrames()",methods.getOnFrames().size(),is(2));
+        assertHasEventMethod("MyEchoSocket.onFrame(BaseFrame)",methods.getOnFrame(BaseFrame.class));
+        assertHasEventMethod("MyEchoSocket.onFrame(BaseFrame)",methods.getOnFrame(TextFrame.class));
     }
 }

jetty-websocket/websocket-core/src/test/java/org/eclipse/jetty/websocket/annotations/FrameSocket.java

@@ -1,13 +1,26 @@
 package org.eclipse.jetty.websocket.annotations;
 
 import org.eclipse.jetty.websocket.frames.BaseFrame;
+import org.eclipse.jetty.websocket.frames.TextFrame;
 
 @WebSocket
 public class FrameSocket
 {
+    /**
+     * The most basic frame type
+     */
     @OnWebSocketFrame
     public void frameMe(BaseFrame frame)
     {
         /* ignore */
     }
+
+    /**
+     * Should allow for a more specific frame type as well.
+     */
+    @OnWebSocketFrame
+    public void textMe(TextFrame frame)
+    {
+        /* ignore */
+    }
 }