RetryAfterParser.java

package com.maybeitssquid.retry;

import jakarta.servlet.http.HttpServletResponse;
import java.math.BigDecimal;
import java.math.RoundingMode;
import java.time.*;
import java.time.format.DateTimeFormatter;
import java.util.Optional;
import java.util.function.Function;
import java.util.regex.Pattern;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/**
 * Parses an HTTP Retry-After header to determine how long to wait before retrying. Parsing of the
 * header is based on a superset of RFC 7231. Anything permitted by that RFC should parse correctly.
 * Minor variations may be allowed, such as fractional seconds, optional day name, single-digit
 * hour/minute/second, optional seconds, extra whitespace before the hour.
 */
public class RetryAfterParser implements Function<HttpServletResponse, Optional<Duration>> {

  private final Function<String, Optional<Duration>>[] parsers;

  /** Logger for errors during parsing, particularly to diagnose a misbehaving server. */
  public static final Logger LOGGER = LoggerFactory.getLogger(RetryAfterParser.class);

  /** The {@code Retry-After} header name. */
  public static final String RETRY_AFTER_HEADER = "Retry-After";

  /**
   * The longest delay this parser will report. A server that asks for more than can be expressed in
   * milliseconds is clamped to this rather than having the value wrap or be discarded, so the delay
   * still exceeds any limit a caller has configured and the retry is refused rather than allowed.
   */
  public static final Duration MAX_DELAY = Duration.ofMillis(Long.MAX_VALUE);

  /** {@link #MAX_DELAY} in milliseconds, for comparison before conversion. */
  private static final BigDecimal MAX_MILLIS = BigDecimal.valueOf(Long.MAX_VALUE);

  /** Longest prefix of a rejected header echoed into the log. */
  private static final int MAX_LOGGED = 128;

  /**
   * Renders a header value for logging. The value comes from the remote server, so control
   * characters are escaped to stop a hostile value from forging log records, and the result is
   * truncated to bound how much a misbehaving server can write to the log.
   *
   * <p>Package-private rather than private so the escaping can be tested directly. The package is
   * exported but the method is not accessible outside it.
   *
   * @param header the raw header value, which may be null.
   * @return a single-line, length-bounded rendering safe to pass to the logger.
   */
  static String forLog(final String header) {
    if (header == null) {
      return "null";
    }
    final boolean truncated = header.length() > MAX_LOGGED;
    final String shown = truncated ? header.substring(0, MAX_LOGGED) : header;
    final StringBuilder escaped = new StringBuilder(shown.length());
    for (int i = 0; i < shown.length(); i++) {
      final char c = shown.charAt(i);
      if (c == '\\') {
        escaped.append("\\\\");
      } else if (c < 0x20 || c == 0x7f) {
        escaped.append(String.format("\\u%04x", (int) c));
      } else {
        escaped.append(c);
      }
    }
    return truncated ? escaped.append("...(truncated)").toString() : escaped.toString();
  }

  private static final DateTimeFormatter RFC_850_FORMATTER =
      DateTimeFormatter.ofPattern("[EEEE, ]d-MMM-yy H:m[:s] z");
  private static final DateTimeFormatter ASCTIME_FORMATTER =
      DateTimeFormatter.ofPattern("[E ]MMM [ ]d H:m[:s] yyyy");

  @SafeVarargs
  private RetryAfterParser(final Function<String, Optional<Duration>>... parsers) {
    this.parsers = parsers;
  }

  /**
   * Parser that accepts only {@code Retry-After} headers with an integer number of seconds to wait.
   * If the header contains some other value, such as a date, it is ignored.
   *
   * @return Parser that accepts only {@code delay-seconds} in the {@code Retry-After} header.
   */
  public static RetryAfterParser secondsOnly() {
    return new RetryAfterParser(STRICT_SECONDS);
  }

  /**
   * Parser that accepts {@code Retry-After} headers with a potentially decimal number of seconds to
   * wait. If the header contains some other value, such as a date, it is ignored.
   *
   * @return Parser that accepts only seconds in the {@code Retry-After} header.
   */
  public static RetryAfterParser decimalSeconds() {
    // Attempt STRICT_SECONDS first, because the DECIMAL_SECONDS parse is expensive
    return new RetryAfterParser(STRICT_SECONDS, DECIMAL_SECONDS);
  }

