<commit_before>package com.relayrides.pushy.apns;

import com.google.gson.Gson;
import com.google.gson.GsonBuilder;
import com.relayrides.pushy.apns.auth.ApnsVerificationKey;
import io.netty.buffer.ByteBuf;
import io.netty.buffer.Unpooled;
import io.netty.channel.ChannelFuture;
import io.netty.channel.ChannelHandlerContext;
import io.netty.channel.ChannelPromise;
import io.netty.handler.codec.http.HttpHeaderNames;
import io.netty.handler.codec.http.HttpMethod;
import io.netty.handler.codec.http.HttpResponseStatus;
import io.netty.handler.codec.http2.*;
import io.netty.util.AsciiString;
import io.netty.util.concurrent.PromiseCombiner;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.security.InvalidKeyException;
import java.security.NoSuchAlgorithmException;
import java.security.SignatureException;
import java.util.*;
import java.util.concurrent.TimeUnit;
import java.util.regex.Matcher;
import java.util.regex.Pattern;

class MockApnsServerHandler extends Http2ConnectionHandler implements Http2FrameListener {

    private final Map<Integer, UUID> requestsWaitingForDataFrame = new HashMap<>();

    private final boolean emulateInternalErrors;

    private final Map<String, Map<String, Date>> deviceTokenExpirationsByTopic;

    private final Map<String, ApnsVerificationKey> verificationKeysByKeyId;
    private final Map<ApnsVerificationKey, Set<String>> topicsByVerificationKey;

    private String expectedTeamId;

    private static final Http2Headers SUCCESS_HEADERS = new DefaultHttp2Headers().status(HttpResponseStatus.OK.codeAsText());

    private static final String APNS_PATH_PREFIX = "/3/device/";
    private static final AsciiString APNS_TOPIC_HEADER = new AsciiString("apns-topic");
    private static final AsciiString APNS_PRIORITY_HEADER = new AsciiString("apns-priority");
    private static final AsciiString APNS_ID_HEADER = new AsciiString("apns-id");
    private static final AsciiString APNS_AUTHORIZATION_HEADER = new AsciiString("authorization");

    private static final int MAX_CONTENT_LENGTH = 4096;
    private static final Pattern TOKEN_PATTERN = Pattern.compile("[0-9a-fA-F]{64}");

    private static final Gson GSON = new GsonBuilder()
            .registerTypeAdapter(Date.class, new DateAsTimeSinceEpochTypeAdapter(TimeUnit.MILLISECONDS))
            .create();

    private static final Logger log = LoggerFactory.getLogger(MockApnsServerHandler.class);

    private enum ErrorReason {
        BAD_COLLAPSE_ID("BadCollapseId", HttpResponseStatus.BAD_REQUEST),
        BAD_DEVICE_TOKEN("BadDeviceToken", HttpResponseStatus.BAD_REQUEST),
        BAD_EXPIRATION_DATE("BadExpirationDate", HttpResponseStatus.BAD_REQUEST),
        BAD_MESSAGE_ID("BadMessageId", HttpResponseStatus.BAD_REQUEST),
        BAD_PRIORITY("BadPriority", HttpResponseStatus.BAD_REQUEST),
        BAD_TOPIC("BadTopic", HttpResponseStatus.BAD_REQUEST),
        DEVICE_TOKEN_NOT_FOR_TOPIC("DeviceTokenNotForTopic", HttpResponseStatus.BAD_REQUEST),
        DUPLICATE_HEADERS("DuplicateHeaders", HttpResponseStatus.BAD_REQUEST),
        IDLE_TIMEOUT("IdleTimeout", HttpResponseStatus.BAD_REQUEST),
        MISSING_DEVICE_TOKEN("MissingDeviceToken", HttpResponseStatus.BAD_REQUEST),
        MISSING_TOPIC("MissingTopic", HttpResponseStatus.BAD_REQUEST),
        PAYLOAD_EMPTY("PayloadEmpty", HttpResponseStatus.BAD_REQUEST),
        TOPIC_DISALLOWED("TopicDisallowed", HttpResponseStatus.BAD_REQUEST),

        BAD_CERTIFICATE("BadCertificate", HttpResponseStatus.FORBIDDEN),
        BAD_CERTIFICATE_ENVIRONMENT("BadCertificateEnvironment", HttpResponseStatus.FORBIDDEN),
        EXPIRED_PROVIDER_TOKEN("ExpiredProviderToken", HttpResponseStatus.FORBIDDEN),
        FORBIDDEN("Forbidden", HttpResponseStatus.FORBIDDEN),
        INVALID_PROVIDER_TOKEN("InvalidProviderToken", HttpResponseStatus.FORBIDDEN),
        MISSING_PROVIDER_TOKEN("MissingProviderToken", HttpResponseStatus.FORBIDDEN),

        BAD_PATH("BadPath", HttpResponseStatus.NOT_FOUND),

        METHOD_NOT_ALLOWED("MethodNotAllowed", HttpResponseStatus.METHOD_NOT_ALLOWED),

        UNREGISTERED("Unregistered", HttpResponseStatus.GONE),

