A faster implementation with regex.

b/478848482

Author: weiminyu

core/src/main/java/google/registry/flows/FeeExtensionXmlTagNormalizer.java

@@ -14,35 +14,19 @@
 
 package google.registry.flows;
 
-import static com.google.common.collect.ImmutableList.toImmutableList;
-import static com.google.common.collect.Streams.stream;
-import static java.nio.charset.StandardCharsets.UTF_8;
+import static google.registry.model.eppcommon.ProtocolDefinition.ServiceExtension.FEE_0_11;
+import static google.registry.model.eppcommon.ProtocolDefinition.ServiceExtension.FEE_0_12;
+import static google.registry.model.eppcommon.ProtocolDefinition.ServiceExtension.FEE_0_6;
+import static google.registry.model.eppcommon.ProtocolDefinition.ServiceExtension.FEE_1_00;
 
-import com.google.common.base.CharMatcher;
-import com.google.common.collect.ImmutableList;
+import com.google.common.annotations.VisibleForTesting;
 import com.google.common.collect.ImmutableSet;
-import com.google.common.collect.Iterables;
-import com.google.common.flogger.FluentLogger;
-import java.io.ByteArrayInputStream;
-import java.io.ByteArrayOutputStream;
-import java.io.UnsupportedEncodingException;
-import java.util.Base64;
-import java.util.Iterator;
+import google.registry.model.eppcommon.EppXmlTransformer;
+import google.registry.model.eppcommon.ProtocolDefinition.ServiceExtension;
 import java.util.Optional;
-import java.util.Set;
+import java.util.regex.Matcher;
+import java.util.regex.Pattern;
 import java.util.stream.Collectors;
