java/modules/extensions/src/test/java/org/apache/synapse/mediators/throttle/ThrottleMediatorTest.java

/*
 *  Licensed to the Apache Software Foundation (ASF) under one
 *  or more contributor license agreements.  See the NOTICE file
 *  distributed with this work for additional information
 *  regarding copyright ownership.  The ASF licenses this file
 *  to you 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.synapse.mediators.throttle;

import junit.framework.TestCase;
import org.apache.axiom.om.OMAbstractFactory;
import org.apache.axiom.om.OMDocument;
import org.apache.axiom.om.OMElement;
import org.apache.axiom.om.OMXMLBuilderFactory;
import org.apache.axiom.om.OMXMLParserWrapper;
import org.apache.axiom.soap.SOAPEnvelope;
import org.apache.synapse.MessageContext;
import org.apache.synapse.SynapseException;
import org.apache.synapse.config.Entry;
import org.apache.synapse.config.SynapseConfiguration;
import org.apache.synapse.core.SynapseEnvironment;
import org.apache.synapse.core.axis2.Axis2MessageContext;
import org.apache.synapse.core.axis2.Axis2SynapseEnvironment;
import org.apache.synapse.mediators.AbstractMediator;
import org.wso2.throttle.*;

import java.io.StringReader;

/**
 * Throttle Mediator Test - This class test throttling when policy has specified as both of
 * InLine and a registry key
 */

public class ThrottleMediatorTest extends TestCase {
    private static final String REMOTE_ADDR = "Remote_Addr";

    private static final String POLICY = " <wsp:Policy xmlns:wsp=\"http://schemas.xmlsoap.org/ws/2004/09/policy\"\n" +
            "                xmlns:throttle=\"http://www.wso2.org/products/wso2commons/throttle\">\n" +
            "        <throttle:ThrottleAssertion>\n" +
            "            <wsp:All>\n" +
            "                <throttle:ID throttle:type=\"IP\">other</throttle:ID>\n" +
            "                <wsp:ExactlyOne>\n" +
            "                    <wsp:All>\n" +
            "                        <throttle:MaximumCount>8</throttle:MaximumCount>\n" +
            "                        <throttle:UnitTime>800000</throttle:UnitTime>\n" +
            "                        <throttle:ProhibitTimePeriod wsp:Optional=\"true\">10</throttle:ProhibitTimePeriod>\n" +
            "                    </wsp:All>\n" +
            "                </wsp:ExactlyOne>\n" +
            "            </wsp:All>\n" +
            "            <wsp:All>\n" +
            "                <throttle:ID throttle:type=\"IP\">192.168.8.200-192.168.8.222</throttle:ID>\n" +
            "                <wsp:ExactlyOne>\n" +
            "                    <wsp:All>\n" +
            "                        <throttle:MaximumCount>3</throttle:MaximumCount>\n" +
            "                        <throttle:UnitTime>800000</throttle:UnitTime>\n" +
            "                        <throttle:ProhibitTimePeriod wsp:Optional=\"true\">10000</throttle:ProhibitTimePeriod>\n" +
            "                    </wsp:All>\n" +
            "                </wsp:ExactlyOne>\n" +
            "            </wsp:All>\n" +
            "            <wsp:All>\n" +
            "                <throttle:ID throttle:type=\"IP\">192.168.8.201</throttle:ID>\n" +
            "                <wsp:ExactlyOne>\n" +
            "                    <wsp:All>\n" +
            "                        <throttle:MaximumCount>200</throttle:MaximumCount>\n" +
            "                        <throttle:UnitTime>600000</throttle:UnitTime>\n" +
            "                        <throttle:ProhibitTimePeriod wsp:Optional=\"true\"></throttle:ProhibitTimePeriod>\n" +
            "                    </wsp:All>\n" +
            "                </wsp:ExactlyOne>\n" +
            "            </wsp:All>\n" +
            "            <wsp:All>\n" +
            "                <throttle:ID throttle:type=\"IP\">192.168.8.198</throttle:ID>\n" +
            "                <wsp:ExactlyOne>\n" +
            "                    <wsp:All>\n" +
            "                        <throttle:MaximumCount>50</throttle:MaximumCount>\n" +
            "                        <throttle:UnitTime>500000</throttle:UnitTime>\n" +
            "                        <throttle:ProhibitTimePeriod wsp:Optional=\"true\"></throttle:ProhibitTimePeriod>\n" +
            "                    </wsp:All>\n" +
            "                </wsp:ExactlyOne>\n" +
            "            </wsp:All>\n" +
            "        </throttle:ThrottleAssertion>\n" +
            "    </wsp:Policy>";

