From 4112294fe1aeb512f72b06a3e9bec940c2d6b938 Mon Sep 17 00:00:00 2001
From: Alessio Soldano <asoldano@redhat.com>
Date: Mon, 21 Nov 2016 13:42:54 +0100
Subject: [PATCH] [JBWS-4040] Fix regression with XMLMessage instances

---
 .../GracefulShutdownInterceptor.java          | 50 +++++++++++--------
 1 file changed, 29 insertions(+), 21 deletions(-)

diff --git a/modules/server/src/main/java/org/jboss/wsf/stack/cxf/interceptor/GracefulShutdownInterceptor.java b/modules/server/src/main/java/org/jboss/wsf/stack/cxf/interceptor/GracefulShutdownInterceptor.java
index 2220a7852..8f711bbba 100644
--- a/modules/server/src/main/java/org/jboss/wsf/stack/cxf/interceptor/GracefulShutdownInterceptor.java
+++ b/modules/server/src/main/java/org/jboss/wsf/stack/cxf/interceptor/GracefulShutdownInterceptor.java
@@ -30,12 +30,12 @@
 import javax.xml.namespace.QName;
 
 import org.apache.cxf.binding.soap.SoapMessage;
-import org.apache.cxf.binding.soap.interceptor.AbstractSoapInterceptor;
 import org.apache.cxf.binding.soap.interceptor.ReadHeadersInterceptor;
-import org.apache.cxf.common.i18n.Message;
 import org.apache.cxf.common.logging.LogUtils;
 import org.apache.cxf.headers.Header;
 import org.apache.cxf.interceptor.Fault;
+import org.apache.cxf.message.Message;
+import org.apache.cxf.phase.AbstractPhaseInterceptor;
 import org.apache.cxf.phase.Phase;
 import org.apache.cxf.transport.http.AbstractHTTPDestination;
 import org.jboss.wsf.spi.deployment.Deployment;
@@ -49,7 +49,7 @@
  * @since 4-Nov-2016
  *
  */
-public class GracefulShutdownInterceptor extends AbstractSoapInterceptor
+public class GracefulShutdownInterceptor extends AbstractPhaseInterceptor<Message>
 {
    private static final Logger LOG = LogUtils.getL7dLogger(GracefulShutdownInterceptor.class);
 
@@ -60,38 +60,46 @@ public GracefulShutdownInterceptor()
    }
 
    @Override
-   public void handleMessage(SoapMessage message) throws Fault
+   public void handleMessage(Message message) throws Fault
    {
       ServletRequest req = (ServletRequest) message.get(AbstractHTTPDestination.HTTP_REQUEST);
       if ("true".equals(req.getAttribute("org.wildfly.suspended")))
       {
-         if (!message.hasHeaders())
+         if (message instanceof SoapMessage)
          {
-            throw createFault();
-         }
-         else
-         {
-            Deployment dep = message.getExchange().get(Endpoint.class).getService().getDeployment();
-            RejectionRule rr = dep.getAttachment(RejectionRule.class);
-            if (rr != null)
+            SoapMessage soapMessage = (SoapMessage)message;
+            if (!soapMessage.hasHeaders())
             {
-               List<Header> headers = message.getHeaders();
-               Map<QName, Object> m = new HashMap<>();
-               for (Header header : headers)
-               {
-                  m.put(header.getName(), header.getObject());
-               }
-               if (rr.rejectMessage(m))
+               throw createFault();
+            }
+            else
+            {
+               Deployment dep = soapMessage.getExchange().get(Endpoint.class).getService().getDeployment();
+               RejectionRule rr = dep.getAttachment(RejectionRule.class);
+               if (rr != null)
                {
-                  throw createFault();
+                  List<Header> headers = soapMessage.getHeaders();
+                  Map<QName, Object> m = new HashMap<>();
+                  for (Header header : headers)
+                  {
+                     m.put(header.getName(), header.getObject());
+                  }
+                  if (rr.rejectMessage(m))
+                  {
+                     throw createFault();
+                  }
                }
             }
          }
+         else
+         {
+            throw createFault();
+         }
       }
    }
 
    private Fault createFault() {
-      Fault f = new Fault(new Message("Server is suspended", LOG));
+      Fault f = new Fault(new org.apache.cxf.common.i18n.Message("Server is suspended", LOG));
       f.setStatusCode(503);
       return f;
    }