//
//  ========================================================================
//  Copyright (c) 1995-2017 Mort Bay Consulting Pty. Ltd.
//  ------------------------------------------------------------------------
//  All rights reserved. This program and the accompanying materials
//  are made available under the terms of the Eclipse Public License v1.0
//  and Apache License v2.0 which accompanies this distribution.
//
//      The Eclipse Public License is available at
//      http://www.eclipse.org/legal/epl-v10.html
//
//      The Apache License v2.0 is available at
//      http://www.opensource.org/licenses/apache2.0.php
//
//  You may elect to redistribute this code under either of these licenses.
//  ========================================================================
//

package org.eclipse.jetty.server.handler;

import java.io.Closeable;
import java.io.IOException;
import java.net.InetAddress;
import java.net.InetSocketAddress;
import java.util.ArrayDeque;
import java.util.Deque;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;
import java.util.concurrent.ExecutionException;

import javax.servlet.AsyncContext;
import javax.servlet.ServletException;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;

import org.eclipse.jetty.http.HostPortHttpField;
import org.eclipse.jetty.http.HttpField;
import org.eclipse.jetty.http.HttpFields;
import org.eclipse.jetty.http.HttpHeader;
import org.eclipse.jetty.http.QuotedCSV;
import org.eclipse.jetty.server.ForwardedRequestCustomizer;
import org.eclipse.jetty.server.Request;
import org.eclipse.jetty.util.IncludeExcludeSet;
import org.eclipse.jetty.util.InetAddressSet;
import org.eclipse.jetty.util.StringUtil;
import org.eclipse.jetty.util.annotation.ManagedAttribute;
import org.eclipse.jetty.util.annotation.ManagedOperation;
import org.eclipse.jetty.util.annotation.Name;
import org.eclipse.jetty.util.log.Log;
import org.eclipse.jetty.util.log.Logger;
import org.eclipse.jetty.util.thread.Locker;


/**
 * <p>Handler to limit the threads per IP address for DOS protection</p>
 * <p>The ThreadLimitHandler applies a limit to the number of Threads 
 * that can be used simultaneously per remote IP address.
 * </p>
 * <p>The handler makes a determination of the remote IP separately to
 * any that may be made by the {@link ForwardedRequestCustomizer} or similar:
 * <ul>
 * <li>This handler will use either only a single style
 * of forwarded header.   This is on the assumption that a trusted local proxy
 * will produce only a single forwarded header and that any additional
 * headers are likely from untrusted client side proxies.</li>
 * <li>If multiple instances of a forwarded header are provided, this
 * handler will use the right-most instance, which will have been set from
 * the trusted local proxy</li>
 * </ul>
 * Requests in excess of the limit will be asynchronously suspended until
 * a thread is available.  
 * <p>This is a simpler alternative to DosFilter</p>
 */
public class ThreadLimitHandler extends HandlerWrapper
{
    private static final Logger LOG = Log.getLogger(ThreadLimitHandler.class);

    private final static String REMOTE = "o.e.j.s.h.TLH.REMOTE";
    private final static String PERMIT = "o.e.j.s.h.TLH.PASS";
    private final boolean _rfc7239;
    private final String _forwardedHeader;
    private final IncludeExcludeSet<String, InetAddress> _includeExcludeSet = new IncludeExcludeSet<>(InetAddressSet.class);
    private final ConcurrentMap<String, Remote> _remotes = new ConcurrentHashMap<>();
    private volatile boolean _enabled;
    private int _threadLimit=10;

    public ThreadLimitHandler()
    {
        this(null,false);
    }
    
    public ThreadLimitHandler(@Name("forwardedHeader") String forwardedHeader)
    {
        this(forwardedHeader,HttpHeader.FORWARDED.is(forwardedHeader));
    }

    public ThreadLimitHandler(@Name("forwardedHeader") String forwardedHeader, @Name("rfc7239")  boolean rfc7239)
    {
        super();
        _rfc7239 = rfc7239;
        _forwardedHeader = forwardedHeader;
        _enabled = true;
    }
    
    @Override
    protected void doStart() throws Exception
    {
        super.doStart();
        LOG.info(String.format("ThreadLimitHandler enable=%b limit=%d include=%s",_enabled,_threadLimit,_includeExcludeSet));
    }

    @ManagedAttribute("true if this handler is enabled")
    public boolean isEnabled()
    {
        return _enabled;
    }

