Simplify the fee extensions.

I added shared base classes to all of the Fee extension types that
make it possible to fully ignore the version in the flows. (You
ask for a FeeCreateCommandExtension, for example, and you get one
without having to worry about which). This is an improvement over
the old code that asked you to provide a list of possible fee
extensions and then ask for the first one in preference order.

As part of this I was able to make the Fee implementation a bit
simpler as well.

-------------
Created by MOE: https://github.com/google/moe
MOE_MIGRATED_REVID=137992390

java/google/registry/flows/ExtensionManager.java

@@ -14,12 +14,16 @@
 
 package google.registry.flows;
 
+import static com.google.common.collect.Iterables.any;
 import static com.google.common.collect.Sets.difference;
 import static com.google.common.collect.Sets.intersection;
 import static google.registry.model.domain.fee.Fee.FEE_EXTENSION_URIS;
 import static google.registry.model.eppcommon.ProtocolDefinition.ServiceExtension.getCommandExtensionUri;
 
 import com.google.common.base.Joiner;
+import com.google.common.base.Predicate;
+import com.google.common.collect.FluentIterable;
+import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableSet;
 import google.registry.flows.EppException.CommandUseErrorException;
 import google.registry.flows.EppException.SyntaxErrorException;
@@ -30,7 +34,6 @@ import google.registry.model.domain.metadata.MetadataExtension;
 import google.registry.model.eppinput.EppInput;
 import google.registry.model.eppinput.EppInput.CommandExtension;
 import google.registry.util.FormattingLogger;
-import java.util.Collection;
 import java.util.Set;
 import javax.inject.Inject;
 
