Add registrar id header to proxy requests (#2791)

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.
    *