# ThrottleHandler.java

    /*
    * Copyright 2005,2006 WSO2, Inc. http://wso2.com
    *
    * 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.synapse.commons.throttle.module.handler;

    import org.apache.axis2.AxisFault;
    import org.apache.axis2.context.ConfigurationContext;
    import org.apache.axis2.context.MessageContext;
    import org.apache.axis2.description.AxisOperation;
    import org.apache.axis2.description.AxisService;
    import org.apache.axis2.handlers.AbstractHandler;
    import org.apache.axis2.transport.http.HTTPConstants;
    import org.apache.commons.logging.Log;
    import org.apache.commons.logging.LogFactory;
    import org.apache.synapse.commons.throttle.core.AccessRateController;
    import org.apache.synapse.commons.throttle.core.ConcurrentAccessController;
    import org.apache.synapse.commons.throttle.core.RoleBasedAccessRateController;
    import org.apache.synapse.commons.throttle.core.Throttle;
    import org.apache.synapse.commons.throttle.core.ThrottleConstants;
    import org.apache.synapse.commons.throttle.core.ThrottleConfiguration;
    import org.apache.synapse.commons.throttle.core.ThrottleException;
    import org.apache.synapse.commons.throttle.core.ThrottleContext;
    import org.apache.synapse.commons.throttle.core.AccessInformation;
    import org.apache.synapse.commons.throttle.module.utils.impl.DummyHandler;
    import org.apache.synapse.commons.throttle.module.utils.impl.DummyAuthenticator;
    import org.apache.synapse.commons.throttle.module.utils.StatCollector;
    import javax.cache.Cache;
    import javax.cache.CacheManager;
    import javax.cache.Caching;
    import javax.servlet.http.HttpServletRequest;
    import javax.xml.namespace.QName;
    import java.util.Map;

    public abstract class ThrottleHandler extends AbstractHandler {

        private static Log log = LogFactory.getLog(ThrottleHandler.class.getName());
        /* The AccessRateController - control(limit) access for a remote caller */
        private AccessRateController accessRateController;

        private RoleBasedAccessRateController roleBasedAccessController;

        public static final String THROTTLING_CACHE_MANAGER = "throttling.cache.manager";

        public static final String THROTTLING_CACHE = "throttling.cache";

        private boolean debugOn;

        public ThrottleHandler() {
            this.debugOn = log.isDebugEnabled();
            this.accessRateController = new AccessRateController();
            this.roleBasedAccessController = new RoleBasedAccessRateController();
        }

        /**
         * @return int - indicates the type of the throttle according to the scope
         */
        protected abstract int getThrottleType();

        /**
         * Loads a throttle metadata for a particular throttle type
         *
         * @param messageContext - The messageContext
         * @param throttleType   - The type of throttle
         * @return IPBaseThrottleConfiguration     - The IPBaseThrottleConfiguration - load from AxisConfiguration
         * @throws ThrottleException Throws if the throttle type is unsupported
         */

        public Throttle loadThrottle(MessageContext messageContext,
                                     int throttleType) throws ThrottleException {

            Throttle throttle = null;
            ConfigurationContext configContext = messageContext.getConfigurationContext();
            //the Parameter which hold throttle ipbase object
            // to get throttles map from the configuration context

            Map throttles = (Map) configContext.getPropertyNonReplicable(ThrottleConstants.THROTTLES_MAP);
            if (throttles == null) {
                if (debugOn) {
                    log.debug("Couldn't find throttles object map .. thottlling will not be occurred ");
                }
                return null;
            }
            switch (throttleType) {
                case ThrottleConstants.GLOBAL_THROTTLE: {
                    throttle =
                            (Throttle) throttles.get(ThrottleConstants.GLOBAL_THROTTLE_KEY);
                    break;
                }
                case ThrottleConstants.OPERATION_BASED_THROTTLE: {
                    AxisOperation axisOperation = messageContext.getAxisOperation();
                    if (axisOperation != null) {
                        QName opName = axisOperation.getName();
                        if (opName != null) {
                            AxisService service = (AxisService) axisOperation.getParent();
                            if (service != null) {
                                String currentServiceName = service.getName();
                                if (currentServiceName != null) {
                                    throttle =
                                            (Throttle) throttles.get(currentServiceName +
                                                    opName.getLocalPart());
                                }
                            }
                        }
                    } else {
                        if (debugOn) {
                            log.debug("Couldn't find axis operation ");
                        }
                        return null;
                    }
                    break;
                }
                case ThrottleConstants.SERVICE_BASED_THROTTLE: {
                    AxisService axisService = messageContext.getAxisService();
                    if (axisService != null) {
                        throttle =
                                (Throttle) throttles.get(axisService.getName());
                    } else {
                        if (debugOn) {
                            log.debug("Couldn't find axis service ");
                        }
                        return null;
                    }
                    break;
                }
                default: {
                    throw new ThrottleException("Unsupported Throttle type");
                }
            }
            return throttle;
        }

        /**
         * processing through the throttle
         * 1) concurrent throttling
         * 2) access rate based throttling - domain or ip
         *
         * @param throttle       The Throttle object - holds all configuration and state data
         *                       of the throttle
         * @param messageContext The MessageContext , that holds all data per message basis
         * @throws AxisFault         Throws when access must deny for caller
         * @throws ThrottleException    ThrottleException
         */
        public void process(Throttle throttle,
                            MessageContext messageContext) throws ThrottleException, AxisFault {

            String throttleId = throttle.getId();
            ConfigurationContext cc = messageContext.getConfigurationContext();

            // acquiring  cache manager.
            Cache<String, ConcurrentAccessController> cache;
            CacheManager cacheManager = Caching.getCacheManagerFactory().getCacheManager(THROTTLING_CACHE_MANAGER);
            if (cacheManager != null) {
                cache = cacheManager.getCache(THROTTLING_CACHE);
            } else {
                cache = Caching.getCacheManager().getCache(THROTTLING_CACHE);
            }
            if (log.isDebugEnabled()) {
                log.debug("created throttling cache : " + cache);
            }
            // Get the concurrent access controller
            ConcurrentAccessController cac;
            String key = null;
            key = ThrottleConstants.THROTTLE_PROPERTY_PREFIX + throttleId
                    + ThrottleConstants.CAC_SUFFIX;
            cac = cache.get(key);

            // check for concurrent access
            boolean canAccess = doConcurrentThrottling(cac, messageContext);

            if (canAccess) {
                // if the concurrent access is success then
                // do the access rate based throttling

                if (messageContext.getFLOW() == MessageContext.IN_FLOW) {
                    //gets the remote caller domain name
                    String domain = null;
                    HttpServletRequest request =
                            (HttpServletRequest) messageContext.getPropertyNonReplicable(
                                    HTTPConstants.MC_HTTP_SERVLETREQUEST);
                    if (request != null) {
                        domain = request.getRemoteHost();
                    }

                    // Domain name based throttling
                    //check whether a configuration has been defined for this domain name or not
                    String callerId = null;
                    if (domain != null) {
                        //loads the ThrottleContext
                        ThrottleContext throttleCtxt =
                                throttle.getThrottleContext(ThrottleConstants.DOMAIN_BASED_THROTTLE_KEY);
                        if (throttleCtxt != null) {
                            //Loads the ThrottleConfiguration
                            ThrottleConfiguration throttleConfig = throttleCtxt.getThrottleConfiguration();
                            if (throttleConfig != null) {
                                //check for configuration for this caller
                                callerId = throttleConfig.getConfigurationKeyOfCaller(domain);
                                if (callerId != null) {
                                    // If this is a clustered env.
                                    throttleCtxt.setThrottleId(throttleId);
                                    AccessInformation infor =
                                            accessRateController.canAccess(throttleCtxt, callerId,
                                                    ThrottleConstants.DOMAIN_BASE);
                                    StatCollector.collect(infor, domain, ThrottleConstants.DOMAIN_BASE);

                                    //check for the permission for access
                                    if (!infor.isAccessAllowed()) {

                                        //In the case of both of concurrency throttling and
                                        //rate based throttling have enabled ,
                                        //if the access rate less than maximum concurrent access ,
                                        //then it is possible to occur death situation.To avoid that reset,
                                        //if the access has denied by rate based throttling
                                        if (cac != null) {
                                            cac.incrementAndGet();
                                            cache.put(key, cac);
                                            if (debugOn) {
                                                log.debug("Added the state of ConcurrentAccessController " +
                                                        "to cache with key : " + key);
                                            }
                                        }
                                        throw new AxisFault(" Access deny for a " +
                                                "caller with Domain " + domain + " " +
                                                " : Reason : " + infor.getFaultReason());
                                    }
                                } else {
                                    if (debugOn) {
                                        log.debug("Could not find the Throttle Context for domain-Based " +
                                                "Throttling for domain name " + domain + " Throttling for this " +
                                                "domain name may not be configured from policy");
                                    }
                                }
                            }
                        }
                    } else {
                        if (debugOn) {
                            log.debug("Could not find the domain of the caller - IP-based throttling may occur");
                        }
                    }

                    //IP based throttling - Only if there is no configuration for caller domain name

                    if (callerId == null) {
                        String ip = (String) messageContext.getProperty(MessageContext.REMOTE_ADDR);
                        if (ip != null) {
                            // loads IP based throttle context
                            ThrottleContext context =
                                    throttle.getThrottleContext(ThrottleConstants.IP_BASED_THROTTLE_KEY);
                            if (context != null) {
                                //Loads the ThrottleConfiguration
                                ThrottleConfiguration config = context.getThrottleConfiguration();
                                if (config != null) {
                                    // check for configuration for this ip
                                    callerId = config.getConfigurationKeyOfCaller(ip);
                                    if (callerId != null) {
                                            context.setThrottleId(throttleId);
                                        AccessInformation infor =
                                                accessRateController.canAccess(context, callerId,
                                                        ThrottleConstants.IP_BASE);
                                        // check for the permission for access
                                        StatCollector.collect(infor,ip,ThrottleConstants.IP_BASE);
                                        if (!infor.isAccessAllowed()) {

                                            //In the case of both of concurrency throttling and
                                            //rate based throttling have enabled ,
                                            //if the access rate less than maximum concurrent access ,
                                            //then it is possible to occur death situation.To avoid that reset,
                                            //if the access has denied by rate based throttling
                                            if (cac != null) {
                                                cac.incrementAndGet();
                                                // set back if this is a clustered env
                                                    cache.put(key, cac);
                                                    if(debugOn) {
                                                        log.debug("Added the state of ConcurrentAccessController " +
                                                                "to cache with key : " + key);
                                                }
                                            }
                                            throw new AxisFault(" Access deny for a " +
                                                    "caller with IP " + ip + " " +
                                                    " : Reason : " + infor.getFaultReason());
                                        }
                                    }
                                }
                            } else {
                                if (debugOn) {
                                    log.debug("Could not find the throttle Context for IP-Based throttling");
                                }
                            }
                        } else {
                            if (debugOn) {
                                log.debug("Could not find the IP address of the caller " +
                                        "- throttling will not occur");
                            }
                        }
                    }
                }
                // all the replication functionality of the access rate based throttling handles by itself
                // just replicate the current state of ConcurrentAccessController
                if (cac != null) {
                    cache.put(key, cac);
                    if(debugOn) {
                        log.debug("Added the state of ConcurrentAccessController " +
                                "to cache with key : " + key);
                    }
                }

                //finally engage rolebased access throttling if available
                doRoleBasedAccessThrottling(throttle, messageContext);
            } else {
                //replicate the current state of ConcurrentAccessController
                if (cac != null) {
                    cache.put(key, cac);
                    if(debugOn) {
                        log.debug("Added the state of ConcurrentAccessController " +
                                "to cache with key : " + key);
                    }
                }
                throw new AxisFault("Access has currently been denied since " +
                        " maximum concurrent access have exceeded");
            }

        }

        /**
         * Helper method for handling concurrent throttling
         *
         * @param concurrentAccessController ConcurrentAccessController
         * @param messageContext             MessageContext - message level states
         * @return true if access is allowed through concurrent throttling ,o.w false
         */
        private boolean doConcurrentThrottling(ConcurrentAccessController concurrentAccessController, MessageContext messageContext) {

            boolean canAccess = true;
            int available;

            if (concurrentAccessController != null) {
                if (messageContext.getFLOW() == MessageContext.IN_FLOW) {
                    available = concurrentAccessController.getAndDecrement();
                    canAccess = available > 0;
                    if (debugOn) {
                        log.debug("Concurrency Throttle : Access " + (canAccess ? "allowed" : "denied") +
                                " :: " + available + " of available of " +
                                concurrentAccessController.getLimit() + " connections");
                    }
                    if (debugOn) {
                        if (!canAccess) {
                            log.debug("Concurrency Throttle : Access has currently been denied since allowed" +
                                    " maximum concurrent access have exceeded");
                        }
                    }
                } else if (messageContext.getFLOW() == MessageContext.OUT_FLOW) {
                    available = concurrentAccessController.incrementAndGet();
                    if (debugOn) {
                        log.debug("Concurrency Throttle : Connection returned" +
                                " :: " + available + " of available of "
                                + concurrentAccessController.getLimit() + " connections");
                    }
                }
            }
            return canAccess;
        }

        /**
         * Helper method for handling role based Access throttling
         *
         * @param messageContext             MessageContext - message level states
         * @return true if access is allowed through concurrent throttling ,o.w false
         */
        private boolean doRoleBasedAccessThrottling(Throttle throttle, MessageContext messageContext) throws
                AxisFault,ThrottleException {

            boolean canAccess = true;
            if (throttle.getThrottleContext(ThrottleConstants.ROLE_BASED_THROTTLE_KEY) == null) {
                //if no role base throttle config return immediately
                return canAccess;
            }
            ConfigurationContext cc = messageContext.getConfigurationContext();
            String throttleId = throttle.getId();

            // acquiring  cache manager.
            Cache<String, ConcurrentAccessController> cache;
            CacheManager cacheManager = Caching.getCacheManagerFactory().getCacheManager(THROTTLING_CACHE_MANAGER);
            if (cacheManager != null) {
                cache = cacheManager.getCache(THROTTLING_CACHE);
            } else {
                cache = Caching.getCacheManager().getCache(THROTTLING_CACHE);
            }
            if (log.isDebugEnabled()) {
                log.debug("created throttling cache : " + cache);
            }
            String key = null;
            ConcurrentAccessController cac = null;
            key = ThrottleConstants.THROTTLE_PROPERTY_PREFIX + throttleId
                    + ThrottleConstants.CAC_SUFFIX;
            cac = cache.get(key);

            if (messageContext.getFLOW() == MessageContext.IN_FLOW) {
                //gets the remote caller role name
                String consumerKey = null;
                boolean isAuthenticated = false;
                String roleID = null;
                HttpServletRequest request =
                        (HttpServletRequest) messageContext.getPropertyNonReplicable(
                                HTTPConstants.MC_HTTP_SERVLETREQUEST);
                if (request != null) {
                    String oAuthHeader = request.getHeader("OAuth");
    //                consumerKey = Utils.extractCustomerKeyFromAuthHeader(oAuthHeader);
    //                roleID = Utils.extractCustomerKeyFromAuthHeader(oAuthHeader);
                    DummyAuthenticator authFuture = new DummyAuthenticator(oAuthHeader);
                    consumerKey = authFuture.getAPIKey();
                    new DummyHandler().authenticateUser(authFuture);
                    roleID = (String) authFuture.getAuthorizedRoles().get(0);
                    isAuthenticated = authFuture.isAuthenticated();
                }

                if(!isAuthenticated){
                    throw new AxisFault(" Access deny for a " +
                            "caller with consumer Key: " + consumerKey + " " +
                            " : Reason : Authentication failure");

                }
                // Domain name based throttling
                //check whether a configuration has been defined for this role name or not
                String consumerRoleID = null;
                if (consumerKey != null && isAuthenticated) {
                    //loads the ThrottleContext
                    ThrottleContext context =
                            throttle.getThrottleContext(ThrottleConstants.ROLE_BASED_THROTTLE_KEY);
                    if (context != null) {
                        //Loads the ThrottleConfiguration
                        ThrottleConfiguration config = context.getThrottleConfiguration();
                        if (config != null) {
                            //check for configuration for this caller
                            consumerRoleID = config.getConfigurationKeyOfCaller(roleID);
                            if (consumerRoleID != null) {
                                context.setThrottleId(throttleId);
                                AccessInformation infor =
                                        roleBasedAccessController.canAccess(context, consumerKey,
                                                consumerRoleID);
                                StatCollector.collect(infor, consumerKey, ThrottleConstants.ROLE_BASE);
                                //check for the permission for access
                                if (!infor.isAccessAllowed()) {

                                    //In the case of both of concurrency throttling and
                                    //rate based throttling have enabled ,
                                    //if the access rate less than maximum concurrent access ,
                                    //then it is possible to occur death situation.To avoid that reset,
                                    //if the access has denied by rate based throttling
                                    if (cac != null) {
                                        cac.incrementAndGet();

                                                cache.put(key, cac);
                                                if(debugOn) {
                                                    log.debug("Added the state of ConcurrentAccessController " +
                                                            "to cache with key : " + key);
                                                }
                                    }
                                    throw new AxisFault(" Access deny for a " +
                                            "caller with Domain " + consumerKey + " " +
                                            " : Reason : " + infor.getFaultReason());
                                }
                            } else {
                                if (debugOn) {
                                    log.debug("Could not find the Throttle Context for role-Based " +
                                            "Throttling for role name " + consumerKey + " Throttling for this " +
                                            "role name may not be configured from policy");
                                }
                            }
                        }
                    }
                } else {
                    if (debugOn) {
                        log.debug("Could not find the role of the caller - role based throttling NOT applied");
                    }
                }
            }
            return canAccess;
        }

        public InvocationResponse invoke(MessageContext msgContext) throws AxisFault {
            //Load throttle
            try {
                Throttle throttle = loadThrottle(msgContext, getThrottleType());
                if (throttle != null) {
                    process(throttle, msgContext);
                }
            } catch (ThrottleException e) {
                log.error(e.getMessage());
                throw new AxisFault(e.getMessage());
            }
            return InvocationResponse.CONTINUE;
        }

    }