    public void setEnabled(boolean enabled)
    {
        _enabled = enabled;
        LOG.info(String.format("ThreadLimitHandler enable=%b limit=%d include=%s",_enabled,_threadLimit,_includeExcludeSet));
    }

    @ManagedAttribute("The maximum threads that can be dispatched per remote IP")
    public int getThreadLimit()
    {
        return _threadLimit;
    }

    public void setThreadLimit(int threadLimit)
    {
        if (threadLimit<=0)
            throw new IllegalArgumentException("limit must be >0");
        _threadLimit = threadLimit;
    }
    
    @ManagedOperation("Include IP in thread limits")
    public void include(String inetAddressPattern)
    {
        _includeExcludeSet.include(inetAddressPattern);
    }

    @ManagedOperation("Exclude IP from thread limits")
    public void exclude(String inetAddressPattern)
    {
        _includeExcludeSet.exclude(inetAddressPattern);
    }
    
    @Override
    public void handle(String target, Request baseRequest, HttpServletRequest request, HttpServletResponse response) throws IOException, ServletException
    {
        // Allow ThreadLimit to be enabled dynamically without restarting server
        if (!_enabled)
        {
            // if disabled, handle normally
            super.handle(target,baseRequest,request,response);
        }
        else
        {
            // Get the remote address of the request
            Remote remote = getRemote(baseRequest);
            if (remote==null)
            {
                // if remote is not known, handle normally
                super.handle(target,baseRequest,request,response);
            }
            else
            {
                // Do we already have a future permit from a previous invocation?
                Closeable permit = (Closeable)baseRequest.getAttribute(PERMIT);
                try
                {
                    if (permit!=null)
                    {
                        // Yes, remove it from any future async cycles.
                        baseRequest.removeAttribute(PERMIT);
                    }
                    else
                    {
                        // No, then lets try to acquire one
                        CompletableFuture<Closeable> future_permit=remote.acquire();

                        // Did we get a permit?
                        if (future_permit.isDone())
                        {
                            // yes
                            permit=future_permit.get();
                        }
                        else
                        {
                            if (LOG.isDebugEnabled())
                                LOG.debug("Threadlimited {} {}",remote,target);
                            // No, lets asynchronously suspend the request
                            AsyncContext async = baseRequest.startAsync();
                            // let's never timeout the async.  If this is a DOS, then good to make them wait, if this is not
                            // then give them maximum time to get a thread.
                            async.setTimeout(0);

                            // dispatch the request when we do eventually get a pass
                            future_permit.thenAccept(c->
                            {
                                baseRequest.setAttribute(PERMIT,c);
                                async.dispatch();
                            });
                            return;
                        }
                    }

                    // Use the permit
                    super.handle(target,baseRequest,request,response);
                }
                catch (InterruptedException | ExecutionException e)
                {
                    throw new ServletException(e);
                }
                finally
                {
                    if (permit!=null)
                        permit.close();
                }
            }
        }
    }
    
    protected int getThreadLimit(String ip)
    {
        if (!_includeExcludeSet.isEmpty())
        {
            try
            {
                if (!_includeExcludeSet.test(InetAddress.getByName(ip)))
                {
                    LOG.debug("excluded {}",ip);
                    return 0;
                }
            }
            catch(Exception e)
            {
                LOG.ignore(e);
            }
        }
        return _threadLimit;
    }

    protected Remote getRemote(Request baseRequest)
    {
        Remote remote = (Remote)baseRequest.getAttribute(REMOTE);
        if (remote!=null)
            return remote;
        
        String ip=getRemoteIP(baseRequest);
        LOG.debug("ip={}",ip);
        if (ip==null)
            return null;
        
        int limit = getThreadLimit(ip);
        if (limit<=0)
            return null;

        remote = _remotes.get(ip);
        if (remote==null)
        {
            Remote r = new Remote(ip,limit);
            remote = _remotes.putIfAbsent(ip,r);
            if (remote==null)
                remote = r;
        }
        
        baseRequest.setAttribute(REMOTE,remote);

        return remote;
    }
    
    
    protected String getRemoteIP(Request baseRequest)
    {
        // Do we have a forwarded header set?
        if (_forwardedHeader!=null && !_forwardedHeader.isEmpty())
        {
            // Yes, then try to get the remote IP from the header
            String remote = _rfc7239?getForwarded(baseRequest):getXForwardedFor(baseRequest);
            if (remote!=null && !remote.isEmpty())
                return remote;
        }
        
        // If no remote IP from a header, determine it directly from the channel
        // Do not use the request methods, as they may have been lied to by the 
        // RequestCustomizer!
        InetSocketAddress inet_addr = baseRequest.getHttpChannel().getRemoteAddress();
        if (inet_addr!=null && inet_addr.getAddress()!=null)
            return inet_addr.getAddress().getHostAddress();
        return null;
    }