    private static final String NEW_POLICY = "<wsp:Policy xmlns:wsp=\"http://schemas.xmlsoap.org/ws/2004/09/policy\"\n" +
            "            xmlns:throttle=\"http://www.wso2.org/products/wso2commons/throttle\">\n" +
            "<throttle:MediatorThrottleAssertion>\n" +
            "    <throttle:MaximumConcurrentAccess>10</throttle:MaximumConcurrentAccess>\n" +
            "    <wsp:Policy>\n" +
            "        <throttle:ID throttle:type=\"IP\">other</throttle:ID>\n" +
            "        <wsp:Policy>\n" +
            "            <throttle:Control>\n" +
            "                <wsp:Policy>\n" +
            "                    <throttle:MaximumCount>8</throttle:MaximumCount>\n" +
            "                    <throttle:UnitTime>800000</throttle:UnitTime>\n" +
            "                    <throttle:ProhibitTimePeriod wsp:Optional=\"true\">10\n" +
            "                    </throttle:ProhibitTimePeriod>\n" +
            "                </wsp:Policy>\n" +
            "            </throttle:Control>\n" +
            "        </wsp:Policy>\n" +
            "    </wsp:Policy>\n" +
            "    <wsp:Policy>\n" +
            "        <throttle:ID throttle:type=\"IP\">192.168.8.200-192.168.8.222</throttle:ID>\n" +
            "        <wsp:Policy>\n" +
            "            <throttle:Control>\n" +
            "                <wsp:Policy>\n" +
            "                    <throttle:MaximumCount>3</throttle:MaximumCount>\n" +
            "                    <throttle:UnitTime>800000</throttle:UnitTime>\n" +
            "                    <throttle:ProhibitTimePeriod wsp:Optional=\"true\">10000\n" +
            "                    </throttle:ProhibitTimePeriod>\n" +
            "                </wsp:Policy>\n" +
            "            </throttle:Control>\n" +
            "        </wsp:Policy>\n" +
            "    </wsp:Policy>\n" +
            "    <wsp:Policy>\n" +
            "        <throttle:ID throttle:type=\"IP\">192.168.8.201</throttle:ID>\n" +
            "        <wsp:Policy>\n" +
            "            <throttle:Control>\n" +
            "                <wsp:Policy>\n" +
            "                    <throttle:MaximumCount>200</throttle:MaximumCount>\n" +
            "                    <throttle:UnitTime>600000</throttle:UnitTime>\n" +
            "                    <throttle:ProhibitTimePeriod wsp:Optional=\"true\"></throttle:ProhibitTimePeriod>\n" +
            "                </wsp:Policy>\n" +
            "            </throttle:Control>\n" +
            "        </wsp:Policy>\n" +
            "    </wsp:Policy>\n" +
            "    <wsp:Policy>\n" +
            "        <throttle:ID throttle:type=\"IP\">192.168.8.198</throttle:ID>\n" +
            "        <wsp:Policy>\n" +
            "            <throttle:Control>\n" +
            "                <wsp:Policy>\n" +
            "                    <throttle:MaximumCount>50</throttle:MaximumCount>\n" +
            "                    <throttle:UnitTime>500000</throttle:UnitTime>\n" +
            "                    <throttle:ProhibitTimePeriod wsp:Optional=\"true\"></throttle:ProhibitTimePeriod>\n" +
            "                </wsp:Policy>\n" +
            "            </throttle:Control>\n" +
            "        </wsp:Policy>\n" +
            "    </wsp:Policy>\n" +
            "</throttle:MediatorThrottleAssertion>\n" +
            "</wsp:Policy>";

    public ThrottleMediatorTest() {
        super(ThrottleMediatorTest.class.getName());
    }

    public static MessageContext createLightweightSynapseMessageContext(
            String payload) throws Exception {
        org.apache.axis2.context.MessageContext mc =
                new org.apache.axis2.context.MessageContext();
        SynapseConfiguration config = new SynapseConfiguration();
        SynapseEnvironment env = new Axis2SynapseEnvironment(config);
        MessageContext synMc = new Axis2MessageContext(mc, config, env);
        SOAPEnvelope envelope =
                OMAbstractFactory.getSOAP11Factory().getDefaultEnvelope();
        OMDocument omDoc =
                OMAbstractFactory.getSOAP11Factory().createOMDocument();
        omDoc.addChild(envelope);

        envelope.getBody().addChild(OMXMLBuilderFactory.createOMBuilder(new StringReader(payload)).getDocumentElement());

        synMc.setEnvelope(envelope);
        return synMc;
    }

    public void testMediate() throws Exception {
        OMXMLParserWrapper builder = OMXMLBuilderFactory.createOMBuilder(new StringReader(POLICY));
        ThrottleTestMediator throttleMediator = new ThrottleTestMediator();
        throttleMediator.setPolicyKey("throttlepolicy");
        MessageContext synCtx = createLightweightSynapseMessageContext("<empty/>");
        synCtx.setProperty(REMOTE_ADDR, "192.168.8.212");
        SynapseConfiguration synCfg = new SynapseConfiguration();
        Entry prop = new Entry();
        prop.setKey("throttlepolicy");
        prop.setType(Entry.INLINE_XML);
        prop.setValue(builder.getDocumentElement());
        synCfg.addEntry("throttlepolicy", prop);
        synCtx.setConfiguration(synCfg);
        for (int i = 0; i < 6; i++) {
            try {
                throttleMediator.mediate(synCtx);
                Thread.sleep(1000);
            }
            catch (Exception e) {

                if (i == 3) {
                    assertTrue(e.getMessage().lastIndexOf("IP_BASE") > 0);
                }
                if (i == 4) {
                    assertTrue(e.getMessage().lastIndexOf("IP_BASE") > 0);
                }
                if (i == 5) {
                    assertTrue(e.getMessage().lastIndexOf("IP_BASE") > 0);
                }
            }
        }

    }