@@ -49,8 +52,6 @@ public final class ExtensionManager {
 
   private final ImmutableSet.Builder<Class<? extends CommandExtension>> implementedBuilder =
       new ImmutableSet.Builder<>();
-  private final ImmutableSet.Builder<ImmutableSet<?>> implementedGroupsBuilder =
-      new ImmutableSet.Builder<>();
 
   @Inject EppInput eppInput;
   @Inject SessionMetadata sessionMetadata;
@@ -64,12 +65,6 @@ public final class ExtensionManager {
     implementedBuilder.add(extension);
   }
 
-  public <T extends CommandExtension> void registerAsGroup(
-      Collection<Class<? extends T>> extensions) {
-    implementedBuilder.addAll(extensions);
-    implementedGroupsBuilder.add(ImmutableSet.copyOf(extensions));
-  }
-
   public void validate() throws EppException {
     ImmutableSet.Builder<Class<? extends CommandExtension>> suppliedBuilder =
         new ImmutableSet.Builder<>();
@@ -79,27 +74,12 @@ public final class ExtensionManager {
     ImmutableSet<Class<? extends CommandExtension>> suppliedExtensions = suppliedBuilder.build();
     ImmutableSet<Class<? extends CommandExtension>> implementedExtensions =
         implementedBuilder.build();
-    ImmutableSet<ImmutableSet<?>> implementedExtensionGroups =
-        implementedGroupsBuilder.build();
-    checkForDuplicateExtensions(suppliedExtensions, implementedExtensionGroups);
+    ImmutableList<CommandExtension> suppliedExtensionInstances =
+        eppInput.getCommandWrapper().getExtensions();
     checkForUndeclaredExtensions(suppliedExtensions);
     checkForRestrictedExtensions(suppliedExtensions);
-    checkForUnimplementedExtensions(suppliedExtensions, implementedExtensions);
-  }
-
-  private void checkForDuplicateExtensions(
-      ImmutableSet<Class<? extends CommandExtension>> suppliedExtensions,
-      ImmutableSet<ImmutableSet<?>> implementedExtensionGroups)
-          throws UnsupportedRepeatedExtensionException {
-    if (suppliedExtensions.size() < eppInput.getCommandWrapper().getExtensions().size()) {
-      throw new UnsupportedRepeatedExtensionException();
-    }
-    // No more than one extension in an extension group can be present.
-    for (ImmutableSet<?> group : implementedExtensionGroups) {
-      if (intersection(suppliedExtensions, group).size() > 1) {
-        throw new UnsupportedRepeatedExtensionException();
-      }
-    }
+    checkForDuplicateExtensions(suppliedExtensionInstances, suppliedExtensions);
+    checkForUnimplementedExtensions(suppliedExtensionInstances, implementedExtensions);
   }
 
   private void checkForUndeclaredExtensions(
@@ -134,22 +114,39 @@ public final class ExtensionManager {
     }
   }
 
-  private void checkForUnimplementedExtensions(
-      ImmutableSet<Class<? extends CommandExtension>> suppliedExtensions,
+  private static void checkForDuplicateExtensions(
+      ImmutableList<CommandExtension> suppliedExtensionInstances,
       ImmutableSet<Class<? extends CommandExtension>> implementedExtensions)
-          throws UnimplementedExtensionException {
-    Set<Class<? extends CommandExtension>> unimplementedExtensions =
-        difference(suppliedExtensions, implementedExtensions);
-    if (!unimplementedExtensions.isEmpty()) {
-      logger.infofmt("Unimplemented extensions: %s", unimplementedExtensions);
-      throw new UnimplementedExtensionException();
+          throws UnsupportedRepeatedExtensionException {
+    for (Class<? extends CommandExtension> implemented : implementedExtensions) {
+      if (FluentIterable.from(suppliedExtensionInstances).filter(implemented).size() > 1) {
+        throw new UnsupportedRepeatedExtensionException();
+      }
     }
   }
 
-  /** Unsupported repetition of an extension. */
-  static class UnsupportedRepeatedExtensionException extends SyntaxErrorException {
-    public UnsupportedRepeatedExtensionException() {
-      super("Unsupported repetition of an extension");
+  private static void checkForUnimplementedExtensions(
+      ImmutableList<CommandExtension> suppliedExtensionInstances,
+      ImmutableSet<Class<? extends CommandExtension>> implementedExtensionClasses)
+          throws UnimplementedExtensionException {
+    ImmutableSet.Builder<Class<? extends CommandExtension>> unimplementedExtensionsBuilder =
+        new ImmutableSet.Builder<>();
+    for (final CommandExtension instance : suppliedExtensionInstances) {
+      if (!any(
+          implementedExtensionClasses,
+          new Predicate<Class<? extends CommandExtension>>() {
+            @Override
+            public boolean apply(Class<? extends CommandExtension> implementedExtensionClass) {
+              return implementedExtensionClass.isInstance(instance);
+            }})) {
+        unimplementedExtensionsBuilder.add(instance.getClass());
+      }
+    }
+    ImmutableSet<Class<? extends CommandExtension>> unimplementedExtensions =
+        unimplementedExtensionsBuilder.build();
+    if (!unimplementedExtensions.isEmpty()) {
+      logger.infofmt("Unimplemented extensions: %s", unimplementedExtensions);
+      throw new UnimplementedExtensionException();
     }
   }
 
@@ -160,4 +157,11 @@ public final class ExtensionManager {
             Joiner.on(", ").join(undeclaredUris)));
     }
   }
+
+  /** Unsupported repetition of an extension. */
+  static class UnsupportedRepeatedExtensionException extends SyntaxErrorException {
+    public UnsupportedRepeatedExtensionException() {
+      super("Unsupported repetition of an extension");
+    }
+  }
 }

java/google/registry/flows/domain/DomainAllocateFlow.java

@@ -30,7 +30,6 @@ import static google.registry.flows.domain.DomainFlowUtils.validateSecDnsExtensi
 import static google.registry.flows.domain.DomainFlowUtils.verifyUnitIsYears;
 import static google.registry.model.EppResourceUtils.createDomainRoid;
 import static google.registry.model.EppResourceUtils.loadDomainApplication;
-import static google.registry.model.domain.fee.Fee.FEE_CREATE_COMMAND_EXTENSIONS_IN_PREFERENCE_ORDER;
 import static google.registry.model.ofy.ObjectifyService.ofy;
 import static google.registry.model.registry.label.ReservedList.matchesAnchorTenantReservation;
 import static google.registry.pricing.PricingEngineProxy.getDomainCreateCost;
@@ -63,7 +62,7 @@ import google.registry.model.domain.DomainResource;
 import google.registry.model.domain.GracePeriod;
 import google.registry.model.domain.Period;
 import google.registry.model.domain.allocate.AllocateCreateExtension;
-import google.registry.model.domain.fee.FeeTransformCommandExtension;
+import google.registry.model.domain.fee.FeeCreateCommandExtension;
 import google.registry.model.domain.fee.FeeTransformResponseExtension;
 import google.registry.model.domain.flags.FlagsCreateCommandExtension;
 import google.registry.model.domain.launch.ApplicationStatus;
@@ -118,11 +117,11 @@ public class DomainAllocateFlow implements TransactionalFlow {
   @Override
   public final EppResponse run() throws EppException {
     extensionManager.register(
+        FeeCreateCommandExtension.class,
         SecDnsCreateExtension.class,
         FlagsCreateCommandExtension.class,
         MetadataExtension.class,
         AllocateCreateExtension.class);
-    extensionManager.registerAsGroup(FEE_CREATE_COMMAND_EXTENSIONS_IN_PREFERENCE_ORDER);
     extensionManager.validate();
     validateClientIsLoggedIn(clientId);
     verifyIsSuperuser();
@@ -374,8 +373,8 @@ public class DomainAllocateFlow implements TransactionalFlow {
       DateTime now, Registry registry, int years) throws EppException {
     EppCommandOperations commandOperations = TldSpecificLogicProxy.getCreatePrice(
         registry, targetId, clientId, now, years, eppInput);
-    FeeTransformCommandExtension feeCreate =
-        eppInput.getFirstExtensionOfClasses(FEE_CREATE_COMMAND_EXTENSIONS_IN_PREFERENCE_ORDER);
+    FeeCreateCommandExtension feeCreate =
+        eppInput.getSingleExtension(FeeCreateCommandExtension.class);
     return (feeCreate == null)
         ? null
         : ImmutableList.of(createFeeCreateResponse(feeCreate, commandOperations));

java/google/registry/flows/domain/DomainApplicationCreateFlow.java

@@ -38,7 +38,6 @@ import static google.registry.flows.domain.DomainFlowUtils.verifyRegistryStateAl
 import static google.registry.flows.domain.DomainFlowUtils.verifySignedMarks;
 import static google.registry.flows.domain.DomainFlowUtils.verifyUnitIsYears;
 import static google.registry.model.EppResourceUtils.createDomainRoid;
-import static google.registry.model.domain.fee.Fee.FEE_CREATE_COMMAND_EXTENSIONS_IN_PREFERENCE_ORDER;
 import static google.registry.model.index.DomainApplicationIndex.loadActiveApplicationsByDomainName;
 import static google.registry.model.ofy.ObjectifyService.ofy;
 import static google.registry.model.registry.label.ReservedList.matchesAnchorTenantReservation;
@@ -65,6 +64,7 @@ import google.registry.model.domain.DomainApplication;
 import google.registry.model.domain.DomainCommand.Create;
 import google.registry.model.domain.DomainResource;
 import google.registry.model.domain.Period;
+import google.registry.model.domain.fee.FeeCreateCommandExtension;
 import google.registry.model.domain.fee.FeeTransformCommandExtension;
 import google.registry.model.domain.flags.FlagsCreateCommandExtension;
 import google.registry.model.domain.launch.ApplicationStatus;
@@ -172,11 +172,11 @@ public final class DomainApplicationCreateFlow implements TransactionalFlow {
   @Override
   public final EppResponse run() throws EppException {
     extensionManager.register(
+        FeeCreateCommandExtension.class,
         SecDnsCreateExtension.class,
         FlagsCreateCommandExtension.class,
         MetadataExtension.class,
         LaunchCreateExtension.class);
-    extensionManager.registerAsGroup(FEE_CREATE_COMMAND_EXTENSIONS_IN_PREFERENCE_ORDER);
     extensionManager.validate();
     validateClientIsLoggedIn(clientId);
     DateTime now = ofy().getTransactionTime();
@@ -212,8 +212,8 @@ public final class DomainApplicationCreateFlow implements TransactionalFlow {
         verifyNotReserved(domainName, isSunriseApplication);
       }
     }
-    FeeTransformCommandExtension feeCreate =
-        eppInput.getFirstExtensionOfClasses(FEE_CREATE_COMMAND_EXTENSIONS_IN_PREFERENCE_ORDER);
+    FeeCreateCommandExtension feeCreate =
+        eppInput.getSingleExtension(FeeCreateCommandExtension.class);
     validateFeeChallenge(targetId, tld, now, feeCreate, commandOperations.getTotalCost());
     SecDnsCreateExtension secDnsCreate =
         validateSecDnsExtension(eppInput.getSingleExtension(SecDnsCreateExtension.class));

java/google/registry/flows/domain/DomainApplicationUpdateFlow.java

@@ -39,7 +39,6 @@ import static google.registry.flows.domain.DomainFlowUtils.verifyApplicationDoma
 import static google.registry.flows.domain.DomainFlowUtils.verifyClientUpdateNotProhibited;
 import static google.registry.flows.domain.DomainFlowUtils.verifyNotInPendingDelete;
 import static google.registry.model.EppResourceUtils.loadDomainApplication;
-import static google.registry.model.domain.fee.Fee.FEE_UPDATE_COMMAND_EXTENSIONS_IN_PREFERENCE_ORDER;
 import static google.registry.model.ofy.ObjectifyService.ofy;
 
 import com.google.common.base.Optional;
@@ -138,7 +137,6 @@ public class DomainApplicationUpdateFlow implements TransactionalFlow {
         LaunchUpdateExtension.class,
         MetadataExtension.class,
         SecDnsUpdateExtension.class);
-    extensionManager.registerAsGroup(FEE_UPDATE_COMMAND_EXTENSIONS_IN_PREFERENCE_ORDER);
     extensionManager.validate();
     validateClientIsLoggedIn(clientId);
     DateTime now = ofy().getTransactionTime();

java/google/registry/flows/domain/DomainCheckFlow.java

@@ -23,7 +23,6 @@ import static google.registry.flows.domain.DomainFlowUtils.validateDomainName;
 import static google.registry.flows.domain.DomainFlowUtils.validateDomainNameWithIdnTables;
 import static google.registry.flows.domain.DomainFlowUtils.verifyNotInPredelegation;
 import static google.registry.model.EppResourceUtils.checkResourcesExist;
-import static google.registry.model.domain.fee.Fee.FEE_CHECK_COMMAND_EXTENSIONS_IN_PREFERENCE_ORDER;
 import static google.registry.model.index.DomainApplicationIndex.loadActiveApplicationsByDomainName;
 import static google.registry.model.registry.label.ReservationType.UNRESERVED;
 import static google.registry.pricing.PricingEngineProxy.isDomainPremium;
@@ -111,8 +110,7 @@ public final class DomainCheckFlow implements Flow {
 
   @Override
   public EppResponse run() throws EppException {
-    extensionManager.register(LaunchCheckExtension.class);
-    extensionManager.registerAsGroup(FEE_CHECK_COMMAND_EXTENSIONS_IN_PREFERENCE_ORDER);
+    extensionManager.register(FeeCheckCommandExtension.class, LaunchCheckExtension.class);
     extensionManager.validate();
     validateClientIsLoggedIn(clientId);
     List<String> targetIds = ((Check) resourceCommand).getTargetIds();
@@ -177,7 +175,7 @@ public final class DomainCheckFlow implements Flow {
   private ImmutableList<? extends ResponseExtension> getResponseExtensions(
       ImmutableMap<String, InternetDomainName> domainNames, DateTime now) throws EppException {
     FeeCheckCommandExtension<?, ?> feeCheck =
-        eppInput.getFirstExtensionOfClasses(FEE_CHECK_COMMAND_EXTENSIONS_IN_PREFERENCE_ORDER);
+        eppInput.getSingleExtension(FeeCheckCommandExtension.class);
     if (feeCheck == null) {
       return null;  // No fee checks were requested.
     }
@@ -185,7 +183,7 @@ public final class DomainCheckFlow implements Flow {
         new ImmutableList.Builder<>();
     for (FeeCheckCommandExtensionItem feeCheckItem : feeCheck.getItems()) {
       for (String domainName : getDomainNamesToCheckForFee(feeCheckItem, domainNames.keySet())) {
-        FeeCheckResponseExtensionItem.Builder builder = feeCheckItem.createResponseBuilder();
+        FeeCheckResponseExtensionItem.Builder<?> builder = feeCheckItem.createResponseBuilder();
         handleFeeRequest(
             feeCheckItem,
             builder,

java/google/registry/flows/domain/DomainCreateFlow.java

@@ -36,7 +36,6 @@ import static google.registry.flows.domain.DomainFlowUtils.verifyPremiumNameIsNo
 import static google.registry.flows.domain.DomainFlowUtils.verifySignedMarks;
 import static google.registry.flows.domain.DomainFlowUtils.verifyUnitIsYears;
 import static google.registry.model.EppResourceUtils.createDomainRoid;
-import static google.registry.model.domain.fee.Fee.FEE_CREATE_COMMAND_EXTENSIONS_IN_PREFERENCE_ORDER;
 import static google.registry.model.index.DomainApplicationIndex.loadActiveApplicationsByDomainName;
 import static google.registry.model.ofy.ObjectifyService.ofy;
 import static google.registry.model.registry.label.ReservedList.matchesAnchorTenantReservation;
@@ -70,7 +69,7 @@ import google.registry.model.domain.DomainCommand.Create;
 import google.registry.model.domain.DomainResource;
 import google.registry.model.domain.GracePeriod;
 import google.registry.model.domain.Period;
-import google.registry.model.domain.fee.FeeTransformCommandExtension;
+import google.registry.model.domain.fee.FeeCreateCommandExtension;
 import google.registry.model.domain.fee.FeeTransformResponseExtension;
 import google.registry.model.domain.flags.FlagsCreateCommandExtension;
 import google.registry.model.domain.launch.LaunchCreateExtension;
@@ -168,11 +167,11 @@ public class DomainCreateFlow implements TransactionalFlow {
   @Override
   public final EppResponse run() throws EppException {
     extensionManager.register(
+        FeeCreateCommandExtension.class,
         SecDnsCreateExtension.class,
         FlagsCreateCommandExtension.class,
         MetadataExtension.class,
         LaunchCreateExtension.class);
-    extensionManager.registerAsGroup(FEE_CREATE_COMMAND_EXTENSIONS_IN_PREFERENCE_ORDER);
     extensionManager.validate();
     validateClientIsLoggedIn(clientId);
     DateTime now = ofy().getTransactionTime();
@@ -199,8 +198,8 @@ public class DomainCreateFlow implements TransactionalFlow {
     if (hasSignedMarks) {
       verifySignedMarksAllowed(tldState, isAnchorTenant);
     }
-    FeeTransformCommandExtension feeCreate =
-        eppInput.getFirstExtensionOfClasses(FEE_CREATE_COMMAND_EXTENSIONS_IN_PREFERENCE_ORDER);
+    FeeCreateCommandExtension feeCreate =
+        eppInput.getSingleExtension(FeeCreateCommandExtension.class);
     EppCommandOperations commandOperations = TldSpecificLogicProxy.getCreatePrice(
         registry, targetId, clientId, now, years, eppInput);
     validateFeeChallenge(
@@ -423,7 +422,7 @@ public class DomainCreateFlow implements TransactionalFlow {
   }
 
   private ImmutableList<FeeTransformResponseExtension> createResponseExtensions(
-      FeeTransformCommandExtension feeCreate, EppCommandOperations commandOperations) {
+      FeeCreateCommandExtension feeCreate, EppCommandOperations commandOperations) {
     return (feeCreate == null)
         ? null
         : ImmutableList.of(createFeeCreateResponse(feeCreate, commandOperations));

java/google/registry/flows/domain/DomainFlowUtils.java

@@ -584,7 +584,7 @@ public class DomainFlowUtils {
    */
   static void handleFeeRequest(
       FeeQueryCommandExtensionItem feeRequest,
-      FeeQueryResponseExtensionItem.Builder builder,
+      FeeQueryResponseExtensionItem.Builder<?, ?> builder,
       InternetDomainName domain,
       String clientId,
       @Nullable CurrencyUnit topLevelCurrency,
@@ -617,7 +617,7 @@ public class DomainFlowUtils {
         .setPeriod(feeRequest.getPeriod())
         .setClass(TldSpecificLogicProxy.getFeeClass(domainNameString, now).orNull());
 
-    List<Fee> fees = ImmutableList.of();
+    ImmutableList<Fee> fees = ImmutableList.of();
     switch (feeRequest.getCommandName()) {
       case CREATE:
         if (isReserved(domain, isSunrise)) {  // Don't return a create price for reserved names.

java/google/registry/flows/domain/DomainRenewFlow.java

@@ -27,7 +27,6 @@ import static google.registry.flows.domain.DomainFlowUtils.validateFeeChallenge;
 import static google.registry.flows.domain.DomainFlowUtils.verifyUnitIsYears;
 import static google.registry.model.domain.DomainResource.MAX_REGISTRATION_YEARS;
 import static google.registry.model.domain.DomainResource.extendRegistrationWithCap;
-import static google.registry.model.domain.fee.Fee.FEE_RENEW_COMMAND_EXTENSIONS_IN_PREFERENCE_ORDER;
 import static google.registry.model.ofy.ObjectifyService.ofy;
 import static google.registry.util.DateTimeUtils.leapSafeAddYears;
 
@@ -53,7 +52,7 @@ import google.registry.model.domain.DomainResource;
 import google.registry.model.domain.GracePeriod;
 import google.registry.model.domain.fee.BaseFee.FeeType;
 import google.registry.model.domain.fee.Fee;
-import google.registry.model.domain.fee.FeeTransformCommandExtension;
+import google.registry.model.domain.fee.FeeRenewCommandExtension;
 import google.registry.model.domain.fee.FeeTransformResponseExtension;
 import google.registry.model.domain.metadata.MetadataExtension;
 import google.registry.model.domain.rgp.GracePeriodStatus;
@@ -117,8 +116,7 @@ public final class DomainRenewFlow implements TransactionalFlow {
 
   @Override
   public final EppResponse run() throws EppException {
-    extensionManager.register(MetadataExtension.class);
-    extensionManager.registerAsGroup(FEE_RENEW_COMMAND_EXTENSIONS_IN_PREFERENCE_ORDER);
+    extensionManager.register(FeeRenewCommandExtension.class, MetadataExtension.class);
     extensionManager.validate();
     validateClientIsLoggedIn(clientId);
     DateTime now = ofy().getTransactionTime();
@@ -127,8 +125,8 @@ public final class DomainRenewFlow implements TransactionalFlow {
     DomainResource existingDomain = loadAndVerifyExistence(DomainResource.class, targetId, now);
     verifyRenewAllowed(authInfo, existingDomain, command);
     int years = command.getPeriod().getValue();
-    FeeTransformCommandExtension feeRenew =
-        eppInput.getFirstExtensionOfClasses(FEE_RENEW_COMMAND_EXTENSIONS_IN_PREFERENCE_ORDER);
+    FeeRenewCommandExtension feeRenew =
+        eppInput.getSingleExtension(FeeRenewCommandExtension.class);
     EppCommandOperations commandOperations = TldSpecificLogicProxy.getRenewPrice(
         Registry.get(existingDomain.getTld()), targetId, clientId, now, years, eppInput);
     validateFeeChallenge(
@@ -217,7 +215,7 @@ public final class DomainRenewFlow implements TransactionalFlow {
   }
 
   private ImmutableList<FeeTransformResponseExtension> createResponseExtensions(
-      Money renewCost, FeeTransformCommandExtension feeRenew) {
+      Money renewCost, FeeRenewCommandExtension feeRenew) {
     return (feeRenew == null) ? null : ImmutableList.of(feeRenew
         .createResponseBuilder()
         .setCurrency(renewCost.getCurrencyUnit())

java/google/registry/flows/domain/DomainRestoreRequestFlow.java

@@ -25,7 +25,6 @@ import static google.registry.flows.domain.DomainFlowUtils.newAutorenewPollMessa
 import static google.registry.flows.domain.DomainFlowUtils.validateFeeChallenge;
 import static google.registry.flows.domain.DomainFlowUtils.verifyNotReserved;
 import static google.registry.flows.domain.DomainFlowUtils.verifyPremiumNameIsNotBlocked;
-import static google.registry.model.domain.fee.Fee.FEE_UPDATE_COMMAND_EXTENSIONS_IN_PREFERENCE_ORDER;
 import static google.registry.model.ofy.ObjectifyService.ofy;
 import static google.registry.util.DateTimeUtils.END_OF_TIME;
 
@@ -53,8 +52,8 @@ import google.registry.model.domain.DomainCommand.Update;
 import google.registry.model.domain.DomainResource;
 import google.registry.model.domain.fee.BaseFee.FeeType;
 import google.registry.model.domain.fee.Fee;
-import google.registry.model.domain.fee.FeeTransformCommandExtension;
 import google.registry.model.domain.fee.FeeTransformResponseExtension;
+import google.registry.model.domain.fee.FeeUpdateCommandExtension;
 import google.registry.model.domain.metadata.MetadataExtension;
 import google.registry.model.domain.rgp.GracePeriodStatus;
 import google.registry.model.domain.rgp.RgpUpdateExtension;
@@ -119,8 +118,10 @@ public final class DomainRestoreRequestFlow implements TransactionalFlow  {
 
   @Override
   public final EppResponse run() throws EppException {
-    extensionManager.register(MetadataExtension.class, RgpUpdateExtension.class);
-    extensionManager.registerAsGroup(FEE_UPDATE_COMMAND_EXTENSIONS_IN_PREFERENCE_ORDER);
+    extensionManager.register(
+        FeeUpdateCommandExtension.class,
+        MetadataExtension.class,
+        RgpUpdateExtension.class);
     extensionManager.validate();
     validateClientIsLoggedIn(clientId);
     Update command = (Update) resourceCommand;
@@ -129,8 +130,8 @@ public final class DomainRestoreRequestFlow implements TransactionalFlow  {
     Money restoreCost = Registry.get(existingDomain.getTld()).getStandardRestoreCost();
     EppCommandOperations renewCommandOperations = TldSpecificLogicProxy.getRenewPrice(
         Registry.get(existingDomain.getTld()), targetId, clientId, now, 1, eppInput);
-    FeeTransformCommandExtension feeUpdate = eppInput.getFirstExtensionOfClasses(
-        FEE_UPDATE_COMMAND_EXTENSIONS_IN_PREFERENCE_ORDER);
+    FeeUpdateCommandExtension feeUpdate =
+        eppInput.getSingleExtension(FeeUpdateCommandExtension.class);
     Money totalCost = renewCommandOperations.getTotalCost();
     verifyRestoreAllowed(command, existingDomain, restoreCost, totalCost, feeUpdate, now);
     HistoryEntry historyEntry = buildHistory(existingDomain, now);
@@ -184,7 +185,7 @@ public final class DomainRestoreRequestFlow implements TransactionalFlow  {
       DomainResource existingDomain,
       Money restoreCost,
       Money renewCost,
-      FeeTransformCommandExtension feeUpdate,
+      FeeUpdateCommandExtension feeUpdate,
       DateTime now) throws EppException {
     verifyOptionalAuthInfo(authInfo, existingDomain);
     if (!isSuperuser) {
@@ -258,7 +259,7 @@ public final class DomainRestoreRequestFlow implements TransactionalFlow  {
   }
 
   private static ImmutableList<FeeTransformResponseExtension> createResponseExtensions(
-      Money restoreCost, Money renewCost, FeeTransformCommandExtension feeUpdate) {
+      Money restoreCost, Money renewCost, FeeUpdateCommandExtension feeUpdate) {
     return (feeUpdate == null) ? null : ImmutableList.of(
         feeUpdate.createResponseBuilder()
             .setCurrency(restoreCost.getCurrencyUnit())

java/google/registry/flows/domain/DomainTransferRequestFlow.java

@@ -30,7 +30,6 @@ import static google.registry.flows.domain.DomainFlowUtils.validateFeeChallenge;
 import static google.registry.flows.domain.DomainFlowUtils.verifyPremiumNameIsNotBlocked;
 import static google.registry.flows.domain.DomainFlowUtils.verifyUnitIsYears;
 import static google.registry.model.domain.DomainResource.extendRegistrationWithCap;
-import static google.registry.model.domain.fee.Fee.FEE_TRANSFER_COMMAND_EXTENSIONS_IN_PREFERENCE_ORDER;
 import static google.registry.model.eppoutput.Result.Code.SUCCESS_WITH_ACTION_PENDING;
 import static google.registry.model.ofy.ObjectifyService.ofy;
 import static google.registry.pricing.PricingEngineProxy.getDomainRenewCost;
@@ -56,7 +55,7 @@ import google.registry.model.domain.DomainResource;
 import google.registry.model.domain.Period;
 import google.registry.model.domain.fee.BaseFee.FeeType;
 import google.registry.model.domain.fee.Fee;
-import google.registry.model.domain.fee.FeeTransformCommandExtension;
+import google.registry.model.domain.fee.FeeTransferCommandExtension;
 import google.registry.model.domain.fee.FeeTransformResponseExtension;
 import google.registry.model.domain.flags.FlagsTransferCommandExtension;
 import google.registry.model.domain.metadata.MetadataExtension;
@@ -130,8 +129,10 @@ public final class DomainTransferRequestFlow implements TransactionalFlow {
 
   @Override
   public final EppResponse run() throws EppException {
-    extensionManager.register(FlagsTransferCommandExtension.class, MetadataExtension.class);
-    extensionManager.registerAsGroup(FEE_TRANSFER_COMMAND_EXTENSIONS_IN_PREFERENCE_ORDER);
+    extensionManager.register(
+        FeeTransferCommandExtension.class,
+        FlagsTransferCommandExtension.class,
+        MetadataExtension.class);
     extensionManager.validate();
     validateClientIsLoggedIn(gainingClientId);
     Period period = ((Transfer) resourceCommand).getPeriod();
@@ -144,8 +145,8 @@ public final class DomainTransferRequestFlow implements TransactionalFlow {
     // The cost of the renewal implied by a transfer.
     Money renewCost = getDomainRenewCost(targetId, now, years);
     // An optional extension from the client specifying what they think the transfer should cost.
-    FeeTransformCommandExtension feeTransfer = eppInput.getFirstExtensionOfClasses(
-        FEE_TRANSFER_COMMAND_EXTENSIONS_IN_PREFERENCE_ORDER);
+    FeeTransferCommandExtension feeTransfer =
+        eppInput.getSingleExtension(FeeTransferCommandExtension.class);
     validateFeeChallenge(targetId, tld, now, feeTransfer, renewCost);
     HistoryEntry historyEntry = buildHistory(period, existingDomain, now);
     DateTime automaticTransferTime = now.plus(registry.getAutomaticTransferLength());
@@ -405,7 +406,7 @@ public final class DomainTransferRequestFlow implements TransactionalFlow {
   }
 
   private ImmutableList<FeeTransformResponseExtension> createResponseExtensions(Money renewCost,
-      FeeTransformCommandExtension feeTransfer) {
+      FeeTransferCommandExtension feeTransfer) {
     return feeTransfer == null
         ? null
         : ImmutableList.of(feeTransfer.createResponseBuilder()

java/google/registry/flows/domain/DomainUpdateFlow.java

@@ -37,7 +37,6 @@ import static google.registry.flows.domain.DomainFlowUtils.validateRegistrantAll
 import static google.registry.flows.domain.DomainFlowUtils.validateRequiredContactsPresent;
 import static google.registry.flows.domain.DomainFlowUtils.verifyClientUpdateNotProhibited;
 import static google.registry.flows.domain.DomainFlowUtils.verifyNotInPendingDelete;
-import static google.registry.model.domain.fee.Fee.FEE_UPDATE_COMMAND_EXTENSIONS_IN_PREFERENCE_ORDER;
 import static google.registry.model.ofy.ObjectifyService.ofy;
 import static google.registry.util.DateTimeUtils.earliestOf;
 
@@ -63,6 +62,7 @@ import google.registry.model.domain.DomainCommand.Update.Change;
 import google.registry.model.domain.DomainResource;
 import google.registry.model.domain.GracePeriod;
 import google.registry.model.domain.fee.FeeTransformCommandExtension;
+import google.registry.model.domain.fee.FeeUpdateCommandExtension;
 import google.registry.model.domain.flags.FlagsUpdateCommandExtension;
 import google.registry.model.domain.metadata.MetadataExtension;
 import google.registry.model.domain.rgp.GracePeriodStatus;
@@ -147,10 +147,10 @@ public final class DomainUpdateFlow implements TransactionalFlow {
   @Override
   public EppResponse run() throws EppException {
     extensionManager.register(
+        FeeUpdateCommandExtension.class,
         FlagsUpdateCommandExtension.class,
         MetadataExtension.class,
         SecDnsUpdateExtension.class);
-    extensionManager.registerAsGroup(FEE_UPDATE_COMMAND_EXTENSIONS_IN_PREFERENCE_ORDER);
     extensionManager.validate();
     validateClientIsLoggedIn(clientId);
     DateTime now = ofy().getTransactionTime();
@@ -201,7 +201,7 @@ public final class DomainUpdateFlow implements TransactionalFlow {
         Registry.get(tld), targetId, clientId, now, eppInput);
 
     FeeTransformCommandExtension feeUpdate =
-        eppInput.getFirstExtensionOfClasses(FEE_UPDATE_COMMAND_EXTENSIONS_IN_PREFERENCE_ORDER);
+        eppInput.getSingleExtension(FeeUpdateCommandExtension.class);
     // If the fee extension is present, validate it (even if the cost is zero, to check for price
     // mismatches). Don't rely on the the validateFeeChallenge check for feeUpdate nullness, because
     // it throws an error if the name is premium, and we don't want to do that here.