        PAYLOAD_TOO_LARGE("PayloadTooLarge", HttpResponseStatus.REQUEST_ENTITY_TOO_LARGE),

        TOO_MANY_PROVIDER_TOKEN_UPDATES("TooManyProviderTokenUpdates", HttpResponseStatus.TOO_MANY_REQUESTS),
        TOO_MANY_REQUESTS("TooManyRequests", HttpResponseStatus.TOO_MANY_REQUESTS),

        INTERNAL_SERVER_ERROR("InternalServerError", HttpResponseStatus.INTERNAL_SERVER_ERROR),

        SERVICE_UNAVAILABLE("ServiceUnavailable", HttpResponseStatus.SERVICE_UNAVAILABLE),
        SHUTDOWN("Shutdown", HttpResponseStatus.SERVICE_UNAVAILABLE);

        private final String reasonText;
        private final HttpResponseStatus httpResponseStatus;

        private ErrorReason(final String reasonText, final HttpResponseStatus httpResponseStatus) {
            this.reasonText = reasonText;
            this.httpResponseStatus = httpResponseStatus;
        }

        public String getReasonText() {
            return this.reasonText;
        }

        public HttpResponseStatus getHttpResponseStatus() {
            return this.httpResponseStatus;
        }
    }

    private static class ExpiredAuthenticationTokenException extends Exception {
        private static final long serialVersionUID = 1L;
    }

    private static class InvalidAuthenticationTokenException extends Exception {
        private static final long serialVersionUID = 1L;

        public InvalidAuthenticationTokenException() {
            super();
        }

        public InvalidAuthenticationTokenException(final Throwable cause) {
            super(cause);
        }
    }

    public static final class MockApnsServerHandlerBuilder extends AbstractHttp2ConnectionHandlerBuilder<MockApnsServerHandler, MockApnsServerHandlerBuilder> {

        private boolean emulateInternalErrors;

        private Map<String, Map<String, Date>> deviceTokenExpirationsByTopic;
        private Map<String, ApnsVerificationKey> verificationKeysByKeyId;
        private Map<ApnsVerificationKey, Set<String>> topicsByVerificationKey;

        @Override
        public MockApnsServerHandlerBuilder initialSettings(final Http2Settings initialSettings) {
            return super.initialSettings(initialSettings);
        }

        public MockApnsServerHandlerBuilder emulateInternalErrors(final boolean emulateInternalErrors) {
            this.emulateInternalErrors = emulateInternalErrors;
            return this;
        }

        public MockApnsServerHandlerBuilder deviceTokenExpirationsByTopic(final Map<String, Map<String, Date>> deviceTokenExpirationsByTopic) {
            this.deviceTokenExpirationsByTopic = deviceTokenExpirationsByTopic;
            return this;
        }

        public Map<String, Map<String, Date>> deviceTokenExpirationsByTopic() {
            return this.deviceTokenExpirationsByTopic;
        }

        public MockApnsServerHandlerBuilder verificationKeysByKeyId(final Map<String, ApnsVerificationKey> verificationKeysByKeyId) {
            this.verificationKeysByKeyId = verificationKeysByKeyId;
            return this;
        }

        public Map<String, ApnsVerificationKey> verificationKeysByKeyId() {
            return this.verificationKeysByKeyId;
        }

        public MockApnsServerHandlerBuilder topicsByVerificationKey(final Map<ApnsVerificationKey, Set<String>> topicsByVerificationKey) {
            this.topicsByVerificationKey = topicsByVerificationKey;
            return this;
        }

        public Map<ApnsVerificationKey, Set<String>> topicsByVerificationKey() {
            return this.topicsByVerificationKey;
        }

        @Override
        public MockApnsServerHandler build(final Http2ConnectionDecoder decoder, final Http2ConnectionEncoder encoder, final Http2Settings initialSettings) {
            final MockApnsServerHandler handler = new MockApnsServerHandler(decoder, encoder, initialSettings, emulateInternalErrors, deviceTokenExpirationsByTopic, verificationKeysByKeyId, topicsByVerificationKey);
            this.frameListener(handler);
            return handler;
        }

        @Override
        public MockApnsServerHandler build() {
            return super.build();
        }
    }

    private static class AcceptNotificationResponse {
        private final int streamId;

        public AcceptNotificationResponse(final int streamId) {
            this.streamId = streamId;
        }

        public int getStreamId() {
            return this.streamId;
        }
    }

    private static class RejectNotificationResponse {
        private final int streamId;
        private final UUID apnsId;
        private final ErrorReason errorReason;
        private final Date timestamp;

        public RejectNotificationResponse(final int streamId, final UUID apnsId, final ErrorReason errorReason) {
            this(streamId, apnsId, errorReason, null);
        }

        public RejectNotificationResponse(final int streamId, final UUID apnsId, final ErrorReason errorReason, final Date timestamp) {
            this.streamId = streamId;
            this.apnsId = apnsId;
            this.errorReason = errorReason;
            this.timestamp = timestamp;
        }

        public int getStreamId() {
            return this.streamId;
        }

        public UUID getApnsId() {
            return this.apnsId;
        }

        public ErrorReason getErrorReason() {
            return this.errorReason;
        }