    private String getForwarded(Request request)
    {
        // Get the right most Forwarded for value.
        // This is the value from the closest proxy and the only one that
        // can be trusted.
        RFC7239 rfc7239 = new RFC7239();
        HttpFields httpFields = request.getHttpFields();
        for (HttpField field : httpFields)
            if (_forwardedHeader.equalsIgnoreCase(field.getName()))
                rfc7239.addValue(field.getValue());
        
        if (rfc7239.getFor()!=null)
            return new HostPortHttpField(rfc7239.getFor()).getHost();
        
        return null;
    }
    
    private String getXForwardedFor(Request request)
    {
        // Get the right most XForwarded-For for value.
        // This is the value from the closest proxy and the only one that
        // can be trusted.
        String forwarded_for = null;
        HttpFields httpFields = request.getHttpFields();
        for (HttpField field : httpFields)
            if (_forwardedHeader.equalsIgnoreCase(field.getName()))
                forwarded_for = field.getValue();
        
        if (forwarded_for==null || forwarded_for.isEmpty())
            return null;
        
        int comma = forwarded_for.lastIndexOf(',');
        return (comma>=0)?forwarded_for.substring(comma+1).trim():forwarded_for;
    }


    private final class Remote implements Closeable
    {
        private final String _ip;
        private final int _limit;
        private final Locker _locker = new Locker();
        private int _permits;
        private Deque<CompletableFuture<Closeable>> _queue = new ArrayDeque<>();
        private final CompletableFuture<Closeable> _permitted = CompletableFuture.completedFuture(this);
        
        public Remote(String ip, int limit)
        {
            _ip=ip;
            _limit=limit;
        }
        
        public CompletableFuture<Closeable> acquire()
        {
            try(Locker.Lock lock = _locker.lock())
            {
                // Do we have available passes?
                if (_permits<_limit)
                {
                    // Yes - increment the allocated passes
                    _permits++;
                    // return the already completed future
                    return _permitted; // TODO is it OK to share/reuse this?
                }
                
                // No pass available, so queue a new future 
                CompletableFuture<Closeable> pass = new CompletableFuture<Closeable>();
                _queue.addLast(pass);
                return pass;
            }
        }
        
        @Override
        public void close() throws IOException
        {
            try(Locker.Lock lock = _locker.lock())
            {
                // reduce the allocated passes
                _permits--;
                while(true)
                {
                    // Are there any future passes waiting?
                    CompletableFuture<Closeable> permit = _queue.pollFirst();
                    
                    // No - we are done
                    if (permit==null)
                        break;
                    
                    // Yes - if we can complete them, we are done
                    if (permit.complete(this))
                    {
                        _permits++;
                        break;
                    }
                    
                    // Somebody else must have completed/failed that future pass,
                    // so let's try for another.
                }
            }
        }
        
        @Override
        public String toString()
        {
            try(Locker.Lock lock = _locker.lock())
            {
                return String.format("R[ip=%s,p=%d,l=%d,q=%d]",_ip,_permits,_limit,_queue.size());
            }
        }
    }

    private final class RFC7239 extends QuotedCSV
    {
        String _for;
        
        private RFC7239()
        {
            super(false);
        }
        
        String getFor()
        {
            return _for;
        }

        @Override
        protected void parsedParam(StringBuffer buffer, int valueLength, int paramName, int paramValue)
        {
            if (valueLength==0 && paramValue>paramName)
            {
                String name=StringUtil.asciiToLowerCase(buffer.substring(paramName,paramValue-1));
                if ("for".equalsIgnoreCase(name))
                {
                    String value=buffer.substring(paramValue);
                    
                    // if unknown, clear any leftward values
                    if ("unknown".equalsIgnoreCase(value))
                        _for = null;
                    // Otherwise accept IP or token(starting with '_') as remote keys
                    else
                        _for=value;
                }
            }
        }
    }
}