Fix error handling in CopyDetailReportsAction (#2793)

* Fix error handling in CopyDetailReportsAction

The action tries to record errors per registrar in an ImmutableMap, without realizing that
there may be duplicate keys due to retries.

Switched to the `buildKeepingLast` method to build the map.

* Addressing comments and rebase

core/src/main/java/google/registry/reporting/billing/CopyDetailReportsAction.java

@@ -18,11 +18,12 @@ import static com.google.common.base.Throwables.getRootCause;
 import static google.registry.request.Action.Method.POST;
 import static jakarta.servlet.http.HttpServletResponse.SC_INTERNAL_SERVER_ERROR;
 import static jakarta.servlet.http.HttpServletResponse.SC_OK;
+import static java.util.stream.Collectors.joining;
 
 import com.google.cloud.storage.BlobId;
 import com.google.common.base.Splitter;
 import com.google.common.collect.ImmutableList;
-import com.google.common.collect.ImmutableMap;
+import com.google.common.collect.ImmutableMultimap;
 import com.google.common.collect.Iterables;
 import com.google.common.flogger.FluentLogger;
 import com.google.common.io.ByteStreams;
@@ -41,7 +42,6 @@ import jakarta.inject.Inject;
 import java.io.IOException;
 import java.io.InputStream;
 import java.util.Optional;
-import java.util.stream.Collectors;
 
 /** Copy all registrar detail reports in a given bucket's subdirectory from GCS to Drive. */
 @Action(
@@ -98,7 +98,8 @@ public final class CopyDetailReportsAction implements Runnable {
       response.setPayload(String.format("Failure, encountered %s", e.getMessage()));
       return;
     }
-    ImmutableMap.Builder<String, Throwable> copyErrorsBuilder = new ImmutableMap.Builder<>();
+    ImmutableMultimap.Builder<String, Throwable> copyErrorsBuilder =
+        new ImmutableMultimap.Builder<>();
     for (String detailReportName : detailReportObjectNames) {
       // The standard report format is "invoice_details_yyyy-MM_registrarId_tld.csv
       // TODO(larryruili): Determine a safer way of enforcing this.
@@ -145,17 +146,18 @@ public final class CopyDetailReportsAction implements Runnable {
     response.setStatus(SC_OK);
     response.setContentType(MediaType.PLAIN_TEXT_UTF_8);
     StringBuilder payload = new StringBuilder().append("Copied detail reports.\n");
-    ImmutableMap<String, Throwable> copyErrors = copyErrorsBuilder.build();
+    ImmutableMultimap<String, Throwable> copyErrors = copyErrorsBuilder.build();
     if (!copyErrors.isEmpty()) {
       payload.append("The following errors were encountered:\n");
-      payload.append(
-          copyErrors.entrySet().stream()
-              .map(
-                  entrySet ->
-                      String.format(
-                          "Registrar: %s\nError: %s\n",
-                          entrySet.getKey(), entrySet.getValue().getMessage()))
-              .collect(Collectors.joining()));
+      for (var registrarId : copyErrors.keySet()) {
+        payload.append(
+            String.format(
+                "Registrar: %s\nError: %s\n",
+                registrarId,
+                copyErrors.get(registrarId).stream()
+                    .map(Throwable::getMessage)
+                    .collect(joining("\n\t"))));
+      }
     }
     response.setPayload(payload.toString());
     emailUtils.sendAlertEmail(payload.toString());

core/src/test/java/google/registry/reporting/billing/CopyDetailReportsActionTest.java

@@ -178,21 +178,56 @@ class CopyDetailReportsActionTest {
     verify(emailUtils)
         .sendAlertEmail(
             """
-                Copied detail reports.
-                The following errors were encountered:
-                Registrar: TheRegistrar
-                Error: java.io.IOException: expected
-                """);
+            Copied detail reports.
+            The following errors were encountered:
+            Registrar: TheRegistrar
+            Error: java.io.IOException: expected
+            """);
     assertThat(response.getStatus()).isEqualTo(SC_OK);
     assertThat(response.getContentType()).isEqualTo(MediaType.PLAIN_TEXT_UTF_8);
     assertThat(response.getPayload())
         .isEqualTo(
             """
-                Copied detail reports.
-                The following errors were encountered:
-                Registrar: TheRegistrar
-                Error: java.io.IOException: expected
-                """);
+            Copied detail reports.
+            The following errors were encountered:
+            Registrar: TheRegistrar
+            Error: java.io.IOException: expected
+            """);
+  }
+
+  @Test
+  void testFail_tooManyFailures_one_registrar_sendsAlertEmail_continues() throws IOException {
+    gcsUtils.createFromBytes(
+        BlobId.of("test-bucket", "results/invoice_details_2017-10_TheRegistrar_hello.csv"),
+        "hola,mundo\n3,4".getBytes(UTF_8));
+
+    gcsUtils.createFromBytes(
+        BlobId.of("test-bucket", "results/invoice_details_2017-10_TheRegistrar_test.csv"),
+        "hello,world\n1,2".getBytes(UTF_8));
+    when(driveConnection.createOrUpdateFile(any(), any(), any(), any()))
+        .thenThrow(new IOException("expected"));
+
+    action.run();
+    verify(emailUtils)
+        .sendAlertEmail(
+            """
+            Copied detail reports.
+            The following errors were encountered:
+            Registrar: TheRegistrar
+            Error: java.io.IOException: expected
+            \tjava.io.IOException: expected
+            """);
+    assertThat(response.getStatus()).isEqualTo(SC_OK);
+    assertThat(response.getContentType()).isEqualTo(MediaType.PLAIN_TEXT_UTF_8);
+    assertThat(response.getPayload())
+        .isEqualTo(
+            """
+            Copied detail reports.
+            The following errors were encountered:
+            Registrar: TheRegistrar
+            Error: java.io.IOException: expected
+            \tjava.io.IOException: expected
+            """);
   }
 
   @Test

proxy/src/main/java/google/registry/proxy/handler/EppServiceHandler.java

@@ -19,8 +19,11 @@ import static com.google.common.base.Preconditions.checkNotNull;
 import static google.registry.networking.handler.SslServerInitializer.CLIENT_CERTIFICATE_PROMISE_KEY;
 import static google.registry.proxy.handler.ProxyProtocolHandler.REMOTE_ADDRESS_KEY;
 import static google.registry.util.X509Utils.getCertificateHash;
+import static java.nio.charset.StandardCharsets.US_ASCII;
 
+import com.google.common.base.Strings;
 import com.google.common.flogger.FluentLogger;
+import com.google.common.io.BaseEncoding;
 import google.registry.proxy.metric.FrontendMetrics;
 import google.registry.util.ProxyHttpHeaders;
 import io.netty.buffer.ByteBuf;
@@ -36,7 +39,11 @@ import io.netty.handler.ssl.SslHandshakeCompletionEvent;
 import io.netty.util.AttributeKey;
 import io.netty.util.concurrent.Promise;
 import java.security.cert.X509Certificate;
+import java.util.Optional;
 import java.util.function.Supplier;
+import java.util.regex.Matcher;
+import java.util.regex.Pattern;
+import javax.annotation.Nullable;
 
 /** Handler that processes EPP protocol logic. */
 public class EppServiceHandler extends HttpsRelayServiceHandler {
@@ -57,6 +64,8 @@ public class EppServiceHandler extends HttpsRelayServiceHandler {
   private String sslClientCertificateHash;
   private String clientAddress;
 
+  private Optional<String> maybeRegistrarId = Optional.empty();
+
   public EppServiceHandler(
       String relayHost,
       String relayPath,
@@ -128,6 +137,9 @@ public class EppServiceHandler extends HttpsRelayServiceHandler {
         .set(ProxyHttpHeaders.FALLBACK_IP_ADDRESS, clientAddress)
         .set(HttpHeaderNames.CONTENT_TYPE, EPP_CONTENT_TYPE)
         .set(HttpHeaderNames.ACCEPT, EPP_CONTENT_TYPE);
+
+    maybeSetRegistrarIdHeader(request);
+
     return request;
   }
 
@@ -142,4 +154,54 @@ public class EppServiceHandler extends HttpsRelayServiceHandler {
     }
     super.write(ctx, msg, promise);
   }
+
+  /**
+   * Sets and caches the Registrar-ID header on the request if the ID can be found.
+   *
+   * <p>This method first checks if the registrar ID has already been determined. If not, it
+   * inspects the cookies for a "SESSION_INFO" cookie, from which it attempts to extract the
+   * registrar ID.
+   *
+   * @param request The {@link FullHttpRequest} on which to potentially set the registrar ID header.
+   * @see #extractRegistrarIdFromSessionInfo(String)
+   */
+  private void maybeSetRegistrarIdHeader(FullHttpRequest request) {
+    if (maybeRegistrarId.isEmpty()) {
+      maybeRegistrarId =
+          cookieStore.entrySet().stream()
+              .map(e -> e.getValue())
+              .filter(cookie -> "SESSION_INFO".equals(cookie.name()))
+              .findFirst()
+              .flatMap(cookie -> extractRegistrarIdFromSessionInfo(cookie.value()));
+    }
+
+    if (maybeRegistrarId.isPresent() && !Strings.isNullOrEmpty(maybeRegistrarId.get())) {
+      request.headers().set(ProxyHttpHeaders.REGISTRAR_ID, maybeRegistrarId.get());
+    }
+  }
+
+  /** Extracts the registrar ID from a Base64-encoded session info string. */
+  private Optional<String> extractRegistrarIdFromSessionInfo(@Nullable String sessionInfo) {
+    if (sessionInfo == null) {
+      return Optional.empty();
+    }
+
+    try {
+      String decodedString = new String(BaseEncoding.base64Url().decode(sessionInfo), US_ASCII);
+      Pattern pattern = Pattern.compile("clientId=([^,\\s]+)?");
+      Matcher matcher = pattern.matcher(decodedString);
+
+      if (matcher.find()) {
+        String maybeRegistrarIdMatch = matcher.group(1);
+        if (!maybeRegistrarIdMatch.equals("null")) {
+          return Optional.of(maybeRegistrarIdMatch);
+        }
+      }
+
+    } catch (Throwable e) {
+      logger.atSevere().withCause(e).log("Failed to decode session info from Base64");
+    }
+
+    return Optional.empty();
+  }
 }

proxy/src/main/java/google/registry/proxy/handler/HttpsRelayServiceHandler.java

@@ -70,7 +70,7 @@ public abstract class HttpsRelayServiceHandler extends ByteToMessageCodec<FullHt
   protected static final ImmutableSet<Class<? extends Exception>> NON_FATAL_OUTBOUND_EXCEPTIONS =
       ImmutableSet.of(NonOkHttpResponseException.class);
 
-  private final Map<String, Cookie> cookieStore = new LinkedHashMap<>();
+  protected final Map<String, Cookie> cookieStore = new LinkedHashMap<>();
   private final String relayHost;
   private final String relayPath;
   private final boolean canary;

proxy/src/test/java/google/registry/proxy/handler/EppServiceHandlerTest.java

@@ -20,6 +20,7 @@ import static google.registry.proxy.TestUtils.assertHttpRequestEquivalent;
 import static google.registry.proxy.TestUtils.makeEppHttpResponse;
 import static google.registry.proxy.handler.ProxyProtocolHandler.REMOTE_ADDRESS_KEY;
 import static google.registry.util.X509Utils.getCertificateHash;
+import static java.nio.charset.StandardCharsets.US_ASCII;
 import static java.nio.charset.StandardCharsets.UTF_8;
 import static org.junit.jupiter.api.Assertions.assertThrows;
 import static org.mockito.Mockito.mock;
@@ -27,6 +28,7 @@ import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.verifyNoMoreInteractions;
 
 import com.google.common.base.Throwables;
+import com.google.common.io.BaseEncoding;
 import google.registry.proxy.TestUtils;
 import google.registry.proxy.handler.HttpsRelayServiceHandler.NonOkHttpResponseException;
 import google.registry.proxy.metric.FrontendMetrics;
@@ -357,4 +359,82 @@ class EppServiceHandlerTest {
     assertThat((Object) channel.readOutbound()).isNull();
     assertThat(channel.isActive()).isTrue();
   }
+
+  @Test
+  void testSuccess_registrarIdHeader_isSetFromSessionInfoCookie() throws Exception {
+    setHandshakeSuccess();
+    channel.readInbound(); // Read and discard the initial hello request.
+
+    // Simulate a server response that sets the SESSION_INFO cookie.
+    String registrarId = "TheRegistrar";
+    String sessionInfoValue =
+        BaseEncoding.base64Url()
+            .encode(("alpha,clientId=" + registrarId + ",beta").getBytes(US_ASCII));
+    Cookie sessionCookie = new DefaultCookie("SESSION_INFO", sessionInfoValue);
+    channel.writeOutbound(
+        makeEppHttpResponse("<epp>greeting</epp>", HttpResponseStatus.OK, sessionCookie));
+    channel.readOutbound(); // Read and discard the response sent to the client.
+
+    // Simulate a subsequent client request and check for the registrar ID header.
+    String clientRequestContent = "<epp>login</epp>";
+    channel.writeInbound(Unpooled.wrappedBuffer(clientRequestContent.getBytes(UTF_8)));
+
+    FullHttpRequest relayedRequest = channel.readInbound();
+    FullHttpRequest expectedRequest = makeEppHttpRequest(clientRequestContent, sessionCookie);
+    expectedRequest.headers().set(ProxyHttpHeaders.REGISTRAR_ID, registrarId);
+
+    assertHttpRequestEquivalent(relayedRequest, expectedRequest);
+    assertThat((Object) channel.readInbound()).isNull();
+    assertThat(channel.isActive()).isTrue();
+  }
+
+  @Test
+  void testSuccess_registrarIdHeader_isNotSetWhenSessionInfoCookieIsMissing() throws Exception {
+    setHandshakeSuccess();
+    channel.readInbound(); // Read and discard the initial hello request.
+
+    // Simulate a server response that does NOT set the SESSION_INFO cookie.
+    Cookie otherCookie = new DefaultCookie("some_other_cookie", "some_value");
+    channel.writeOutbound(
+        makeEppHttpResponse("<epp>greeting</epp>", HttpResponseStatus.OK, otherCookie));
+    channel.readOutbound(); // Read and discard the response sent to the client.
+
+    // Simulate a subsequent client request and verify the header is absent.
+    String clientRequestContent = "<epp>login</epp>";
+    channel.writeInbound(Unpooled.wrappedBuffer(clientRequestContent.getBytes(UTF_8)));
+
+    FullHttpRequest relayedRequest = channel.readInbound();
+    FullHttpRequest expectedRequest = makeEppHttpRequest(clientRequestContent, otherCookie);
+
+    assertHttpRequestEquivalent(relayedRequest, expectedRequest);
+    assertThat(relayedRequest.headers().contains(ProxyHttpHeaders.REGISTRAR_ID)).isFalse();
+    assertThat((Object) channel.readInbound()).isNull();
+    assertThat(channel.isActive()).isTrue();
+  }
+
+  @Test
+  void testSuccess_registrarIdHeader_isNotSetWhenClientIdIsNull() throws Exception {
+    setHandshakeSuccess();
+    channel.readInbound(); // Read and discard the initial hello request.
+
+    // Simulate a server response with a SESSION_INFO cookie where clientId is "null".
+    String sessionInfoValue =
+        BaseEncoding.base64Url().encode("alpha,clientId=null,beta".getBytes(US_ASCII));
+    Cookie sessionCookie = new DefaultCookie("SESSION_INFO", sessionInfoValue);
+    channel.writeOutbound(
+        makeEppHttpResponse("<epp>greeting</epp>", HttpResponseStatus.OK, sessionCookie));
+    channel.readOutbound(); // Read and discard the response sent to the client.
+
+    // Simulate a subsequent client request and verify the header is absent.
+    String clientRequestContent = "<epp>login</epp>";
+    channel.writeInbound(Unpooled.wrappedBuffer(clientRequestContent.getBytes(UTF_8)));
+
+    FullHttpRequest relayedRequest = channel.readInbound();
+    FullHttpRequest expectedRequest = makeEppHttpRequest(clientRequestContent, sessionCookie);
+
+    assertHttpRequestEquivalent(relayedRequest, expectedRequest);
+    assertThat(relayedRequest.headers().contains(ProxyHttpHeaders.REGISTRAR_ID)).isFalse();
+    assertThat((Object) channel.readInbound()).isNull();
+    assertThat(channel.isActive()).isTrue();
+  }
 }

util/src/main/java/google/registry/util/ProxyHttpHeaders.java

@@ -30,6 +30,9 @@ public final class ProxyHttpHeaders {
   /** HTTP header name used to pass the client IP address from the proxy to Nomulus. */
   public static final String IP_ADDRESS = "Nomulus-Client-Address";
 
+  /** HTTP header name used to pass the Registrar Id from the proxy to Nomulus. */
+  public static final String REGISTRAR_ID = "Nomulus-Registrar-Id";
+
   /**
    * Fallback HTTP header name used to pass the client IP address from the proxy to Nomulus.
    *