        public Date getTimestamp() {
            return this.timestamp;
        }
    }

    private static class InternalServerErrorResponse {
        private final int streamId;

        public InternalServerErrorResponse(final int streamId) {
            this.streamId = streamId;
        }

        public int getStreamId() {
            return this.streamId;
        }
    }

    protected MockApnsServerHandler(final Http2ConnectionDecoder decoder,
                                    final Http2ConnectionEncoder encoder,
                                    final Http2Settings initialSettings,
                                    final boolean emulateInternalErrors,
                                    final Map<String, Map<String, Date>> deviceTokenExpirationsByTopic,
                                    final Map<String, ApnsVerificationKey> verificationKeysByKeyId,
                                    final Map<ApnsVerificationKey, Set<String>> topicsByVerificationKey) {

        super(decoder, encoder, initialSettings);

        this.emulateInternalErrors = emulateInternalErrors;

        this.deviceTokenExpirationsByTopic = deviceTokenExpirationsByTopic;

        this.verificationKeysByKeyId = verificationKeysByKeyId;
        this.topicsByVerificationKey = topicsByVerificationKey;
    }

    @Override
    public int onDataRead(final ChannelHandlerContext context, final int streamId, final ByteBuf data, final int padding, final boolean endOfStream) throws Http2Exception {
        final int bytesProcessed = data.readableBytes() + padding;

        if (endOfStream) {
            // Presumably, we spotted an error earlier and sent a response immediately if we don't have an entry in the
            // "waiting for data frame" map.
            if (this.requestsWaitingForDataFrame.containsKey(streamId)) {
                final UUID apnsId = this.requestsWaitingForDataFrame.remove(streamId);

                if (data.readableBytes() <= MAX_CONTENT_LENGTH) {
                    context.channel().writeAndFlush(new AcceptNotificationResponse(streamId));
                } else {
                    context.channel().writeAndFlush(new RejectNotificationResponse(streamId, apnsId, ErrorReason.PAYLOAD_TOO_LARGE));
                }
            }
        }

        return bytesProcessed;
    }