-import javax.xml.namespace.QName;
-import javax.xml.stream.XMLEventFactory;
-import javax.xml.stream.XMLEventReader;
-import javax.xml.stream.XMLEventWriter;
-import javax.xml.stream.XMLInputFactory;
-import javax.xml.stream.XMLOutputFactory;
-import javax.xml.stream.XMLStreamException;
-import javax.xml.stream.events.Attribute;
-import javax.xml.stream.events.EndElement;
-import javax.xml.stream.events.Namespace;
-import javax.xml.stream.events.StartElement;
-import javax.xml.stream.events.XMLEvent;
 
 /**
  * Normalizes Fee extension namespace tags in EPP XML response messages.
@@ -53,236 +37,92 @@
  *
  * <p>Some registrars are not XML namespace-aware and rely on the XML tags being specific literals.
  * This makes it difficult to perform seamless rollout of new versions: if Nomulus reassigns a tag
- * literal to a different version, it effectively forces all these registrars to upgrade.
+ * literal to a different version, it effectively forces all these registrars to upgrade at the time
+ * of the deployment.
  *
  * <p>This class can be used to normalize the namespace tag in EPP responses. Since every response
  * message may use at most one version of the Fee extension, we can remove declared but unused
  * versions from the message, thus freeing up the canonical tag ('fee') for the active version.
  */
 public class FeeExtensionXmlTagNormalizer {
 
-  private static final FluentLogger logger = FluentLogger.forEnclosingClass();
-
   // So far we only have Fee extensions to process
   private static final String CANONICAL_FEE_TAG = "fee";
-  private static final ImmutableSet FEE_EXTENSIONS =
-      ImmutableSet.of(
-          "urn:ietf:params:xml:ns:fee-0.6",
-          "urn:ietf:params:xml:ns:fee-0.11",
-          "urn:ietf:params:xml:ns:fee-0.12",
-          "urn:ietf:params:xml:ns:epp:fee-1.0");
-
-  private static final XMLInputFactory XML_INPUT_FACTORY = createXmlInputFactory();
-  private static final XMLOutputFactory XML_OUTPUT_FACTORY = XMLOutputFactory.newFactory();
-  private static final XMLEventFactory XML_EVENT_FACTORY = XMLEventFactory.newFactory();
-
-  /**
-   * Returns an EPP XML message with normalized Fee extension tags.
-   *
-   * <p>The output always begins with version and encoding declarations no matter if the input
-   * includes them. If encoding is not declared by input, UTF-8 will be used according to XML
-   * standard.
-   */
-  public static String normalizeFeeExtensionTag(byte[] inputXmlBytes) {
-    try {
-      // Keep exactly one newline at end of sanitized string.
-      return CharMatcher.whitespace().trimTrailingFrom(normalize(inputXmlBytes)) + "\n";
-    } catch (XMLStreamException | UnsupportedEncodingException e) {
-      logger.atWarning().withCause(e).log("Failed to sanitize EPP XML message.");
-      return Base64.getMimeEncoder().encodeToString(inputXmlBytes);
-    }
-  }
-
-  private static String normalize(byte[] inputXmlBytes)
-      throws XMLStreamException, UnsupportedEncodingException {
-    ParseResults parseResults = findFeeExtensionInUse(inputXmlBytes);
 
-    if (parseResults.feeExtensionInUse.isEmpty()) {
-      // Fee extension not present. Return as is.
-      return new String(inputXmlBytes, UTF_8);
-    }
-
-    ByteArrayOutputStream outputXmlBytes = new ByteArrayOutputStream();
-    XMLEventWriter xmlEventWriter =
-        XML_OUTPUT_FACTORY.createXMLEventWriter(outputXmlBytes, UTF_8.name());
+  private static final ImmutableSet<ServiceExtension> FEE_EXTENSIONS =
+      ImmutableSet.of(FEE_0_6, FEE_0_11, FEE_0_12, FEE_1_00);
 
-    for (XMLEvent event : parseResults.xmlEvents()) {
-      xmlEventWriter.add(normalizeXmlEvent(event, parseResults.feeExtensionInUse));
-      // Most standard Java StAX implementations omits the content between the XML header and the
-      // root element. Add a "\n" between them to improve readability.
-      if (event.isStartDocument()) {
-        xmlEventWriter.add(XML_EVENT_FACTORY.createCharacters("\n"));
-      }
-    }
+  private static final Pattern FEE_EXTENSION_IN_USE_PATTERN =
+      Pattern.compile(feeExtensionInUseRegex());
 
-    xmlEventWriter.flush();
-    return outputXmlBytes.toString(UTF_8);
+  @VisibleForTesting
+  static String feeExtensionInUseRegex() {
+    return FEE_EXTENSIONS.stream()
+        .map(ServiceExtension::getXmlTag)
+        .map(tag -> String.format("\\b(%s):", tag))
+        .collect(Collectors.joining("|"));
   }
 
   /**
-   * Holds intermediate results during XML processing.
-   *
-   * @param feeExtensionInUse The fee extension namespace URI in the EPP response, if found
-   * @param xmlEvents The parsed XML objects found in a pass, saved for reuse
-   */
-  private record ParseResults(
-      Optional<String> feeExtensionInUse, ImmutableList<XMLEvent> xmlEvents) {}
-
-  /**
-   * Makes one pass of the input XML and returns parsed data the Fee extension in use.
+   * Returns a EPP response that uses the canonical tag ({@code fee}) for the fee extension.
    *
-   * <p>Each XML message should use at most one Fee extension. This method returns it if found. The
-   * {@link XMLEvent} objects returned by the parser are also saved for reuse.
+   * <p>This method replaces any versioned tag, e.g., {@code fee12} with the canonical tag. It also
+   * removes unused namespace declarations and update the tag in the remaining declaration.
    *
-   * @throws IllegalArgumentException if more than one Fee extension version is found
+   * <p>The input {@code xml} must be an EPP response message generated by the {@link
+   * EppXmlTransformer}. With this assumption, we can use regular expressions which is 10X faster
+   * than XML stream parsers.
    */
-  private static ParseResults findFeeExtensionInUse(byte[] inputXmlBytes)
-      throws XMLStreamException {
-    XMLEventReader xmlEventReader =
-        XML_INPUT_FACTORY.createXMLEventReader(new ByteArrayInputStream(inputXmlBytes));
-
-    ImmutableList.Builder<XMLEvent> eventBuffer = new ImmutableList.Builder<>();
-    Optional<String> feeExtensionInUse = Optional.empty();
-
-    // Make one pass through the message to identify the Fee extension in use.
-    while (xmlEventReader.hasNext()) {
-      XMLEvent xmlEvent = xmlEventReader.nextEvent();
-      Optional<String> eventFeeExtensionUri = getXmlEventFeeExtensionUri(xmlEvent);
-
-      if (feeExtensionInUse.isEmpty()) {
-        feeExtensionInUse = eventFeeExtensionUri;
-      } else if (eventFeeExtensionUri.isPresent()
-          && !feeExtensionInUse.equals(eventFeeExtensionUri)) {
-        throw new IllegalArgumentException(
-            String.format(
-                "Expecting one Fee extension, found two: %s -- %s",
-                feeExtensionInUse, eventFeeExtensionUri.get()));
-      }
-      eventBuffer.add(xmlEvent);
-    }
-    return new ParseResults(feeExtensionInUse, eventBuffer.build());
-  }
-
-  private static XMLEvent normalizeXmlEvent(XMLEvent xmlEvent, Optional<String> feeExtensionInUse) {
-    if (xmlEvent.isStartElement()) {
-      return normalizeStartElement(xmlEvent.asStartElement(), feeExtensionInUse);
-    } else if (xmlEvent.isEndElement()) {
-      return normalizeEndElement(xmlEvent.asEndElement(), feeExtensionInUse);
-    } else {
-      return xmlEvent;
+  public static String normalize(String xml) {
+    Optional<String> maybeFeeTagInUse = findFeeExtensionInUse(xml);
+    if (maybeFeeTagInUse.isEmpty()) {
+      return xml;
     }
-  }
-
-  private static Optional<String> getXmlEventFeeExtensionUri(XMLEvent xmlEvent) {
-    if (xmlEvent.isStartElement()) {
-      return getFeeExtensionUri(xmlEvent.asStartElement());
-    }
-    if (xmlEvent.isEndElement()) {
-      String extension = xmlEvent.asEndElement().getName().getNamespaceURI();
-      if (FEE_EXTENSIONS.contains(extension)) {
-        return Optional.of(extension);
+    String feeTagInUse = maybeFeeTagInUse.get();
+    String normalized = xml;
+    for (ServiceExtension serviceExtension : FEE_EXTENSIONS) {
+      if (serviceExtension.getXmlTag().equals(feeTagInUse)) {
+        normalized = normalizeExtensionInUse(feeTagInUse, serviceExtension.getUri(), normalized);
+      } else {
+        normalized =
+            removeUnusedExtension(
+                serviceExtension.getXmlTag(), serviceExtension.getUri(), normalized);
       }
     }
-    return Optional.empty();
-  }
-
-  private static Optional<String> getFeeExtensionUri(StartElement startElement) {
-    Set<String> attrs =
-        stream(startElement.asStartElement().getAttributes())
-            .map(Attribute::getName)
-            .map(FeeExtensionXmlTagNormalizer::getFeeExtensionUri)
-            .flatMap(Optional::stream)
-            .collect(Collectors.toSet());
-    var qName = startElement.asStartElement().getName();
-    if (FEE_EXTENSIONS.contains(qName.getNamespaceURI())) {
-      attrs.add(qName.getNamespaceURI());
-    }
-    if (attrs.size() > 1) {
-      throw new IllegalArgumentException("Multiple Fee extension in use: " + attrs);
-    }
-    if (attrs.isEmpty()) {
-      return Optional.empty();
-    }
-    // attrs.size == 1
-    return Optional.of(Iterables.getOnlyElement(attrs));
-  }
-
-  private static Optional<String> getFeeExtensionUri(QName name) {
-    String extensionUri = name.getNamespaceURI();
-    if (FEE_EXTENSIONS.contains(extensionUri)) {
-      return Optional.of(extensionUri);
-    }
-    return Optional.empty();
+    return normalized;
   }
 
-  private static XMLEvent normalizeStartElement(
-      StartElement startElement, Optional<String> feeExtensionInUse) {
-    QName name = normalizeName(startElement.getName());
-    ImmutableList<Namespace> namespaces =
-        normalizeNamespaces(startElement.getNamespaces(), feeExtensionInUse);
-    ImmutableList<Attribute> attributes = normalizeAttributes(startElement.getAttributes());
-
-    return XML_EVENT_FACTORY.createStartElement(name, attributes.iterator(), namespaces.iterator());
-  }
-
-  private static XMLEvent normalizeEndElement(
-      EndElement endElement, Optional<String> feeExtensionInUse) {
-    QName name = normalizeName(endElement.getName());
-    ImmutableList<Namespace> namespaces =
-        normalizeNamespaces(endElement.getNamespaces(), feeExtensionInUse);
-
-    return XML_EVENT_FACTORY.createEndElement(name, namespaces.iterator());
+  static String removeUnusedExtension(String tag, String uri, String xml) {
+    String declaration = String.format("xmlns:%s=\"%s\"", tag, uri);
+    // There must be a leading whitespace, and it can be safely removed with the declaration.
+    return xml.replaceAll(String.format("\\s%s", declaration), "");
   }
 
-  private static QName normalizeName(QName name) {
-    if (!FEE_EXTENSIONS.contains(name.getNamespaceURI())
-        || name.getPrefix().equals(CANONICAL_FEE_TAG)) {
-      return name;
+  static String normalizeExtensionInUse(String tagInUse, String uriInUse, String xml) {
+    if (tagInUse.equals(CANONICAL_FEE_TAG)) {
+      return xml;
     }
-    return new QName(name.getNamespaceURI(), name.getLocalPart(), CANONICAL_FEE_TAG);
+    // Change the tag in the namespace declaration:
+    String currentDeclaration = String.format("xmlns:%s=\"%s\"", tagInUse, uriInUse);
+    String desiredDeclaraion = String.format("xmlns:fee=\"%s\"", uriInUse);
+    // The new tag at each site of use, with trailing colon:
+    String newTagWithColon = CANONICAL_FEE_TAG + ":";
+    return xml.replaceAll(String.format("\\b%s:", tagInUse), newTagWithColon)
+        .replaceAll(currentDeclaration, desiredDeclaraion);
   }
 
-  private static Attribute normalizeAttribute(Attribute attribute) {
-    QName name = normalizeName(attribute.getName());
-    return XML_EVENT_FACTORY.createAttribute(name, attribute.getValue());
-  }
+  static Optional<String> findFeeExtensionInUse(String xml) {
+    Matcher matcher = FEE_EXTENSION_IN_USE_PATTERN.matcher(xml);
 
-  private static Optional<Namespace> normalizeNamespace(
-      Namespace namespace, Optional<String> feeExtensionInUse) {
-    var extension = namespace.getNamespaceURI();
-    if (!FEE_EXTENSIONS.contains(extension)) {
-      return Optional.of(namespace);
+    if (!matcher.find()) {
+      return Optional.empty();
     }
-    if (feeExtensionInUse.isPresent() && extension.equals(feeExtensionInUse.get())) {
-      if (namespace.getPrefix().equals(CANONICAL_FEE_TAG)) {
-        return Optional.of(namespace);
+    // We know only one extension is in use, so we can return on the first match
+    for (int i = 1; i <= matcher.groupCount(); i++) {
+      if (matcher.group(i) != null) {
+        return Optional.of(matcher.group(i));
       }
-      return Optional.of(XML_EVENT_FACTORY.createNamespace(CANONICAL_FEE_TAG, extension));
     }
-    return Optional.empty();
-  }
-
-  private static ImmutableList<Attribute> normalizeAttributes(Iterator<Attribute> attributes) {
-    return stream(attributes).map(attr -> normalizeAttribute(attr)).collect(toImmutableList());
-  }
-
-  private static ImmutableList<Namespace> normalizeNamespaces(
-      Iterator<Namespace> namespaces, Optional<String> feeExtensionInUse) {
-    return stream(namespaces)
-        .map(namespace -> normalizeNamespace(namespace, feeExtensionInUse))
-        .flatMap(Optional::stream)
-        .collect(toImmutableList());
-  }
-
-  private static XMLInputFactory createXmlInputFactory() {
-    XMLInputFactory xmlInputFactory = XMLInputFactory.newFactory();
-    // Coalesce adjacent data, so that all chars in a string will be grouped as one item.
-    xmlInputFactory.setProperty(XMLInputFactory.IS_COALESCING, true);
-    // Preserve Name Space information.
-    xmlInputFactory.setProperty(XMLInputFactory.IS_NAMESPACE_AWARE, true);
-    // Prevent XXE attacks.
-    xmlInputFactory.setProperty(XMLInputFactory.IS_SUPPORTING_EXTERNAL_ENTITIES, false);
-    xmlInputFactory.setProperty(XMLInputFactory.SUPPORT_DTD, false);
-    return xmlInputFactory;
+    throw new IllegalStateException("Should not reach here. Bad FEE_EXTENSION_IN_USE_PATTERN?");
   }
 }

core/src/main/java/google/registry/model/eppcommon/ProtocolDefinition.java

@@ -18,6 +18,7 @@
 import static com.google.common.collect.Maps.uniqueIndex;
 
 import com.google.common.annotations.VisibleForTesting;
+import com.google.common.base.VerifyException;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
 import google.registry.model.domain.fee06.FeeCheckCommandExtensionV06;
@@ -87,6 +88,7 @@ public enum ServiceExtension {
     private final Class<? extends CommandExtension> commandExtensionClass;
     private final Class<? extends ResponseExtension> responseExtensionClass;
     private final String uri;
+    private final String xmlTag;
     private final ServiceExtensionVisibility visibility;
 
     ServiceExtension(
@@ -96,6 +98,7 @@ public enum ServiceExtension {
       this.commandExtensionClass = commandExtensionClass;
       this.responseExtensionClass = responseExtensionClass;
       this.uri = getCommandExtensionUri(commandExtensionClass);
+      this.xmlTag = getCommandExtensionXmlTag(commandExtensionClass);
       this.visibility = visibility;
     }
 
@@ -111,11 +114,27 @@ public String getUri() {
       return uri;
     }
 
+    public String getXmlTag() {
+      return xmlTag;
+    }
+
     /** Returns the namespace URI of the command extension class. */
     public static String getCommandExtensionUri(Class<? extends CommandExtension> clazz) {
       return clazz.getPackage().getAnnotation(XmlSchema.class).namespace();
     }
 
+    /** Returns the XML tag for this extension in the response message. */
+    public static String getCommandExtensionXmlTag(Class<? extends CommandExtension> clazz) {
+      var xmlSchema = clazz.getPackage().getAnnotation(XmlSchema.class);
+      var xmlns = xmlSchema.xmlns();
+      if (xmlns == null || xmlns.length != 1) {
+        throw new VerifyException(
+            String.format(
+                "Expecting exactly one NS declaration in %s", clazz.getPackage().getName()));
+      }
+      return xmlns[0].prefix();
+    }
+
     public boolean isVisible() {
       return switch (visibility) {
         case ALL -> true;