  /**
   * Parser that accepts only {@code Retry-After} headers that meet a reasonably strict
   * interpretation of RFC-9110.
   *
   * @param clock the clock to use to compute offsets when the header is a date. Useful for testing.
   * @return Parser for RFC-9110 {@code Retry-After} headers.
   */
  public static RetryAfterParser strict(final InstantSource clock) {
    return new RetryAfterParser(
        STRICT_SECONDS, wait(clock, IMF_FIXDATE), wait(clock, RFC_850), wait(clock, ASCTIME));
  }

  /**
   * Parser that accepts only {@code Retry-After} headers that meet a reasonably strict
   * interpretation of RFC-9110. Uses system clock to compute offsets when the header is a date.
   *
   * @return Parser for RFC-9110 {@code Retry-After} headers.
   */
  public static RetryAfterParser strict() {
    return strict(InstantSource.system());
  }

  /**
   * Parser that accepts {@code Retry-After} headers that meet a superset of RFC-9110, including
   * headers with a decimal number of seconds delay and ISO-8601 instants.
   *
   * @param clock the clock to use to compute offsets when the header is a date. Useful for testing.
   * @return Parser for extended {@code Retry-After} headers.
   */
  public static RetryAfterParser extended(final InstantSource clock) {
    return new RetryAfterParser(
        STRICT_SECONDS,
        DECIMAL_SECONDS,
        wait(clock, IMF_FIXDATE),
        wait(clock, RFC_850),
        wait(clock, ASCTIME),
        wait(clock, ISO));
  }

  /**
   * Parser that accepts {@code Retry-After} headers that meet a superset of RFC-9110, including
   * headers with a decimal number of seconds delay. Uses system clock to compute offsets when the
   * header is a date.
   *
   * @return Parser for extended {@code Retry-After} headers.
   */
  public static RetryAfterParser extended() {
    return extended(InstantSource.system());
  }

  /**
   * Accept {@code Retry-After} header that matches only strict <a
   * href="https://datatracker.ietf.org/doc/html/rfc7231#section-7.1.3">RFC 7231</a> {@code
   * delay-seconds}.
   *
   * <p>A delay too large to express in milliseconds is clamped to {@link #MAX_DELAY} rather than
   * wrapped or discarded, so every {@link Duration} this parser returns is safe to pass to {@link
   * Duration#toMillis()} and still exceeds any sane limit.
   */
  public static final Function<String, Optional<Duration>> STRICT_SECONDS =
      new PatternGuarded<>("^\\d+$", RetryAfterParser::strictSeconds);

  /**
   * Accept extended {@code Retry-After} header that allows decimal seconds.
   *
   * <p>Precision finer than a millisecond is truncated, and an oversized delay is clamped to {@link
   * #MAX_DELAY}, matching {@link #STRICT_SECONDS}.
   */
  public static final Function<String, Optional<Duration>> DECIMAL_SECONDS =
      new PatternGuarded<>("^\\d+(\\.\\d*)?$", h -> clamped(new BigDecimal(h)));

  /**
   * Converts an integer number of seconds to a duration, clamping a delay too large to express in
   * milliseconds.
   *
   * @param header the delay in seconds, already known to contain only digits.
   * @return the delay, clamped to {@link #MAX_DELAY}.
   */
  private static Duration strictSeconds(final String header) {
    try {
      return Duration.ofMillis(Math.multiplyExact(Long.parseLong(header), 1000L));
    } catch (final ArithmeticException | NumberFormatException tooLarge) {
      // The guard pattern admits only digits, so the only way to get here is a value too large
      return MAX_DELAY;
    }
  }

  /**
   * Converts a decimal number of seconds to a duration, truncating precision finer than a
   * millisecond and clamping a delay too large to express.
   *
   * @param seconds the delay in seconds.
   * @return the delay, clamped to {@link #MAX_DELAY}.
   */
  private static Duration clamped(final BigDecimal seconds) {
    final BigDecimal millis = seconds.movePointRight(3).setScale(0, RoundingMode.DOWN);
    return millis.compareTo(MAX_MILLIS) >= 0
        ? MAX_DELAY
        : Duration.ofMillis(millis.longValueExact());
  }

  /**
   * Forgiving parser for a superset of IMF-fixdate using the builtin {@link
   * DateTimeFormatter#RFC_1123_DATE_TIME}
   *
   * <p>Example: "Thu, 02 Jan 2003 01:23:45 GMT"
   */
  public static final Function<String, Optional<ZonedDateTime>> IMF_FIXDATE =
      new PatternGuarded<>(
          "^(\\w{3},\\s)?\\d{1,2}\\s\\w{3}\\s\\d{4}\\s\\d{1,2}:\\d{2}(:\\d{2})?\\sGMT$",
          h -> ZonedDateTime.parse(h, DateTimeFormatter.RFC_1123_DATE_TIME));