    @Override
    public void onHeadersRead(final ChannelHandlerContext context, final int streamId, final Http2Headers headers, final int padding, final boolean endOfStream) throws Http2Exception {
        if (this.emulateInternalErrors) {
            context.channel().writeAndFlush(new InternalServerErrorResponse(streamId));
        }

        if (!HttpMethod.POST.asciiName().contentEquals(headers.get(Http2Headers.PseudoHeaderName.METHOD.value()))) {
            context.channel().writeAndFlush(new RejectNotificationResponse(streamId, null, ErrorReason.METHOD_NOT_ALLOWED));
            return;
        }

        final UUID apnsId;
        {
            final CharSequence apnsIdSequence = headers.get(APNS_ID_HEADER);

            if (apnsIdSequence != null) {
                try {
                    apnsId = UUID.fromString(apnsIdSequence.toString());
                } catch (final IllegalArgumentException e) {
                    context.channel().writeAndFlush(new RejectNotificationResponse(streamId, null, ErrorReason.BAD_MESSAGE_ID));
                    return;
                }
            } else {
                // If the client didn't send us a UUID, make one up (for now)
                apnsId = UUID.randomUUID();
            }
        }

        if (endOfStream) {
            context.channel().writeAndFlush(new RejectNotificationResponse(streamId, apnsId, ErrorReason.PAYLOAD_EMPTY));
            return;
        }

        final String base64EncodedAuthenticationToken;
        {
            final CharSequence authorizationSequence = headers.get(APNS_AUTHORIZATION_HEADER);

            if (authorizationSequence != null) {
                final String authorizationString = authorizationSequence.toString();

                if (authorizationString.startsWith("bearer")) {
                    base64EncodedAuthenticationToken = authorizationString.substring("bearer".length()).trim();
                } else {
                    base64EncodedAuthenticationToken = null;
                }
            } else {
                base64EncodedAuthenticationToken = null;
            }
        }

        final String topic;
        {
            final CharSequence topicSequence = headers.get(APNS_TOPIC_HEADER);
            topic = (topicSequence != null) ? topicSequence.toString() : null;
        }

        if (topic == null) {
            context.channel().writeAndFlush(new RejectNotificationResponse(streamId, apnsId, ErrorReason.MISSING_TOPIC));
            return;
        }

        // TODO Restore caching
        final AuthenticationToken authenticationToken = new AuthenticationToken(base64EncodedAuthenticationToken);
        final ApnsVerificationKey verificationKey = this.verificationKeysByKeyId.get(authenticationToken.getKeyId());

        // Have we ever heard of the key in question?
        if (verificationKey == null) {
            context.channel().writeAndFlush(new RejectNotificationResponse(streamId, apnsId, ErrorReason.INVALID_PROVIDER_TOKEN));
            return;
        }

        try {
            if (!authenticationToken.verifySignature(verificationKey)) {
                context.channel().writeAndFlush(new RejectNotificationResponse(streamId, apnsId, ErrorReason.INVALID_PROVIDER_TOKEN));
                return;
            }
        } catch (NoSuchAlgorithmException | InvalidKeyException | SignatureException e) {
            // This should never happen (here, at least) because we check keys at construction time. If something's
            // going to go wrong, it will go wrong before we ever get here.
            log.error("Failed to verify signature.", e);
            throw new RuntimeException(e);
        }

        // At this point, we've verified that the token is signed by somebody with the named team's private key. The
        // real APNs server only allows one team per connection, so if this is our first notification, we want to keep
        // track of the team that sent it so we can reject notifications from other teams, even if they're signed
        // correctly.
        if (this.expectedTeamId == null) {
            this.expectedTeamId = authenticationToken.getTeamId();
        }

        if (!this.expectedTeamId.equals(authenticationToken.getTeamId())) {
            context.channel().writeAndFlush(new RejectNotificationResponse(streamId, apnsId, ErrorReason.INVALID_PROVIDER_TOKEN));
            return;
        }

        if (authenticationToken.getIssuedAt().getTime() + MockApnsServer.AUTHENTICATION_TOKEN_EXPIRATION_MILLIS < System.currentTimeMillis()) {
            context.channel().writeAndFlush(new RejectNotificationResponse(streamId, apnsId, ErrorReason.EXPIRED_PROVIDER_TOKEN));
            return;
        }

        final Set<String> topics = this.topicsByVerificationKey.get(verificationKey);

        if (topics == null || !topics.contains(topic)) {
            context.channel().writeAndFlush(new RejectNotificationResponse(streamId, apnsId, ErrorReason.DEVICE_TOKEN_NOT_FOR_TOPIC));
            return;
        }

        {
            final Integer priorityCode = headers.getInt(APNS_PRIORITY_HEADER);

            if (priorityCode != null) {
                try {
                    DeliveryPriority.getFromCode(priorityCode);
                } catch (final IllegalArgumentException e) {
                    context.channel().writeAndFlush(new RejectNotificationResponse(streamId, apnsId, ErrorReason.BAD_PRIORITY));
                    return;
                }
            }
        }

        {
            final CharSequence pathSequence = headers.get(Http2Headers.PseudoHeaderName.PATH.value());

            if (pathSequence != null) {
                final String pathString = pathSequence.toString();

                if (pathString.startsWith(APNS_PATH_PREFIX)) {
                    final String tokenString = pathString.substring(APNS_PATH_PREFIX.length());

                    final Matcher tokenMatcher = TOKEN_PATTERN.matcher(tokenString);

                    if (!tokenMatcher.matches()) {
                        context.channel().writeAndFlush(new RejectNotificationResponse(streamId, apnsId, ErrorReason.BAD_DEVICE_TOKEN));
                        return;
                    }

                    final boolean deviceTokenRegisteredForTopic;
                    final Date expirationTimestamp;
                    {
                        final Map<String, Date> expirationTimestampsByDeviceToken = this.deviceTokenExpirationsByTopic.get(topic);

                        expirationTimestamp = expirationTimestampsByDeviceToken != null ? expirationTimestampsByDeviceToken.get(tokenString) : null;
                        deviceTokenRegisteredForTopic = expirationTimestampsByDeviceToken != null && expirationTimestampsByDeviceToken.containsKey(tokenString);

                    }

                    if (expirationTimestamp != null) {
                        context.channel().writeAndFlush(new RejectNotificationResponse(streamId, apnsId, ErrorReason.UNREGISTERED, expirationTimestamp));
                        return;
                    }

                    if (!deviceTokenRegisteredForTopic) {
                        context.channel().writeAndFlush(new RejectNotificationResponse(streamId, apnsId, ErrorReason.DEVICE_TOKEN_NOT_FOR_TOPIC));
                        return;
                    }
                }
            } else {
                context.channel().writeAndFlush(new RejectNotificationResponse(streamId, apnsId, ErrorReason.BAD_PATH));
                return;
            }
        }

        this.requestsWaitingForDataFrame.put(streamId, apnsId);
    }

    @Override
    public void onHeadersRead(final ChannelHandlerContext ctx, final int streamId, final Http2Headers headers, final int streamDependency,
            final short weight, final boolean exclusive, final int padding, final boolean endOfStream) throws Http2Exception {

        this.onHeadersRead(ctx, streamId, headers, padding, endOfStream);
    }

    @Override
    public void onPriorityRead(final ChannelHandlerContext ctx, final int streamId, final int streamDependency, final short weight, final boolean exclusive) throws Http2Exception {
    }

    @Override
    public void onRstStreamRead(final ChannelHandlerContext ctx, final int streamId, final long errorCode) throws Http2Exception {
    }

    @Override
    public void onSettingsAckRead(final ChannelHandlerContext ctx) throws Http2Exception {
    }