    public void testMediateWithInLineXML() throws Exception {
        OMXMLParserWrapper builder = OMXMLBuilderFactory.createOMBuilder(new StringReader(POLICY));
        ThrottleTestMediator throttleMediator = new ThrottleTestMediator();
        throttleMediator.setInLinePolicy(builder.getDocumentElement());
        MessageContext synCtx = createLightweightSynapseMessageContext("<empty/>");
        synCtx.setProperty(REMOTE_ADDR, "192.168.8.212");
        SynapseConfiguration synCfg = new SynapseConfiguration();
        synCtx.setConfiguration(synCfg);
        for (int i = 0; i < 6; i++) {
            try {
                throttleMediator.mediate(synCtx);
                Thread.sleep(1000);
            }
            catch (Exception e) {

                if (i == 3) {
                    assertTrue(e.getMessage().lastIndexOf("IP_BASE") > 0);
                }
                if (i == 4) {
                    assertTrue(e.getMessage().lastIndexOf("IP_BASE") > 0);
                }
                if (i == 5) {
                    assertTrue(e.getMessage().lastIndexOf("IP_BASE") > 0);
                }
            }
        }

    }

    public class ThrottleTestMediator extends AbstractMediator {

        private String policyKey = null;

        Throttle throttle = null;
        private OMElement inLinePolicy = null;

        public boolean mediate(MessageContext synCtx) {

            init(synCtx);
            try {
                return canAccess(synCtx);
            }
            catch (
                    ThrottleException e) {
                throw new SynapseException(e.getMessage());
            }

        }

        protected boolean canAccess(MessageContext synContext)
                throws SynapseException, ThrottleException {

            if (throttle == null) {
                throw new ThrottleException("Can not find a throttle");

            }
            //IP based throttling
            String remoteIP = (String) synContext.getProperty(REMOTE_ADDR);
            if (remoteIP == null) {
                throw new ThrottleException("IP address of the caller can not find - " +
                        "Currently only support caller-IP base access control" +
                        "- Throttling will not happen ");

            } else {
                ThrottleContext throttleContext
                        = throttle.getThrottleContext(ThrottleConstants.IP_BASED_THROTTLE_KEY);
                if (throttleContext == null) {
                    throw new ThrottleException("Can not find a configuration for " +
                            "IP Based Throttle");

                }
                AccessRateController accessController;
                try {
                    accessController = new AccessRateController();
                    boolean canAccess = accessController.canAccess(
                            throttleContext, remoteIP, ThrottleConstants.IP_BASE).isAccessAllowed();
                    if (!canAccess) {
                        throw new SynapseException("Access has currently been denied by" +
                                " the IP_BASE throttle for IP :\t" + remoteIP);
                    }
                    return canAccess;
                }
                catch (ThrottleException e) {
                    return true;
                }
            }
        }

        protected void init(MessageContext synContext) {

            boolean reCreate = false; // It is not need to recreate ,if property is not dyanamic
            OMElement policyOmElement = null;

            if (policyKey != null) {
                Entry entry = synContext.getConfiguration().getEntryDefinition(policyKey);
                if (entry == null) {
                    return;
                }
                Object entryValue = entry.getValue();

                if (!(entryValue instanceof OMElement)) {
                    return;
                }
                // if entry is dynamic, need to check wheather updated or not
                if ((!entry.isCached() || entry.isExpired())) {
                    reCreate = true;
                }
                policyOmElement = (OMElement) entryValue;
            } else if (inLinePolicy != null) {
                policyOmElement = inLinePolicy;
            }
            if (policyOmElement == null) {
                return;
            }
            if (!reCreate) {
                //The first time creation
                if (throttle == null) {
                    createThrottleMetaData(policyOmElement);
                }
            } else {
                createThrottleMetaData(policyOmElement);
            }

        }

        protected void createThrottleMetaData(OMElement policyOmElement) {

            try {
                throttle = ThrottleFactory.createMediatorThrottle(policyOmElement);
            }
            catch (ThrottleException e) {

                throw new SynapseException(e.getMessage());
            }
        }

        public String getType() {
            return "ThrottleMediator";
        }


        public String getPolicyKey() {
            return policyKey;
        }

        public void setPolicyKey(String policyKey) {
            this.policyKey = policyKey;
        }

        public OMElement getInLinePolicy() {
            return inLinePolicy;
        }

        public void setInLinePolicy(OMElement inLinePolicy) {
            this.inLinePolicy = inLinePolicy;
        }
    }
}