org.subethamail.smtp.internal.server.ConcurrentSessionsBySourceLimiter Maven / Gradle / Ivy

Go to download

Show more of this group Show more artifacts with this name
Show all versions of subethasmtp Show documentation

A fork of a fork (!) of SubEtha, an easy-to-use server-side SMTP library for Java.

There is a newer version: 7.1.3

package org.subethamail.smtp.internal.server;

import java.net.InetAddress;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;

import org.subethamail.smtp.server.Session;
import org.subethamail.smtp.server.SessionHandler;

/**
 * A {@link SessionHandler} to track and limit connection counts by remote addresses.
 *
 * @author Diego Salvi
 */
public final class ConcurrentSessionsBySourceLimiter implements SessionHandler {

    /** Session drop response */
    private final SessionAcceptance drop;

    private final int maxConnectionsPerSource;
    private final ConcurrentMap counts;

    /**
     * Create a new {@link ConcurrentSessionsBySourceLimiter} with default reject message:
     * {@code "421 Too many connections, try again later"}.
     *
     * @param maxConnectionsPerSource maximum number of concurrent connection per remote source ip
     */
    public ConcurrentSessionsBySourceLimiter(int maxConnectionsPerSource) {
        this(maxConnectionsPerSource, 421, "Too many connections, try again later");
    }

    /**
     * Create a new {@link ConcurrentSessionsBySourceLimiter} with custom reject message
     * @param maxConnectionsPerSource maximum number of concurrent connection per remote source ip
     * @param code SMTP code
     * @param message SMTP message
     */
    public ConcurrentSessionsBySourceLimiter(int maxConnectionsPerSource, int code, String message) {
        this.maxConnectionsPerSource = maxConnectionsPerSource;
        this.drop = SessionAcceptance.failure(code, message);
        this.counts = new ConcurrentHashMap<>();
    }

    @Override
    public SessionAcceptance accept(Session session) {
        try {
            counts.compute(toInetAddress(session), (k, v) -> {
                if (v == null) {
                    return 1;
                } else {

                    if (v == maxConnectionsPerSource) {
                        throw LimitReachedException.INSTANCE;
                    } else {
                        return ++v;
                    }
                }
            });
        } catch (LimitReachedException limit) {
            return drop;
        }
        return SessionAcceptance.success();
    }

    @Override
    public void onSessionEnd(Session session) {
        counts.compute(toInetAddress(session), (k, v) -> {
            if (--v == 0) {
                return null;
            } else {
                return v;
            }
        });
    }

    private static InetAddress toInetAddress(Session session) {
        return session.getSocket().getInetAddress();
    }

    /**
     * A lightweight exception to avoid to create non useful stacktraces, this exception is use only internally
     * and never leaked out this class.
     */
    @SuppressWarnings("serial")
    private static final class LimitReachedException extends RuntimeException {
        public static final LimitReachedException INSTANCE = new LimitReachedException();

        public LimitReachedException() {
            /* Disables stacktraces and suppressions */
            super("Limit reached", null, false, false);
        }
    }

}