  /**
   * Forgiving parer for a superset of RFC-850 dates.
   *
   * <p>Example: "Thursday, 02-Jan-03 01:23:45 GMT"
   */
  public static final Function<String, Optional<ZonedDateTime>> RFC_850 =
      new PatternGuarded<>(
          "^(\\w+,\\s)?\\d{1,2}-\\w{3}-\\d{2}\\s\\d{1,2}:\\d{2}(:\\d{2})?\\s\\w+$",
          h -> ZonedDateTime.parse(h, RFC_850_FORMATTER));

  /**
   * Forgiving parer for a superset of ASCTIME dates.
   *
   * <p>Example: "Thu Jan 2 01:23:45 2003"
   */
  public static final Function<String, Optional<ZonedDateTime>> ASCTIME =
      new PatternGuarded<>(
          "^(\\w{3}\\s+)?\\w{3}\\s+\\d+\\s\\d{1,2}:\\d{2}(:\\d{2}+)?\\s+\\d{4}$",
          h -> LocalDateTime.parse(h, ASCTIME_FORMATTER).atZone(ZoneOffset.UTC));

  /**
   * Parser for ISO-8601 dates. Any number of fractional second digits from 1 to 9 is accepted, not
   * only whole groups of three.
   *
   * <p>Example: "2011-12-03T10:15:30Z" Example: "2011-12-03T10:15:30.1Z" Example:
   * "2011-12-03T10:15:30.123Z" Example: "2011-12-03T10:15:30.123456789Z"
   */
  public static final Function<String, Optional<ZonedDateTime>> ISO =
      new PatternGuarded<>(
          "^\\d{4}-\\d{2}-\\d{2}T\\d{2}:\\d{2}:\\d{2}(\\.\\d{1,9})?Z$",
          h -> ZonedDateTime.parse(h, DateTimeFormatter.ISO_OFFSET_DATE_TIME));

  /**
   * Parser to recognize the Retry-After formats defined in section 5.6.6 of RFC-9110 and convert
   * the value to a duration. If the header specified a date, the duration is relative to the
   * current time.
   *
   * <p>The source of current time defaults to the system clock, but can be overridden by supplying
   * an {@link InstantSource} to the constructor.
   *
   * @param response the raw HTTP servlet response
   * @return duration until a retry is allowed
   */
  @Override
  public Optional<Duration> apply(final HttpServletResponse response) {
    if (response == null) return Optional.empty();

    final String rawHeader = response.getHeader(RETRY_AFTER_HEADER);
    if (rawHeader == null) return Optional.empty();

    final String header = rawHeader.trim();
    if (header.isEmpty()) {
      LOGGER.warn("Received empty Retry-After header \"{}\"", forLog(rawHeader));
      return Optional.empty();
    }

    // Return the result of the first successful parse
    for (final Function<String, Optional<Duration>> parser : parsers) {
      final Optional<Duration> retryAfter = parser.apply(header);
      if (retryAfter.isPresent()) {
        return retryAfter;
      }
    }

    LOGGER.warn("Received unrecognized Retry-After header \"{}\"", forLog(rawHeader));
    return Optional.empty();
  }

  /**
   * Converts header that arrives as a date into an offset from q clock.
   *
   * @param clock the clock to obtain the current time.
   * @param parser parser for the date in the header.
   * @return function that converts a date header into an offset.
   */
  private static Function<String, Optional<Duration>> wait(
      final InstantSource clock, final Function<String, Optional<ZonedDateTime>> parser) {

    return h ->
        parser
            .apply(h)
            .map(
                t -> {
                  final Duration difference = Duration.between(clock.instant(), t.toInstant());
                  return difference.isNegative() ? Duration.ZERO : difference;
                });
  }

  /**
   * Creates a function that tests an input string against a pattern before attempting to parse the
   * string.
   *
   * @param <T> the type generated by the parser
   */
  private static class PatternGuarded<T> implements Function<String, Optional<T>> {
    private final Pattern guard;
    private final Function<String, T> parser;

    public PatternGuarded(final String pattern, final Function<String, T> parser) {
      this.guard = Pattern.compile(pattern);
      this.parser = parser;
    }

    @Override
    public Optional<T> apply(final String header) {
      try {
        return guard.matcher(header).matches()
            ? Optional.of(parser.apply(header))
            : Optional.empty();
      } catch (final RuntimeException e) {
        LOGGER.warn("Failed to parse Retry-After header \"{}\"", forLog(header), e);
        return Optional.empty();
      }
    }
  }
}