    @Override
    public void onSettingsRead(final ChannelHandlerContext ctx, final Http2Settings settings) throws Http2Exception {
    }<commit_msg>Moved to an exception-based header-checking system.<commit_after>package com.relayrides.pushy.apns;

import com.google.gson.Gson;
import com.google.gson.GsonBuilder;
import com.relayrides.pushy.apns.auth.ApnsVerificationKey;
import io.netty.buffer.ByteBuf;
import io.netty.buffer.Unpooled;
import io.netty.channel.ChannelFuture;
import io.netty.channel.ChannelHandlerContext;
import io.netty.channel.ChannelPromise;
import io.netty.handler.codec.http.HttpHeaderNames;
import io.netty.handler.codec.http.HttpMethod;
import io.netty.handler.codec.http.HttpResponseStatus;
import io.netty.handler.codec.http2.*;
import io.netty.util.AsciiString;
import io.netty.util.concurrent.PromiseCombiner;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.security.InvalidKeyException;
import java.security.NoSuchAlgorithmException;
import java.security.SignatureException;
import java.util.*;
import java.util.concurrent.TimeUnit;
import java.util.regex.Matcher;
import java.util.regex.Pattern;

class MockApnsServerHandler extends Http2ConnectionHandler implements Http2FrameListener {

    private final Map<Integer, UUID> requestsWaitingForDataFrame = new HashMap<>();

    private final boolean emulateInternalErrors;

    private final Map<String, Map<String, Date>> deviceTokenExpirationsByTopic;

    private final Map<String, ApnsVerificationKey> verificationKeysByKeyId;
    private final Map<ApnsVerificationKey, Set<String>> topicsByVerificationKey;

    private String expectedTeamId;

    private static final Http2Headers SUCCESS_HEADERS = new DefaultHttp2Headers().status(HttpResponseStatus.OK.codeAsText());

    private static final String APNS_PATH_PREFIX = "/3/device/";
    private static final AsciiString APNS_TOPIC_HEADER = new AsciiString("apns-topic");
    private static final AsciiString APNS_PRIORITY_HEADER = new AsciiString("apns-priority");
    private static final AsciiString APNS_ID_HEADER = new AsciiString("apns-id");
    private static final AsciiString APNS_AUTHORIZATION_HEADER = new AsciiString("authorization");

    private static final int MAX_CONTENT_LENGTH = 4096;
    private static final Pattern TOKEN_PATTERN = Pattern.compile("[0-9a-fA-F]{64}");

    private static final Gson GSON = new GsonBuilder()
            .registerTypeAdapter(Date.class, new DateAsTimeSinceEpochTypeAdapter(TimeUnit.MILLISECONDS))
            .create();

    private static final Logger log = LoggerFactory.getLogger(MockApnsServerHandler.class);

    private enum ErrorReason {
        BAD_COLLAPSE_ID("BadCollapseId", HttpResponseStatus.BAD_REQUEST),
        BAD_DEVICE_TOKEN("BadDeviceToken", HttpResponseStatus.BAD_REQUEST),
        BAD_EXPIRATION_DATE("BadExpirationDate", HttpResponseStatus.BAD_REQUEST),
        BAD_MESSAGE_ID("BadMessageId", HttpResponseStatus.BAD_REQUEST),
        BAD_PRIORITY("BadPriority", HttpResponseStatus.BAD_REQUEST),
        BAD_TOPIC("BadTopic", HttpResponseStatus.BAD_REQUEST),
        DEVICE_TOKEN_NOT_FOR_TOPIC("DeviceTokenNotForTopic", HttpResponseStatus.BAD_REQUEST),
        DUPLICATE_HEADERS("DuplicateHeaders", HttpResponseStatus.BAD_REQUEST),
        IDLE_TIMEOUT("IdleTimeout", HttpResponseStatus.BAD_REQUEST),
        MISSING_DEVICE_TOKEN("MissingDeviceToken", HttpResponseStatus.BAD_REQUEST),
        MISSING_TOPIC("MissingTopic", HttpResponseStatus.BAD_REQUEST),
        PAYLOAD_EMPTY("PayloadEmpty", HttpResponseStatus.BAD_REQUEST),
        TOPIC_DISALLOWED("TopicDisallowed", HttpResponseStatus.BAD_REQUEST),

        BAD_CERTIFICATE("BadCertificate", HttpResponseStatus.FORBIDDEN),
        BAD_CERTIFICATE_ENVIRONMENT("BadCertificateEnvironment", HttpResponseStatus.FORBIDDEN),
        EXPIRED_PROVIDER_TOKEN("ExpiredProviderToken", HttpResponseStatus.FORBIDDEN),
        FORBIDDEN("Forbidden", HttpResponseStatus.FORBIDDEN),
        INVALID_PROVIDER_TOKEN("InvalidProviderToken", HttpResponseStatus.FORBIDDEN),
        MISSING_PROVIDER_TOKEN("MissingProviderToken", HttpResponseStatus.FORBIDDEN),

        BAD_PATH("BadPath", HttpResponseStatus.NOT_FOUND),

        METHOD_NOT_ALLOWED("MethodNotAllowed", HttpResponseStatus.METHOD_NOT_ALLOWED),

        UNREGISTERED("Unregistered", HttpResponseStatus.GONE),

        PAYLOAD_TOO_LARGE("PayloadTooLarge", HttpResponseStatus.REQUEST_ENTITY_TOO_LARGE),

