/*
    *  Copyright (c) 2008 Rodrigo Ruiz
    *
    *  Licensed under the Apache License, Version 2.0 (the "License");
    *  you may not use this file except in compliance with the License.
    *  You may obtain a copy of the License at
    *
    *      http://www.apache.org/licenses/LICENSE-2.0
    *
   *  Unless required by applicable law or agreed to in writing, software
   *  distributed under the License is distributed on an "AS IS" BASIS,
   *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   *  See the License for the specific language governing permissions and
   *  limitations under the License.
   */
  package org.apache.ws.addressing.handler;
  
  import java.net.MalformedURLException;
  import java.net.URL;
  import java.util.ArrayList;
  import java.util.Iterator;
  import java.util.List;
  
  import javax.xml.namespace.QName;
  import javax.xml.rpc.Call;
  import javax.xml.rpc.JAXRPCException;
  import javax.xml.rpc.Service;
  import javax.xml.rpc.ServiceException;
  import javax.xml.rpc.ServiceFactory;
  import javax.xml.rpc.handler.MessageContext;
  import javax.xml.rpc.handler.soap.SOAPMessageContext;
  import javax.xml.soap.Node;
  import javax.xml.soap.SOAPBody;
  import javax.xml.soap.SOAPConnection;
  import javax.xml.soap.SOAPConnectionFactory;
  import javax.xml.soap.SOAPElement;
  import javax.xml.soap.SOAPException;
  import javax.xml.soap.SOAPMessage;
  import javax.xml.transform.TransformerFactory;
  
  import org.apache.axis.message.addressing.Action;
  import org.apache.axis.message.addressing.AddressingHeaders;
  import org.apache.axis.message.addressing.AddressingVersion;
  import org.apache.axis.message.addressing.AttributedURI;
  import org.apache.axis.message.addressing.Constants;
  import org.apache.axis.message.addressing.EndpointReference;
  import org.apache.axis.message.addressing.MessageID;
  import org.apache.axis.message.addressing.To;
  import org.apache.axis.message.addressing.util.AddressingUtils;
  import org.apache.axis.types.URI;
  import org.apache.axis.types.URI.MalformedURIException;
  import org.apache.commons.logging.Log;
  import org.apache.commons.logging.LogFactory;
  
  /**
   * A server-side JAX-RPC {@link javax.xml.rpc.handler.Handler} that extracts
   * WS-Addressing headers from incoming SOAP requests and inserts them into
   * outgoing SOAP responses.
   *
   * @author Davanum Srinivas
   * @author Ian P. Springer
   * @version $Revision: 14 $
   */
  public class ServerSideAddressingHandler extends AbstractAddressingHandler {
  
    /**
     * Class logger.
     */
    private static final Log LOG = LogFactory.getLog(ServerSideAddressingHandler.class);
  
    /**
     * Thread local that gives each thread its own TransformerFactory (since it is
     * not thread-safe).
     */
    public static final FactoryThreadLocal TRANSFORMER_FACTORY = new FactoryThreadLocal();
  
    /**
     * {@inheritDoc}
     */
    @Override
    public boolean handleRequest(MessageContext msgContext) {
      SOAPMessageContext soapMsgContext = (SOAPMessageContext) msgContext;
      try {
        SOAPMessage msg = soapMsgContext.getMessage();
        if (msg == null) {
          return CONTINUE_HANDLER_CHAIN_PROCESSING;
        }
        AddressingHeaders headers = new AddressingHeaders(msg.getSOAPPart().getEnvelope(),
          getActor(), true, isRemoveHeadersEnabled(), false, getReferencePropertyQNames());
  
        if (headers.getTo() == null) {
          // not a WS-A request; let it pass thru w/out doing anything
          return CONTINUE_HANDLER_CHAIN_PROCESSING;
        }
        if (headers.getAction() == null) {
          // should we throw a SOAPFaultException here?
          LOG.debug("WSA request to endpoint " + headers.getTo()
            + " is missing the required wsa:Action header.");
        }
       msgContext.setProperty(Constants.ENV_ADDRESSING_REQUEST_HEADERS, headers);
 
       // set the target service based on To header if it hasn't already been
       // determined. NOTE: May want to add an option to override later.
       setTargetService(soapMsgContext, headers);
     } catch (Exception e) {
       if (LOG.isDebugEnabled()) {
         e.printStackTrace();
       }
       throw new JAXRPCException("Unexpected error in handleRequest()", e);
     }
     return CONTINUE_HANDLER_CHAIN_PROCESSING;
   }
 
   /**
    * {@inheritDoc}
    */
   @Override
   public boolean handleResponse(MessageContext ctx) {
     SOAPMessageContext soapMsgContext = (SOAPMessageContext) ctx;
     try {
       SOAPMessage msg = soapMsgContext.getMessage();
       if (msg == null) {
         return CONTINUE_HANDLER_CHAIN_PROCESSING;
       }
 
       AddressingVersion version = AddressingUtils.getAddressingVersion(ctx);
 
       AddressingHeaders reqHeaders = (AddressingHeaders) ctx
         .getProperty(Constants.ENV_ADDRESSING_REQUEST_HEADERS);
 
       if (reqHeaders == null) {
         // not a WS-A response; let it pass thru w/out doing anything
         return CONTINUE_HANDLER_CHAIN_PROCESSING;
       }
       AddressingHeaders resHeaders = AddressingUtils.getResponseHeaders(ctx);
       resHeaders.setSetMustUnderstand(isMustUnderstandEnabled(ctx));
       processFromHeader(reqHeaders, resHeaders);
       processActionHeader(reqHeaders, resHeaders);
       processToHeader(version, reqHeaders, resHeaders);
       processRelatesToHeader(version, reqHeaders, resHeaders);
       processMessageIdHeader(resHeaders);
       resHeaders.toEnvelope(msg.getSOAPPart().getEnvelope(), getActor());
       processReplyToHeader(reqHeaders, soapMsgContext);
     } catch (Exception e) {
       if (LOG.isDebugEnabled()) {
         e.printStackTrace();
       }
       throw new JAXRPCException("Unexpected error in handleResponse()", e);
     }
     return CONTINUE_HANDLER_CHAIN_PROCESSING;
   }
 
   /**
    * {@inheritDoc}
    */
   @Override
   public boolean handleFault(MessageContext ctx) {
     SOAPMessageContext soapMsgContext = (SOAPMessageContext) ctx;
     AddressingVersion version = AddressingUtils.getAddressingVersion(ctx);
 
     try {
       AddressingHeaders reqHeaders = (AddressingHeaders) ctx
         .getProperty(Constants.ENV_ADDRESSING_REQUEST_HEADERS);
       if (reqHeaders == null) {
         // not a WSA fault; let it pass thru w/out doing anything
         return CONTINUE_HANDLER_CHAIN_PROCESSING;
       }
       AddressingHeaders resHeaders = AddressingUtils.getResponseHeaders(ctx);
       processFromHeader(reqHeaders, resHeaders);
       resHeaders.setAction(new Action(version.getFaultActionURI()));
       addRelatesToHeader(version, reqHeaders, resHeaders);
       addMessageIdHeader(resHeaders);
       addAddressingHeadersToSOAPEnvelope(soapMsgContext, resHeaders);
       processFaultToHeader(reqHeaders, soapMsgContext);
     } catch (Exception e) {
       if (LOG.isDebugEnabled()) {
         e.printStackTrace();
       }
       throw new JAXRPCException("Unexpected error in handleFault()", e);
     }
     return CONTINUE_HANDLER_CHAIN_PROCESSING;
   }
 
   /**
    * Forward the SOAP message contained in the specified SOAP message context on
    * to the specified end-point reference.
    *
    * @param ctx Message context to get the message from
    * @param ref End-point reference to forward to
    */
   protected void forwardMessage(SOAPMessageContext ctx, EndpointReference ref) {
     try {
       SOAPConnection soapConn = SOAPConnectionFactory.newInstance().createConnection();
       soapConn.call(ctx.getMessage(), new URL(ref.getAddress().toString()));
       soapConn.close();
       SOAPBody responseBody = ctx.getMessage().getSOAPPart().getEnvelope().getBody();
       removeAllChildElements(responseBody);
     } catch (SOAPException e) {
       throw new JAXRPCException("Failed to forward SOAP message.", e);
     } catch (MalformedURLException e) {
       throw new JAXRPCException("Failed to forward SOAP message.", e);
     }
   }
 
   /**
    * Override this method if you need something other than the default Service.
    * <p>
    * The service returned by this method is used in creating the new Call
    * object. Something like: <br>
    *
    * <pre>
    * Service service = getService(msgContext);
    * Call call = service.createCall()
    * call.setTargetEndpointAddress(toEndPointReference.getAddress().toString());
    * </pre>
    *
    * @param ctx Message context
    * @return Service instance
    */
   protected Service getService(MessageContext ctx) throws ServiceException {
     return ServiceFactory.newInstance().createService(new QName(""));
   }
 
   /**
    * Override this method to prepare the new call, for instance to add
    * properties from the old MessageContext that may be needed by other
    * handlers.
    *
    * @param call Call object about to be invoked
    * @param oldContext MessageContext of the original request/response.
    */
   protected void configureCall(Call call, SOAPMessageContext oldContext) {
     // intentionally empty
   }
 
   /**
    * Can be overridden by subclasses to customize how the wsa:to header is
    * interpreted.
    *
    * @param headers Headers to extract the information from
    * @return The target service name
    * @throws Exception If an error occurs
    */
   protected String getTargetServiceName(AddressingHeaders headers) throws Exception {
     To toURI = headers.getTo();
     if (toURI == null) {
       return null;
     }
     String to = toURI.getPath();
     if (to == null) {
       return null;
     }
     // set the target service
     return (to.substring(to.lastIndexOf('/') + 1));
   }
 
   /**
    * Platform-specific subclasses can optionally implement this method.
    *
    * @param ctx Context to get information from
    * @param headers Headers to use
    * @throws Exception If an error occurs
    */
   protected void setTargetService(SOAPMessageContext ctx, AddressingHeaders headers)
     throws Exception {
     // intentionally empty
   }
 
   /**
    * Processes the request From header and sets the appropriate values at the
    * corresponding response headers.
    *
    * @param reqHeaders Request header container
    * @param resHeaders Response header container
    */
   private void processFromHeader(AddressingHeaders reqHeaders,
     AddressingHeaders resHeaders) {
     EndpointReference fromEPR = resHeaders.getFrom();
     if (fromEPR == null) {
       To toURI = reqHeaders.getTo();
       if (toURI != null) {
         fromEPR = new EndpointReference(toURI);
         fromEPR.setProperties(reqHeaders.getReferenceProperties());
         resHeaders.setFrom(fromEPR);
       }
     }
   }
 
   /**
    * Processes the request Action header and sets the appropriate values at the
    * corresponding response headers.
    *
    * @param reqHeaders Request header container
    * @param resHeaders Response header container
    */
   private void processActionHeader(AddressingHeaders reqHeaders,
     AddressingHeaders resHeaders) throws URI.MalformedURIException {
     Action action = resHeaders.getAction();
     if (action == null) {
       // not set - try request headers
       action = reqHeaders.getAction();
       if (action != null) {
         resHeaders
           .setAction(new Action(new URI(action.toString() + "Response")));
       }
     }
   }
 
   /**
    * Processes the request To header and sets the appropriate values at the
    * corresponding response headers.
    *
    * @param version    WS-Addressing version to use
    * @param reqHeaders Request header container
    * @param resHeaders Response header container
    */
   private void processToHeader(AddressingVersion version, AddressingHeaders reqHeaders,
     AddressingHeaders resHeaders) throws MalformedURIException {
     if (resHeaders.getFrom() == null && reqHeaders.getFrom() != null) {
       resHeaders.setTo(reqHeaders.getFrom().getAddress());
     } else {
       resHeaders.setTo(new To(version.getAnonymousRoleURI()));
     }
   }
 
   /**
    * Processes the request RelatesTo header and sets the appropriate values at
    * the corresponding response headers.
    *
    * @param version    WS-Addressing version to use
    * @param reqHeaders Request header container
    * @param resHeaders Response header container
    */
   private void processRelatesToHeader(AddressingVersion version,
     AddressingHeaders reqHeaders, AddressingHeaders resHeaders)
     throws MalformedURIException {
     MessageID msgID = reqHeaders.getMessageID();
     if (msgID != null) {
       resHeaders.addRelatesTo(msgID.toString(), version.getResponseRelationshipType());
     }
   }
 
   /**
    * Sets the appropriate values at the MessageID response header.
    *
    * @param resHeaders Response header container
    */
   private void processMessageIdHeader(AddressingHeaders resHeaders)
     throws URI.MalformedURIException {
     resHeaders.setMessageID(createMessageID());
   }
 
   /**
    * Processes the request ReplyTo header and sets the appropriate values at the
    * message context.
    *
    * @param reqHeaders Request header container
    * @param ctx        Message context information
    */
   private void processReplyToHeader(AddressingHeaders reqHeaders,
     SOAPMessageContext ctx) throws Exception {
     EndpointReference replyTo = reqHeaders.getReplyTo();
     if (replyTo != null) {
       AttributedURI address = replyTo.getAddress();
       if (address != null) {
         String uri = address.toString();
         if (uri != null && !uri.equals(AddressingUtils.getAnonymousRoleURI(ctx))) {
           forwardMessage(ctx, replyTo);
         }
       }
     }
   }
 
   /**
    * Adds the response headers to the current response SOAP envelope.
    *
    * @param soapMsgContext Context information (for retrieving the SOAP
    *        envelope)
    * @param resHeaders Response header container
    * @throws Exception If an error occurs
    */
   private void addAddressingHeadersToSOAPEnvelope(SOAPMessageContext soapMsgContext,
     AddressingHeaders resHeaders) throws Exception {
     SOAPMessage msg = soapMsgContext.getMessage();
     if (msg == null) {
       throw new JAXRPCException("Unable to get response message from message context.");
     }
     resHeaders.toEnvelope(msg.getSOAPPart().getEnvelope());
   }
 
   /**
    * Processes the request FaultTo header and sets the appropriate values at the
    * message context.
    *
    * @param reqHeaders Request header container
    * @param soapMsgContext Message context
    */
   private void processFaultToHeader(AddressingHeaders reqHeaders,
     SOAPMessageContext soapMsgContext) throws Exception {
     EndpointReference faultTo = reqHeaders.getFaultTo();
     if (faultTo != null) {
       AttributedURI address = faultTo.getAddress();
       if (address != null && address.toString() != null) {
         forwardMessage(soapMsgContext, faultTo);
       }
     }
   }
 
   /**
    * Processes the request RelatesTo header and sets the appropriate values at
    * the corresponding response headers.
    *
    * @param version    WS-Addressing version to use
    * @param reqHeaders Request header container
    * @param resHeaders Response header container
    */
   private void addRelatesToHeader(AddressingVersion version,
     AddressingHeaders reqHeaders, AddressingHeaders resHeaders)
     throws MalformedURIException {
     MessageID msgID = reqHeaders.getMessageID();
     if (msgID != null) {
       resHeaders.addRelatesTo(msgID.toString(), version.getResponseRelationshipType());
     }
   }
 
   /**
    * Adds the MessageID header.
    *
    * @param resHeaders Response header container
    * @throws URI.MalformedURIException If an error occurs
    */
   private void addMessageIdHeader(AddressingHeaders resHeaders)
     throws URI.MalformedURIException {
     MessageID msgID = new MessageID(new URI("uuid:" + generateUUId()));
     resHeaders.setMessageID(msgID);
   }
 
   /**
    * Removes all child elements from the specified
    * {@link javax.xml.soap.SOAPElement}.
    *
    * @param soapElem Element to strip
    */
   protected static void removeAllChildElements(SOAPElement soapElem) {
     Iterator<?> iter = soapElem.getChildElements();
     // NOTE: Convert iterator to list to avoid ConcurrentModificationExceptions
     // caused by modifying items in an iterator during iteration
     List<?> children = toList(iter);
     for (int i = 0; i < children.size(); i++) {
       Node child = (Node) children.get(i);
       if (child.getParentElement() != null) {
         child.detachNode();
         child.recycleNode();
       }
     }
   }
 
   /**
    * Iterates over the specified iterator and puts all its elements into a list.
    *
    * @param <T> Iterator type
    * @param iter Iterator
    * @return List
    */
   private static <T> List<T> toList(Iterator<T> iter) {
     List<T> list = new ArrayList<T>();
     if (iter != null) {
       while (iter.hasNext()) {
         list.add(iter.next());
       }
     }
     return list;
   }
 
   /**
    * ThreadLocalExtension type.
    *
    * @author Rodrigo Ruiz
    */
   private static final class FactoryThreadLocal extends ThreadLocal<TransformerFactory> {
     /**
      * {@inheritDoc}
      */
     @Override
     protected synchronized TransformerFactory initialValue() {
       return TransformerFactory.newInstance();
     }
   }
 }