diff --git a/core/src/main/java/google/registry/config/RegistryConfig.java b/core/src/main/java/google/registry/config/RegistryConfig.java index 793b341a4..9ce647915 100644 --- a/core/src/main/java/google/registry/config/RegistryConfig.java +++ b/core/src/main/java/google/registry/config/RegistryConfig.java @@ -1587,6 +1587,12 @@ public final class RegistryConfig { return config.eppServer.readTimeoutSeconds; } + @Provides + @Config("eppServerPreLoginReadTimeoutSeconds") + public static int provideEppServerPreLoginReadTimeoutSeconds(RegistryConfigSettings config) { + return config.eppServer.preLoginReadTimeoutSeconds; + } + @Provides @Config("eppServerMaxConnectionsPerIp") public static int provideEppServerMaxConnectionsPerIp(RegistryConfigSettings config) { @@ -1594,9 +1600,9 @@ public final class RegistryConfig { } @Provides - @Config("eppServerMaxConnectionsPerCert") - public static int provideEppServerMaxConnectionsPerCert(RegistryConfigSettings config) { - return config.eppServer.maxConnectionsPerCert; + @Config("eppServerMaxConnectionsPerRegistrar") + public static int provideEppServerMaxConnectionsPerRegistrar(RegistryConfigSettings config) { + return config.eppServer.maxConnectionsPerRegistrar; } @Provides diff --git a/core/src/main/java/google/registry/config/RegistryConfigSettings.java b/core/src/main/java/google/registry/config/RegistryConfigSettings.java index d6bda8f0e..8c7ba2475 100644 --- a/core/src/main/java/google/registry/config/RegistryConfigSettings.java +++ b/core/src/main/java/google/registry/config/RegistryConfigSettings.java @@ -217,8 +217,9 @@ public class RegistryConfigSettings { public int maxMessageLengthBytes; public int headerLengthBytes; public int readTimeoutSeconds; + public int preLoginReadTimeoutSeconds; public int maxConnectionsPerIp; - public int maxConnectionsPerCert; + public int maxConnectionsPerRegistrar; public int serverCertificateCacheSeconds; public Quota quota; } diff --git a/core/src/main/java/google/registry/config/files/default-config.yaml b/core/src/main/java/google/registry/config/files/default-config.yaml index 322d6ca87..370a68eee 100644 --- a/core/src/main/java/google/registry/config/files/default-config.yaml +++ b/core/src/main/java/google/registry/config/files/default-config.yaml @@ -457,20 +457,22 @@ eppServer: headerLengthBytes: 4 # Time after which an idle connection will be closed. readTimeoutSeconds: 3600 + # Time after which an idle connection will be closed before login. + preLoginReadTimeoutSeconds: 10 # Max concurrent connections per IP address. maxConnectionsPerIp: 10 - # Max concurrent connections per authenticated certificate. - maxConnectionsPerCert: 10 + # Max concurrent connections per authenticated registrar. + maxConnectionsPerRegistrar: 10 # Server certificate cache duration. serverCertificateCacheSeconds: 1800 # Quota configuration for EPP quota: - refreshSeconds: 0 + # Default quota applies individually to any IP or registrar NOT listed in customQuota defaultQuota: userId: [] tokenAmount: 100 - refillSeconds: 0 + refillSeconds: 10 # To implement a shared quota group across multiple registrars, place a virtual # group name as the FIRST element of the userId list. # e.g., userId: ["my_group", "registrar1", "registrar2"] diff --git a/core/src/main/java/google/registry/eppserver/handler/EppProxyProtocolHandler.java b/core/src/main/java/google/registry/eppserver/handler/EppProxyProtocolHandler.java index 60c88f8b0..b1b5f86d8 100644 --- a/core/src/main/java/google/registry/eppserver/handler/EppProxyProtocolHandler.java +++ b/core/src/main/java/google/registry/eppserver/handler/EppProxyProtocolHandler.java @@ -18,6 +18,7 @@ import static com.google.common.base.Preconditions.checkState; import static java.nio.charset.StandardCharsets.US_ASCII; import com.google.common.flogger.FluentLogger; +import com.google.common.net.InetAddresses; import io.netty.buffer.ByteBuf; import io.netty.channel.ChannelHandlerContext; import io.netty.handler.codec.ByteToMessageDecoder; @@ -78,9 +79,17 @@ public class EppProxyProtocolHandler extends ByteToMessageDecoder { logger.atFine().log("PROXY HEADER for channel %s: %s", ctx.channel(), proxyHeader); String[] headerArray = proxyHeader.split(" ", -1); if (headerArray.length == 6) { - remoteIP = headerArray[2]; - logger.atFine().log( - "Header parsed, using %s as remote IP for channel %s", remoteIP, ctx.channel()); + String parsedIP = headerArray[2]; + if (InetAddresses.isInetAddress(parsedIP)) { + remoteIP = parsedIP; + logger.atFine().log( + "Header parsed, using %s as remote IP for channel %s", remoteIP, ctx.channel()); + } else { + logger.atWarning().log( + "Invalid IP address in PROXY header: %s, falling back to source IP for channel %s", + parsedIP, ctx.channel()); + remoteIP = getSourceIP(ctx); + } // If the header is "PROXY UNKNOWN" // (see https://www.haproxy.org/download/1.8/doc/proxy-protocol.txt), likely when the // remote connection to the external load balancer is through special means, make it diff --git a/core/src/main/java/google/registry/eppserver/handler/EppServiceHandler.java b/core/src/main/java/google/registry/eppserver/handler/EppServiceHandler.java index 09960a74f..e9cfd3b1b 100644 --- a/core/src/main/java/google/registry/eppserver/handler/EppServiceHandler.java +++ b/core/src/main/java/google/registry/eppserver/handler/EppServiceHandler.java @@ -42,10 +42,12 @@ import io.netty.channel.SimpleChannelInboundHandler; import io.netty.util.AttributeKey; import io.netty.util.concurrent.Future; import io.netty.util.concurrent.Promise; +import io.netty.util.concurrent.ScheduledFuture; import jakarta.inject.Inject; import jakarta.inject.Named; import java.security.cert.X509Certificate; import java.util.UUID; +import java.util.concurrent.TimeUnit; import java.util.function.Supplier; import java.util.regex.Matcher; import java.util.regex.Pattern; @@ -73,14 +75,17 @@ public class EppServiceHandler extends SimpleChannelInboundHandler { private final QuotaManager commandQuotaManager; private final Supplier idTokenSupplier; private final String projectId; + private final int preLoginReadTimeoutSeconds; private String sslClientCertificateHash; private String clientAddress; private String registrarId; // The clID extracted from login + private String authenticatedRegistrarId; // The verified registrar ID after successful login private String sessionCookie; private boolean ipAcquired = false; - private boolean certAcquired = false; + private boolean registrarAcquired = false; + private ScheduledFuture preLoginTimeoutTask; @VisibleForTesting RequestHandler requestHandler = RegistryServlet.component.requestHandler(); @@ -91,13 +96,15 @@ public class EppServiceHandler extends SimpleChannelInboundHandler { LocalConnectionLimiter localConnectionLimiter, @CommandQuota QuotaManager commandQuotaManager, @Named("idToken") Supplier idTokenSupplier, - @Config("projectId") String projectId) { + @Config("projectId") String projectId, + @Config("eppServerPreLoginReadTimeoutSeconds") int preLoginReadTimeoutSeconds) { this.helloBytes = helloBytes.clone(); this.metrics = metrics; this.localConnectionLimiter = localConnectionLimiter; this.commandQuotaManager = commandQuotaManager; this.idTokenSupplier = idTokenSupplier; this.projectId = projectId; + this.preLoginReadTimeoutSeconds = preLoginReadTimeoutSeconds; } @Override @@ -110,8 +117,7 @@ public class EppServiceHandler extends SimpleChannelInboundHandler { ctx.executor().execute(() -> onSslHandshakeComplete(ctx, promise.getNow())); } else { logger.atWarning().withCause(promise.cause()).log("SSL handshake failed"); - @SuppressWarnings("unused") - Future unusedFuture = ctx.close(); + closeConnection(ctx); } }); } @@ -119,26 +125,36 @@ public class EppServiceHandler extends SimpleChannelInboundHandler { } private void onSslHandshakeComplete(ChannelHandlerContext ctx, X509Certificate cert) { + if (!ctx.channel().isActive()) { + return; + } sslClientCertificateHash = getCertificateHash(cert); clientAddress = ctx.channel().attr(REMOTE_ADDRESS_KEY).get(); ctx.channel().attr(CLIENT_CERTIFICATE_HASH_KEY).set(sslClientCertificateHash); - // 1. Connection throttling (IP and Certificate) + // 1. Connection throttling (IP only pre-login) if (!localConnectionLimiter.acquireIp(clientAddress)) { metrics.registerQuotaRejection("epp_connection_ip", clientAddress); - @SuppressWarnings("unused") - Future unusedFuture = ctx.close(); + closeConnection(ctx); return; } ipAcquired = true; - if (!localConnectionLimiter.acquireCert(sslClientCertificateHash)) { - metrics.registerQuotaRejection("epp_connection", sslClientCertificateHash); - @SuppressWarnings("unused") - Future unusedFuture = ctx.close(); - return; - } - certAcquired = true; + // Schedule login timeout + preLoginTimeoutTask = + ctx.executor() + .schedule( + () -> { + if (!registrarAcquired) { + logger.atWarning().log( + "EPP login timeout expired for channel %s, closing connection", + ctx.channel()); + metrics.registerQuotaRejection("epp_login_timeout", clientAddress); + closeConnection(ctx); + } + }, + preLoginReadTimeoutSeconds, + TimeUnit.SECONDS); metrics.registerActiveConnection("epp", sslClientCertificateHash, ctx.channel()); @@ -154,7 +170,32 @@ public class EppServiceHandler extends SimpleChannelInboundHandler { private void handleEppFrame(ChannelHandlerContext ctx, ByteBuf frame) { String xml = frame.toString(UTF_8); - // 1. Maturing Identity: If we don't have clID yet, try to extract it from a login command. + extractRegistrarId(xml); + + if (!acquireCommandQuota(ctx)) { + return; + } + + FakeHttpServletRequest req = buildServletRequest(xml); + FakeHttpServletResponse rsp = new FakeHttpServletResponse(); + String traceId = + String.format( + "projects/%s/traces/%s", projectId, UUID.randomUUID().toString().replace("-", "")); + setCurrentTraceId(traceId); + setCurrentRequest("POST", "/_dr/epp", "Netty-EPP", "EPP/1.0"); + try { + requestHandler.handleRequest(req, rsp); + processServletResponse(ctx, rsp); + } catch (Exception e) { + logger.atSevere().withCause(e).log("Internal EPP processing error"); + closeConnection(ctx); + } finally { + setCurrentTraceId(null); + unsetCurrentRequest(); + } + } + + private void extractRegistrarId(String xml) { if (registrarId == null) { Matcher matcher = CLID_PATTERN.matcher(xml); if (matcher.find()) { @@ -162,20 +203,22 @@ public class EppServiceHandler extends SimpleChannelInboundHandler { logger.atInfo().log("Identified registrar: %s", registrarId); } } + } - // 2. Command-level rate limiting - // Use clID if identified, otherwise fallback to cert hash (for the login command itself). - String throttleId = (registrarId != null) ? registrarId : sslClientCertificateHash; + private boolean acquireCommandQuota(ChannelHandlerContext ctx) { + String throttleId = + (authenticatedRegistrarId != null) ? authenticatedRegistrarId : sslClientCertificateHash; if (throttleId != null) { if (!commandQuotaManager.acquireQuota(new QuotaManager.QuotaRequest(throttleId)).success()) { metrics.registerQuotaRejection("epp_command", throttleId); - @SuppressWarnings("unused") - Future unusedFuture = ctx.close(); - return; + closeConnection(ctx); + return false; } } + return true; + } - // 3. Execute command in-process + private FakeHttpServletRequest buildServletRequest(String xml) { FakeHttpServletRequest req = new FakeHttpServletRequest(); req.setRequestUri("/_dr/epp"); req.setBody(xml.getBytes(UTF_8)); @@ -188,42 +231,60 @@ public class EppServiceHandler extends SimpleChannelInboundHandler { req.setHeader("Cookie", sessionCookie); } req.setHeader("Authorization", "Bearer " + idTokenSupplier.get()); + return req; + } - FakeHttpServletResponse rsp = new FakeHttpServletResponse(); - String traceId = - String.format( - "projects/%s/traces/%s", projectId, UUID.randomUUID().toString().replace("-", "")); - setCurrentTraceId(traceId); - setCurrentRequest("POST", "/_dr/epp", "Netty-EPP", "EPP/1.0"); - try { - requestHandler.handleRequest(req, rsp); - String setCookie = rsp.getHeader("Set-Cookie"); - if (setCookie != null) { - sessionCookie = setCookie; - } - - ByteBuf out = Unpooled.wrappedBuffer(rsp.getPayload()); - if ("close".equals(rsp.getHeader(ProxyHttpHeaders.EPP_SESSION))) { - @SuppressWarnings("unused") - Future unusedFuture = ctx.writeAndFlush(out).addListener(ChannelFutureListener.CLOSE); - } else { - @SuppressWarnings("unused") - Future unusedFuture = ctx.writeAndFlush(out); - } - } catch (Exception e) { - logger.atSevere().withCause(e).log("Internal EPP processing error"); - @SuppressWarnings("unused") - Future unusedFuture = ctx.close(); - } finally { - setCurrentTraceId(null); - unsetCurrentRequest(); + private void processServletResponse(ChannelHandlerContext ctx, FakeHttpServletResponse rsp) { + String setCookie = rsp.getHeader("Set-Cookie"); + if (setCookie != null) { + sessionCookie = setCookie; } + + String authRegistrarId = rsp.getHeader(ProxyHttpHeaders.LOGGED_IN_REGISTRAR); + if (authRegistrarId != null && !registrarAcquired) { + logger.atInfo().log("Registrar %s successfully authenticated", authRegistrarId); + if (!localConnectionLimiter.acquireRegistrar(authRegistrarId)) { + logger.atWarning().log( + "Registrar %s exceeded concurrent connection limit, closing connection", + authRegistrarId); + metrics.registerQuotaRejection("epp_connection_registrar", authRegistrarId); + closeConnection(ctx); + return; + } + registrarAcquired = true; + authenticatedRegistrarId = authRegistrarId; + registrarId = authRegistrarId; + + // Cancel pre-login timeout task + if (preLoginTimeoutTask != null) { + preLoginTimeoutTask.cancel(false); + preLoginTimeoutTask = null; + } + } + + ByteBuf out = Unpooled.wrappedBuffer(rsp.getPayload()); + if ("close".equals(rsp.getHeader(ProxyHttpHeaders.EPP_SESSION))) { + @SuppressWarnings("unused") + Future unusedFuture = ctx.writeAndFlush(out).addListener(ChannelFutureListener.CLOSE); + } else { + @SuppressWarnings("unused") + Future unusedFuture = ctx.writeAndFlush(out); + } + } + + private void closeConnection(ChannelHandlerContext ctx) { + @SuppressWarnings("unused") + Future unusedFuture = ctx.close(); } @Override public void channelInactive(ChannelHandlerContext ctx) throws Exception { - if (certAcquired) { - localConnectionLimiter.releaseCert(sslClientCertificateHash); + if (preLoginTimeoutTask != null) { + preLoginTimeoutTask.cancel(false); + preLoginTimeoutTask = null; + } + if (registrarAcquired) { + localConnectionLimiter.releaseRegistrar(authenticatedRegistrarId); } if (ipAcquired) { localConnectionLimiter.releaseIp(clientAddress); diff --git a/core/src/main/java/google/registry/eppserver/quota/LocalConnectionLimiter.java b/core/src/main/java/google/registry/eppserver/quota/LocalConnectionLimiter.java index 08f87694a..a0cd51911 100644 --- a/core/src/main/java/google/registry/eppserver/quota/LocalConnectionLimiter.java +++ b/core/src/main/java/google/registry/eppserver/quota/LocalConnectionLimiter.java @@ -22,24 +22,24 @@ import javax.annotation.concurrent.ThreadSafe; /** * Thread-safe, in-memory rate limiter for restricting the number of concurrent connections allowed - * per IP address and per authenticated certificate. + * per IP address and per authenticated registrar. */ @ThreadSafe @Singleton public class LocalConnectionLimiter { private final int maxConnectionsPerIp; - private final int maxConnectionsPerCert; + private final int maxConnectionsPerRegistrar; private final ConcurrentHashMap ipConnections = new ConcurrentHashMap<>(); - private final ConcurrentHashMap certConnections = new ConcurrentHashMap<>(); + private final ConcurrentHashMap registrarConnections = new ConcurrentHashMap<>(); @Inject public LocalConnectionLimiter( @Config("eppServerMaxConnectionsPerIp") int maxConnectionsPerIp, - @Config("eppServerMaxConnectionsPerCert") int maxConnectionsPerCert) { + @Config("eppServerMaxConnectionsPerRegistrar") int maxConnectionsPerRegistrar) { this.maxConnectionsPerIp = maxConnectionsPerIp; - this.maxConnectionsPerCert = maxConnectionsPerCert; + this.maxConnectionsPerRegistrar = maxConnectionsPerRegistrar; } /** Attempts to acquire a slot for the given IP address. */ @@ -52,14 +52,14 @@ public class LocalConnectionLimiter { release(ipAddress, ipConnections); } - /** Attempts to acquire a slot for the given certificate hash. */ - public boolean acquireCert(String certHash) { - return acquire(certHash, certConnections, maxConnectionsPerCert); + /** Attempts to acquire a slot for the given registrar ID. */ + public boolean acquireRegistrar(String registrarId) { + return acquire(registrarId, registrarConnections, maxConnectionsPerRegistrar); } - /** Releases a slot for the given certificate hash. */ - public void releaseCert(String certHash) { - release(certHash, certConnections); + /** Releases a slot for the given registrar ID. */ + public void releaseRegistrar(String registrarId) { + release(registrarId, registrarConnections); } private boolean acquire(String key, ConcurrentHashMap map, int limit) { diff --git a/core/src/main/java/google/registry/flows/EppRequestHandler.java b/core/src/main/java/google/registry/flows/EppRequestHandler.java index e38e82ed7..8e4b166c2 100644 --- a/core/src/main/java/google/registry/flows/EppRequestHandler.java +++ b/core/src/main/java/google/registry/flows/EppRequestHandler.java @@ -75,6 +75,15 @@ public class EppRequestHandler { // closed by the proxy. Whether the EPP proxy actually terminates the connection with the // client is up to its implementation. // See: https://tools.ietf.org/html/rfc5734#section-2 + String authRegistrarId = null; + try { + authRegistrarId = sessionMetadata.getRegistrarId(); + } catch (IllegalStateException e) { + // Session was invalidated (e.g. during logout) + } + if (authRegistrarId != null) { + response.setHeader(ProxyHttpHeaders.LOGGED_IN_REGISTRAR, authRegistrarId); + } if (eppOutput.isResponse() && eppOutput.getResponse().getResult().getCode() == SUCCESS_AND_CLOSE) { response.setHeader(ProxyHttpHeaders.EPP_SESSION, "close"); diff --git a/core/src/test/java/google/registry/eppserver/handler/EppProxyProtocolHandlerTest.java b/core/src/test/java/google/registry/eppserver/handler/EppProxyProtocolHandlerTest.java index d974e5711..ef6e64be2 100644 --- a/core/src/test/java/google/registry/eppserver/handler/EppProxyProtocolHandlerTest.java +++ b/core/src/test/java/google/registry/eppserver/handler/EppProxyProtocolHandlerTest.java @@ -19,11 +19,23 @@ import static com.google.common.truth.Truth.assertThat; import io.netty.buffer.ByteBuf; import io.netty.buffer.Unpooled; import io.netty.channel.embedded.EmbeddedChannel; +import java.net.InetAddress; +import java.net.InetSocketAddress; +import java.net.SocketAddress; import java.nio.charset.StandardCharsets; import org.junit.jupiter.api.Test; class EppProxyProtocolHandlerTest { + private EmbeddedChannel createChannel(EppProxyProtocolHandler handler) { + return new EmbeddedChannel(handler) { + @Override + public SocketAddress remoteAddress() { + return new InetSocketAddress(InetAddress.getLoopbackAddress(), 12345); + } + }; + } + @Test void testProxyProtocol_parsesValidHeader() { EppProxyProtocolHandler handler = new EppProxyProtocolHandler(); @@ -39,6 +51,21 @@ class EppProxyProtocolHandlerTest { assertThat(channel.pipeline().get(EppProxyProtocolHandler.class)).isNull(); } + @Test + void testProxyProtocol_invalidIP_fallsBackToSource() { + EppProxyProtocolHandler handler = new EppProxyProtocolHandler(); + EmbeddedChannel channel = createChannel(handler); + + String proxyHeader = "PROXY TCP4 invalid_ip_address 10.0.0.1 50000 443\r\n"; + ByteBuf buffer = Unpooled.wrappedBuffer(proxyHeader.getBytes(StandardCharsets.US_ASCII)); + + channel.writeInbound(buffer); + + String remoteAddress = channel.attr(EppProxyProtocolHandler.REMOTE_ADDRESS_KEY).get(); + assertThat(remoteAddress).isEqualTo("127.0.0.1"); // Falls back to mocked remoteAddress + assertThat(channel.pipeline().get(EppProxyProtocolHandler.class)).isNull(); + } + @Test void testProxyProtocol_unknownHeader() { EppProxyProtocolHandler handler = new EppProxyProtocolHandler(); @@ -57,7 +84,7 @@ class EppProxyProtocolHandlerTest { @Test void testProxyProtocol_noHeader_notProxied() { EppProxyProtocolHandler handler = new EppProxyProtocolHandler(); - EmbeddedChannel channel = new EmbeddedChannel(handler); + EmbeddedChannel channel = createChannel(handler); String normalData = "NOT_A_PROXY_HEADER"; ByteBuf buffer = Unpooled.wrappedBuffer(normalData.getBytes(StandardCharsets.US_ASCII)); @@ -65,8 +92,7 @@ class EppProxyProtocolHandlerTest { channel.writeInbound(buffer); String remoteAddress = channel.attr(EppProxyProtocolHandler.REMOTE_ADDRESS_KEY).get(); - // In EmbeddedChannel without remoteAddress mock, getSourceIP returns null - assertThat(remoteAddress).isNull(); + assertThat(remoteAddress).isEqualTo("127.0.0.1"); // Falls back to mocked remoteAddress assertThat(channel.pipeline().get(EppProxyProtocolHandler.class)).isNull(); ByteBuf passedOn = channel.readInbound(); diff --git a/core/src/test/java/google/registry/eppserver/handler/EppServiceHandlerTest.java b/core/src/test/java/google/registry/eppserver/handler/EppServiceHandlerTest.java index 21c8893bf..97e401dc0 100644 --- a/core/src/test/java/google/registry/eppserver/handler/EppServiceHandlerTest.java +++ b/core/src/test/java/google/registry/eppserver/handler/EppServiceHandlerTest.java @@ -17,11 +17,15 @@ package google.registry.eppserver.handler; import static google.registry.eppserver.handler.EppProxyProtocolHandler.REMOTE_ADDRESS_KEY; import static google.registry.networking.handler.SslServerInitializer.CLIENT_CERTIFICATE_PROMISE_KEY; import static java.nio.charset.StandardCharsets.UTF_8; +import static org.junit.jupiter.api.Assertions.assertNotNull; import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyLong; import static org.mockito.ArgumentMatchers.argThat; import static org.mockito.ArgumentMatchers.eq; import static org.mockito.Mockito.doAnswer; +import static org.mockito.Mockito.lenient; import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.times; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; @@ -34,6 +38,7 @@ import google.registry.request.RequestHandler; import google.registry.util.FakeHttpServletRequest; import google.registry.util.FakeHttpServletResponse; import google.registry.util.ProxyHttpHeaders; +import google.registry.util.X509Utils; import io.netty.buffer.ByteBuf; import io.netty.buffer.Unpooled; import io.netty.channel.Channel; @@ -42,13 +47,17 @@ import io.netty.channel.ChannelFutureListener; import io.netty.channel.ChannelHandlerContext; import io.netty.util.Attribute; import io.netty.util.concurrent.DefaultPromise; +import io.netty.util.concurrent.EventExecutor; import io.netty.util.concurrent.ImmediateEventExecutor; import io.netty.util.concurrent.Promise; +import io.netty.util.concurrent.ScheduledFuture; import java.security.cert.X509Certificate; +import java.util.concurrent.TimeUnit; import java.util.function.Supplier; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.extension.ExtendWith; +import org.mockito.ArgumentCaptor; import org.mockito.Mock; import org.mockito.junit.jupiter.MockitoExtension; @@ -61,6 +70,8 @@ class EppServiceHandlerTest { @Mock private Supplier idTokenSupplier; @Mock private ChannelHandlerContext ctx; @Mock private Channel channel; + @Mock private EventExecutor executor; + @Mock private ScheduledFuture scheduledFuture; @Mock private RequestHandler requestHandler; @Mock private Attribute> certPromiseAttr; @@ -80,12 +91,32 @@ class EppServiceHandlerTest { localConnectionLimiter, commandQuotaManager, idTokenSupplier, - "test-project"); + "test-project", + 10); // preLoginReadTimeoutSeconds handler.requestHandler = requestHandler; when(ctx.channel()).thenReturn(channel); - when(ctx.executor()).thenReturn(ImmediateEventExecutor.INSTANCE); + when(ctx.executor()).thenReturn(executor); + lenient().when(channel.isActive()).thenReturn(true); + + doAnswer( + invocation -> { + Runnable runnable = invocation.getArgument(0); + runnable.run(); + return null; + }) + .when(executor) + .execute(any(Runnable.class)); + + lenient() + .doReturn(scheduledFuture) + .when(executor) + .schedule(any(Runnable.class), anyLong(), any(TimeUnit.class)); + + lenient() + .when(commandQuotaManager.acquireQuota(any(QuotaRequest.class))) + .thenReturn(new QuotaResponse(true)); } private void setUpSuccessfulHandshake() throws Exception { @@ -101,7 +132,6 @@ class EppServiceHandlerTest { when(certificate.getEncoded()).thenReturn(new byte[] {1, 2, 3}); when(localConnectionLimiter.acquireIp(any(String.class))).thenReturn(true); - when(localConnectionLimiter.acquireCert(any(String.class))).thenReturn(true); certPromise.setSuccess(certificate); } @@ -156,24 +186,31 @@ class EppServiceHandlerTest { } @Test - void testChannelActive_certQuotaRejected() throws Exception { - certPromise = new DefaultPromise<>(ImmediateEventExecutor.INSTANCE); - when(channel.attr(CLIENT_CERTIFICATE_PROMISE_KEY)).thenReturn(certPromiseAttr); - when(certPromiseAttr.get()).thenReturn(certPromise); + void testChannelRead0_registrarQuotaRejected() throws Exception { + setUpSuccessfulHandshake(); - handler.channelActive(ctx); + when(idTokenSupplier.get()).thenReturn("fake_id_token"); + when(commandQuotaManager.acquireQuota(any(QuotaRequest.class))) + .thenReturn(new QuotaResponse(true)); - when(channel.attr(REMOTE_ADDRESS_KEY)).thenReturn(remoteAddressAttr); - when(remoteAddressAttr.get()).thenReturn("192.168.1.1"); - when(channel.attr(EppServiceHandler.CLIENT_CERTIFICATE_HASH_KEY)).thenReturn(certHashAttr); - when(certificate.getEncoded()).thenReturn(new byte[] {1, 2, 3}); + String eppLoginXml = "RegistrarA"; + ByteBuf inFrame = Unpooled.wrappedBuffer(eppLoginXml.getBytes(UTF_8)); - when(localConnectionLimiter.acquireIp(any(String.class))).thenReturn(true); - when(localConnectionLimiter.acquireCert(any(String.class))).thenReturn(false); + doAnswer( + invocation -> { + FakeHttpServletResponse rsp = invocation.getArgument(1); + rsp.setHeader(ProxyHttpHeaders.LOGGED_IN_REGISTRAR, "RegistrarA"); + rsp.getWriter().write("success"); + return null; + }) + .when(requestHandler) + .handleRequest(any(FakeHttpServletRequest.class), any(FakeHttpServletResponse.class)); - certPromise.setSuccess(certificate); + when(localConnectionLimiter.acquireRegistrar("RegistrarA")).thenReturn(false); - verify(metrics).registerQuotaRejection(eq("epp_connection"), any(String.class)); + handler.channelRead0(ctx, inFrame); + + verify(metrics).registerQuotaRejection(eq("epp_connection_registrar"), eq("RegistrarA")); verify(ctx).close(); } @@ -194,17 +231,24 @@ class EppServiceHandlerTest { FakeHttpServletRequest req = invocation.getArgument(0); FakeHttpServletResponse rsp = invocation.getArgument(1); - rsp.setHeader("Set-Cookie", "SESSION_INFO=xyz123"); + rsp.setHeader("Set-Cookie", "SESSION_INFO=Y2xpZW50SWQ9UmVnaXN0cmFyQQ=="); + rsp.setHeader(ProxyHttpHeaders.LOGGED_IN_REGISTRAR, "RegistrarA"); rsp.getWriter().write("success"); return null; }) .when(requestHandler) .handleRequest(any(FakeHttpServletRequest.class), any(FakeHttpServletResponse.class)); + // Mock successful registrar connection acquisition + when(localConnectionLimiter.acquireRegistrar("RegistrarA")).thenReturn(true); + handler.channelRead0(ctx, inFrame); - // Verify command quota was requested for the extracted clID "RegistrarA" - verify(commandQuotaManager).acquireQuota(eq(new QuotaRequest("RegistrarA"))); + // Verify command quota was requested for the cert hash pre-login + String certHash = X509Utils.getCertificateHash(certificate); + verify(commandQuotaManager, times(2)).acquireQuota(eq(new QuotaRequest(certHash))); + verify(localConnectionLimiter).acquireRegistrar("RegistrarA"); + verify(scheduledFuture).cancel(eq(false)); // Verify the response from the servlet was written back to the channel verify(ctx) @@ -223,7 +267,7 @@ class EppServiceHandlerTest { invocation -> { FakeHttpServletRequest req = invocation.getArgument(0); // Verify the cookie was properly propagated - if (!"SESSION_INFO=xyz123".equals(req.getHeader("Cookie"))) { + if (!"SESSION_INFO=Y2xpZW50SWQ9UmVnaXN0cmFyQQ==".equals(req.getHeader("Cookie"))) { throw new AssertionError("Missing or incorrect cookie"); } // Verify the registrar ID was properly propagated @@ -236,6 +280,9 @@ class EppServiceHandlerTest { .handleRequest(any(FakeHttpServletRequest.class), any(FakeHttpServletResponse.class)); handler.channelRead0(ctx, inFrame2); + + // Verify command quota was requested for the authenticated registrar post-login + verify(commandQuotaManager).acquireQuota(eq(new QuotaRequest("RegistrarA"))); } @Test @@ -285,13 +332,64 @@ class EppServiceHandlerTest { } @Test - void testChannelInactive_releasesQuotas() throws Exception { + void testChannelInactive_releasesIp() throws Exception { setUpSuccessfulHandshake(); handler.channelInactive(ctx); - // Verify the in-memory limiter releases both IP and Cert + // Verify the in-memory limiter releases IP verify(localConnectionLimiter).releaseIp(eq("192.168.1.1")); - verify(localConnectionLimiter).releaseCert(any(String.class)); + } + + @Test + void testChannelInactive_postLogin_releasesIpAndRegistrar() throws Exception { + setUpSuccessfulHandshake(); + + when(idTokenSupplier.get()).thenReturn("fake_id_token"); + when(commandQuotaManager.acquireQuota(any(QuotaRequest.class))) + .thenReturn(new QuotaResponse(true)); + when(localConnectionLimiter.acquireRegistrar("RegistrarA")).thenReturn(true); + + String eppLoginXml = "RegistrarA"; + ByteBuf inFrame = Unpooled.wrappedBuffer(eppLoginXml.getBytes(UTF_8)); + + doAnswer( + invocation -> { + FakeHttpServletResponse rsp = invocation.getArgument(1); + rsp.setHeader("Set-Cookie", "SESSION_INFO=Y2xpZW50SWQ9UmVnaXN0cmFyQQ=="); + rsp.setHeader(ProxyHttpHeaders.LOGGED_IN_REGISTRAR, "RegistrarA"); + rsp.getWriter().write("success"); + return null; + }) + .when(requestHandler) + .handleRequest(any(FakeHttpServletRequest.class), any(FakeHttpServletResponse.class)); + + handler.channelRead0(ctx, inFrame); + + handler.channelInactive(ctx); + + // Verify the in-memory limiter releases both IP and Registrar + verify(localConnectionLimiter).releaseIp(eq("192.168.1.1")); + verify(localConnectionLimiter).releaseRegistrar(eq("RegistrarA")); + } + + @Test + void testChannelActive_loginTimeoutTriggered() throws Exception { + ArgumentCaptor timeoutTaskCaptor = ArgumentCaptor.forClass(Runnable.class); + when(executor.schedule(timeoutTaskCaptor.capture(), eq(10L), eq(TimeUnit.SECONDS))) + .thenReturn(null); + + setUpSuccessfulHandshake(); + + Runnable timeoutTask = timeoutTaskCaptor.getValue(); + assertNotNull(timeoutTask); + + ChannelFuture closeFuture = mock(ChannelFuture.class); + when(ctx.close()).thenReturn(closeFuture); + + timeoutTask.run(); + + verify(metrics).registerQuotaRejection("epp_login_timeout", "192.168.1.1"); + verify(ctx).close(); } } diff --git a/core/src/test/java/google/registry/eppserver/quota/LocalConnectionLimiterTest.java b/core/src/test/java/google/registry/eppserver/quota/LocalConnectionLimiterTest.java index 0249ab692..6922c5afe 100644 --- a/core/src/test/java/google/registry/eppserver/quota/LocalConnectionLimiterTest.java +++ b/core/src/test/java/google/registry/eppserver/quota/LocalConnectionLimiterTest.java @@ -65,37 +65,37 @@ class LocalConnectionLimiterTest { } @Test - void testAcquireCert_successUpToLimit() { - assertThat(limiter.acquireCert("cert_hash_1")).isTrue(); - assertThat(limiter.acquireCert("cert_hash_1")).isTrue(); + void testAcquireRegistrar_successUpToLimit() { + assertThat(limiter.acquireRegistrar("registrar_1")).isTrue(); + assertThat(limiter.acquireRegistrar("registrar_1")).isTrue(); } @Test - void testAcquireCert_rejectsOverLimit() { - assertThat(limiter.acquireCert("cert_hash_1")).isTrue(); - assertThat(limiter.acquireCert("cert_hash_1")).isTrue(); - // 3rd attempt from same cert should be rejected - assertThat(limiter.acquireCert("cert_hash_1")).isFalse(); + void testAcquireRegistrar_rejectsOverLimit() { + assertThat(limiter.acquireRegistrar("registrar_1")).isTrue(); + assertThat(limiter.acquireRegistrar("registrar_1")).isTrue(); + // 3rd attempt from same registrar should be rejected + assertThat(limiter.acquireRegistrar("registrar_1")).isFalse(); } @Test - void testAcquireCert_independentAcrossCerts() { - assertThat(limiter.acquireCert("cert_hash_1")).isTrue(); - assertThat(limiter.acquireCert("cert_hash_1")).isTrue(); - assertThat(limiter.acquireCert("cert_hash_1")).isFalse(); + void testAcquireRegistrar_independentAcrossRegistrars() { + assertThat(limiter.acquireRegistrar("registrar_1")).isTrue(); + assertThat(limiter.acquireRegistrar("registrar_1")).isTrue(); + assertThat(limiter.acquireRegistrar("registrar_1")).isFalse(); - // A different cert should still be allowed - assertThat(limiter.acquireCert("cert_hash_2")).isTrue(); + // A different registrar should still be allowed + assertThat(limiter.acquireRegistrar("registrar_2")).isTrue(); } @Test - void testReleaseCert_freesSlot() { - assertThat(limiter.acquireCert("cert_hash_1")).isTrue(); - assertThat(limiter.acquireCert("cert_hash_1")).isTrue(); - assertThat(limiter.acquireCert("cert_hash_1")).isFalse(); + void testReleaseRegistrar_freesSlot() { + assertThat(limiter.acquireRegistrar("registrar_1")).isTrue(); + assertThat(limiter.acquireRegistrar("registrar_1")).isTrue(); + assertThat(limiter.acquireRegistrar("registrar_1")).isFalse(); - limiter.releaseCert("cert_hash_1"); + limiter.releaseRegistrar("registrar_1"); // Now we should be able to acquire again - assertThat(limiter.acquireCert("cert_hash_1")).isTrue(); + assertThat(limiter.acquireRegistrar("registrar_1")).isTrue(); } } diff --git a/util/src/main/java/google/registry/util/ProxyHttpHeaders.java b/util/src/main/java/google/registry/util/ProxyHttpHeaders.java index 7a3a9a46a..471905e8c 100644 --- a/util/src/main/java/google/registry/util/ProxyHttpHeaders.java +++ b/util/src/main/java/google/registry/util/ProxyHttpHeaders.java @@ -33,14 +33,11 @@ public final class ProxyHttpHeaders { /** 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. - * - *

Note that Java 17's servlet implementation may inject some seemingly unrelated addresses - * into this header. We only use this as a fallback so the proxy can transition to use the above - * header that should not be interfered with. - */ + /** Fallback HTTP header name used to pass the client IP address from the proxy to Nomulus. */ public static final String FALLBACK_IP_ADDRESS = HttpHeaders.X_FORWARDED_FOR; + /** HTTP header name used to pass the authenticated Registrar Id from Nomulus to GKE. */ + public static final String LOGGED_IN_REGISTRAR = "Nomulus-Logged-In-Registrar"; + private ProxyHttpHeaders() {} }