        TOO_MANY_PROVIDER_TOKEN_UPDATES("TooManyProviderTokenUpdates", HttpResponseStatus.TOO_MANY_REQUESTS),
        TOO_MANY_REQUESTS("TooManyRequests", HttpResponseStatus.TOO_MANY_REQUESTS),

        INTERNAL_SERVER_ERROR("InternalServerError", HttpResponseStatus.INTERNAL_SERVER_ERROR),

        SERVICE_UNAVAILABLE("ServiceUnavailable", HttpResponseStatus.SERVICE_UNAVAILABLE),
        SHUTDOWN("Shutdown", HttpResponseStatus.SERVICE_UNAVAILABLE);

        private final String reasonText;
        private final HttpResponseStatus httpResponseStatus;

        private ErrorReason(final String reasonText, final HttpResponseStatus httpResponseStatus) {
            this.reasonText = reasonText;
            this.httpResponseStatus = httpResponseStatus;
        }

        public String getReasonText() {
            return this.reasonText;
        }

        public HttpResponseStatus getHttpResponseStatus() {
            return this.httpResponseStatus;
        }
    }

    protected static class RejectedNotificationException extends Exception {
        private final ErrorReason errorReason;
        private final Date deviceTokenExpirationTimestamp;

        public RejectedNotificationException(final ErrorReason errorReason) {
            this(errorReason, null);
        }

        public RejectedNotificationException(final ErrorReason errorReason, final Date deviceTokenExpirationTimestamp) {
            Objects.requireNonNull(errorReason, "Error reason must not be null.");

            this.errorReason = errorReason;
            this.deviceTokenExpirationTimestamp = deviceTokenExpirationTimestamp;
        }

        public ErrorReason getErrorReason() {
            return errorReason;
        }

        public Date getDeviceTokenExpirationTimestamp() {
            return deviceTokenExpirationTimestamp;
        }
    }

    private static class ExpiredAuthenticationTokenException extends Exception {
        private static final long serialVersionUID = 1L;
    }

    private static class InvalidAuthenticationTokenException extends Exception {
        private static final long serialVersionUID = 1L;

        public InvalidAuthenticationTokenException() {
            super();
        }

        public InvalidAuthenticationTokenException(final Throwable cause) {
            super(cause);
        }
    }

    public static final class MockApnsServerHandlerBuilder extends AbstractHttp2ConnectionHandlerBuilder<MockApnsServerHandler, MockApnsServerHandlerBuilder> {

        private boolean emulateInternalErrors;

        private Map<String, Map<String, Date>> deviceTokenExpirationsByTopic;
        private Map<String, ApnsVerificationKey> verificationKeysByKeyId;
        private Map<ApnsVerificationKey, Set<String>> topicsByVerificationKey;

        @Override
        public MockApnsServerHandlerBuilder initialSettings(final Http2Settings initialSettings) {
            return super.initialSettings(initialSettings);
        }

        public MockApnsServerHandlerBuilder emulateInternalErrors(final boolean emulateInternalErrors) {
            this.emulateInternalErrors = emulateInternalErrors;
            return this;
        }

        public MockApnsServerHandlerBuilder deviceTokenExpirationsByTopic(final Map<String, Map<String, Date>> deviceTokenExpirationsByTopic) {
            this.deviceTokenExpirationsByTopic = deviceTokenExpirationsByTopic;
            return this;
        }

        public Map<String, Map<String, Date>> deviceTokenExpirationsByTopic() {
            return this.deviceTokenExpirationsByTopic;
        }

        public MockApnsServerHandlerBuilder verificationKeysByKeyId(final Map<String, ApnsVerificationKey> verificationKeysByKeyId) {
            this.verificationKeysByKeyId = verificationKeysByKeyId;
            return this;
        }

        public Map<String, ApnsVerificationKey> verificationKeysByKeyId() {
            return this.verificationKeysByKeyId;
        }

        public MockApnsServerHandlerBuilder topicsByVerificationKey(final Map<ApnsVerificationKey, Set<String>> topicsByVerificationKey) {
            this.topicsByVerificationKey = topicsByVerificationKey;
            return this;
        }

        public Map<ApnsVerificationKey, Set<String>> topicsByVerificationKey() {
            return this.topicsByVerificationKey;
        }

        @Override
        public MockApnsServerHandler build(final Http2ConnectionDecoder decoder, final Http2ConnectionEncoder encoder, final Http2Settings initialSettings) {
            final MockApnsServerHandler handler = new MockApnsServerHandler(decoder, encoder, initialSettings, emulateInternalErrors, deviceTokenExpirationsByTopic, verificationKeysByKeyId, topicsByVerificationKey);
            this.frameListener(handler);
            return handler;
        }

        @Override
        public MockApnsServerHandler build() {
            return super.build();
        }
    }

    private static class AcceptNotificationResponse {
        private final int streamId;

        public AcceptNotificationResponse(final int streamId) {
            this.streamId = streamId;
        }

        public int getStreamId() {
            return this.streamId;
        }
    }

    private static class RejectNotificationResponse {
        private final int streamId;
        private final UUID apnsId;
        private final ErrorReason errorReason;
        private final Date timestamp;

        public RejectNotificationResponse(final int streamId, final UUID apnsId, final ErrorReason errorReason, final Date timestamp) {
            this.streamId = streamId;
            this.apnsId = apnsId;
            this.errorReason = errorReason;
            this.timestamp = timestamp;
        }

        public int getStreamId() {
            return this.streamId;
        }

        public UUID getApnsId() {
            return this.apnsId;
        }

        public ErrorReason getErrorReason() {
            return this.errorReason;
        }

        public Date getTimestamp() {
            return this.timestamp;
        }
    }

    private static class InternalServerErrorResponse {
        private final int streamId;

        public InternalServerErrorResponse(final int streamId) {
            this.streamId = streamId;
        }

        public int getStreamId() {
            return this.streamId;
        }
    }

    protected MockApnsServerHandler(final Http2ConnectionDecoder decoder,
                                    final Http2ConnectionEncoder encoder,
                                    final Http2Settings initialSettings,
                                    final boolean emulateInternalErrors,
                                    final Map<String, Map<String, Date>> deviceTokenExpirationsByTopic,
                                    final Map<String, ApnsVerificationKey> verificationKeysByKeyId,
                                    final Map<ApnsVerificationKey, Set<String>> topicsByVerificationKey) {

        super(decoder, encoder, initialSettings);

        this.emulateInternalErrors = emulateInternalErrors;

        this.deviceTokenExpirationsByTopic = deviceTokenExpirationsByTopic;

        this.verificationKeysByKeyId = verificationKeysByKeyId;
        this.topicsByVerificationKey = topicsByVerificationKey;
    }

    @Override
    public int onDataRead(final ChannelHandlerContext context, final int streamId, final ByteBuf data, final int padding, final boolean endOfStream) throws Http2Exception {
        final int bytesProcessed = data.readableBytes() + padding;

        if (endOfStream) {
            // Presumably, we spotted an error earlier and sent a response immediately if we don't have an entry in the
            // "waiting for data frame" map.
            if (this.requestsWaitingForDataFrame.containsKey(streamId)) {
                final UUID apnsId = this.requestsWaitingForDataFrame.remove(streamId);

                if (data.readableBytes() <= MAX_CONTENT_LENGTH) {
                    context.channel().writeAndFlush(new AcceptNotificationResponse(streamId));
                } else {
                    context.channel().writeAndFlush(new RejectNotificationResponse(streamId, apnsId, ErrorReason.PAYLOAD_TOO_LARGE, null));
                }
            }
        }

        return bytesProcessed;
    }

    @Override
    public void onHeadersRead(final ChannelHandlerContext context, final int streamId, final Http2Headers headers, final int padding, final boolean endOfStream) throws Http2Exception {
        if (this.emulateInternalErrors) {
            context.channel().writeAndFlush(new InternalServerErrorResponse(streamId));
        } else {
            UUID apnsId = null;

            try {
                {
                    final CharSequence apnsIdSequence = headers.get(APNS_ID_HEADER);

                    if (apnsIdSequence != null) {
                        try {
                            apnsId = UUID.fromString(apnsIdSequence.toString());
                        } catch (final IllegalArgumentException e) {
                            throw new RejectedNotificationException(ErrorReason.BAD_MESSAGE_ID);
                        }
                    } else {
                        // If the client didn't send us a UUID, make one up (for now)
                        apnsId = UUID.randomUUID();
                    }
                }

                if (!HttpMethod.POST.asciiName().contentEquals(headers.get(Http2Headers.PseudoHeaderName.METHOD.value()))) {
                    throw new RejectedNotificationException(ErrorReason.METHOD_NOT_ALLOWED);
                }

                if (endOfStream) {
                    throw new RejectedNotificationException(ErrorReason.PAYLOAD_EMPTY);
                }

                this.verifyHeaders(headers);
                this.requestsWaitingForDataFrame.put(streamId, apnsId);
            } catch (final RejectedNotificationException e) {
                context.channel().writeAndFlush(new RejectNotificationResponse(streamId, apnsId, e.getErrorReason(), e.getDeviceTokenExpirationTimestamp()));
            }
        }
    }

    protected void verifyHeaders(final Http2Headers headers) throws RejectedNotificationException {
        final String base64EncodedAuthenticationToken;
        {
            final CharSequence authorizationSequence = headers.get(APNS_AUTHORIZATION_HEADER);

            if (authorizationSequence != null) {
                final String authorizationString = authorizationSequence.toString();

                if (authorizationString.startsWith("bearer")) {
                    base64EncodedAuthenticationToken = authorizationString.substring("bearer".length()).trim();
                } else {
                    base64EncodedAuthenticationToken = null;
                }
            } else {
                base64EncodedAuthenticationToken = null;
            }
        }

        final String topic;
        {
            final CharSequence topicSequence = headers.get(APNS_TOPIC_HEADER);
            topic = (topicSequence != null) ? topicSequence.toString() : null;
        }

        if (topic == null) {
            throw new RejectedNotificationException(ErrorReason.MISSING_TOPIC);
        }

        // TODO Restore caching
        final AuthenticationToken authenticationToken = new AuthenticationToken(base64EncodedAuthenticationToken);
        final ApnsVerificationKey verificationKey = this.verificationKeysByKeyId.get(authenticationToken.getKeyId());

        // Have we ever heard of the key in question?
        if (verificationKey == null) {
            throw new RejectedNotificationException(ErrorReason.INVALID_PROVIDER_TOKEN);
        }

        try {
            if (!authenticationToken.verifySignature(verificationKey)) {
                throw new RejectedNotificationException(ErrorReason.INVALID_PROVIDER_TOKEN);
            }
        } catch (NoSuchAlgorithmException | InvalidKeyException | SignatureException e) {
            // This should never happen (here, at least) because we check keys at construction time. If something's
            // going to go wrong, it will go wrong before we ever get here.
            log.error("Failed to verify signature.", e);
            throw new RuntimeException(e);
        }

        // At this point, we've verified that the token is signed by somebody with the named team's private key. The
        // real APNs server only allows one team per connection, so if this is our first notification, we want to keep
        // track of the team that sent it so we can reject notifications from other teams, even if they're signed
        // correctly.
        if (this.expectedTeamId == null) {
            this.expectedTeamId = authenticationToken.getTeamId();
        }

        if (!this.expectedTeamId.equals(authenticationToken.getTeamId())) {
            throw new RejectedNotificationException(ErrorReason.INVALID_PROVIDER_TOKEN);
        }

        if (authenticationToken.getIssuedAt().getTime() + MockApnsServer.AUTHENTICATION_TOKEN_EXPIRATION_MILLIS < System.currentTimeMillis()) {
            throw new RejectedNotificationException(ErrorReason.EXPIRED_PROVIDER_TOKEN);
        }

        final Set<String> topics = this.topicsByVerificationKey.get(verificationKey);

        if (topics == null || !topics.contains(topic)) {
            throw new RejectedNotificationException(ErrorReason.DEVICE_TOKEN_NOT_FOR_TOPIC);
        }

        {
            final Integer priorityCode = headers.getInt(APNS_PRIORITY_HEADER);

            if (priorityCode != null) {
                try {
                    DeliveryPriority.getFromCode(priorityCode);
                } catch (final IllegalArgumentException e) {
                    throw new RejectedNotificationException(ErrorReason.BAD_PRIORITY);
                }
            }
        }

        {
            final CharSequence pathSequence = headers.get(Http2Headers.PseudoHeaderName.PATH.value());

            if (pathSequence != null) {
                final String pathString = pathSequence.toString();

                if (pathString.startsWith(APNS_PATH_PREFIX)) {
                    final String tokenString = pathString.substring(APNS_PATH_PREFIX.length());

                    final Matcher tokenMatcher = TOKEN_PATTERN.matcher(tokenString);

                    if (!tokenMatcher.matches()) {
                        throw new RejectedNotificationException(ErrorReason.BAD_DEVICE_TOKEN);
                    }

                    final boolean deviceTokenRegisteredForTopic;
                    final Date expirationTimestamp;
                    {
                        final Map<String, Date> expirationTimestampsByDeviceToken = this.deviceTokenExpirationsByTopic.get(topic);

                        expirationTimestamp = expirationTimestampsByDeviceToken != null ? expirationTimestampsByDeviceToken.get(tokenString) : null;
                        deviceTokenRegisteredForTopic = expirationTimestampsByDeviceToken != null && expirationTimestampsByDeviceToken.containsKey(tokenString);
                    }

                    if (expirationTimestamp != null) {
                        throw new RejectedNotificationException(ErrorReason.UNREGISTERED, expirationTimestamp);
                    }

                    if (!deviceTokenRegisteredForTopic) {
                        throw new RejectedNotificationException(ErrorReason.DEVICE_TOKEN_NOT_FOR_TOPIC);
                    }
                }
            } else {
                throw new RejectedNotificationException(ErrorReason.BAD_PATH);
            }
        }
    }

    @Override
    public void onHeadersRead(final ChannelHandlerContext ctx, final int streamId, final Http2Headers headers, final int streamDependency,
            final short weight, final boolean exclusive, final int padding, final boolean endOfStream) throws Http2Exception {

        this.onHeadersRead(ctx, streamId, headers, padding, endOfStream);
    }

    @Override
    public void onPriorityRead(final ChannelHandlerContext ctx, final int streamId, final int streamDependency, final short weight, final boolean exclusive) throws Http2Exception {
    }

    @Override
    public void onRstStreamRead(final ChannelHandlerContext ctx, final int streamId, final long errorCode) throws Http2Exception {
    }

    @Override
    public void onSettingsAckRead(final ChannelHandlerContext ctx) throws Http2Exception {
    }

    @Override
    public void onSettingsRead(final ChannelHandlerContext ctx, final Http2Settings settings) throws Http2Exception {
    }
{'fraction_non_alphanumeric': 0.06968354277136146, 'fraction_numerical': 0.002348611413767221, 'mean_word_length': 2.7853542296764733, 'pattern_counts': {'":': 0, '<': 67, '<?xml version=': 0, '>': 63, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 12, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '18351', 'n_tokens_mistral': 10511, 'n_tokens_neox': 10067, 'n_words': 2201}