Add an HTTP header to response from Nomulus after successful login (#879)

* Add a logged-in response header

* small fixes

* Refactor EPP test cases to check for headers

* small change

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

@@ -74,6 +74,14 @@ public class EppRequestHandler {
           && eppOutput.getResponse().getResult().getCode() == SUCCESS_AND_CLOSE) {
         response.setHeader("Epp-Session", "close");
       }
+      // If a login request returns a success, a logged-in header is added to the response to inform
+      // the proxy that it is no longer necessary to send the full client certificate to the backend
+      // for this connection.
+      if (eppOutput.isResponse()
+          && eppOutput.getResponse().isLoginResponse()
+          && eppOutput.isSuccess()) {
+        response.setHeader("Logged-In", "true");
+      }
     } catch (Exception e) {
       logger.atWarning().withCause(e).log("handleEppCommand general exception");
       response.setStatus(SC_BAD_REQUEST);

core/src/main/java/google/registry/flows/session/LoginFlow.java

@@ -141,7 +141,7 @@ public class LoginFlow implements Flow {
     sessionMetadata.resetFailedLoginAttempts();
     sessionMetadata.setClientId(login.getClientId());
     sessionMetadata.setServiceExtensionUris(serviceExtensionUrisBuilder.build());
-    return responseBuilder.build();
+    return responseBuilder.setIsLoginResponse().build();
   }
 
   /** Registrar with this client ID could not be found. */

core/src/main/java/google/registry/model/eppoutput/EppResponse.java

@@ -65,6 +65,7 @@ import javax.xml.bind.annotation.XmlElement;
 import javax.xml.bind.annotation.XmlElementRef;
 import javax.xml.bind.annotation.XmlElementRefs;
 import javax.xml.bind.annotation.XmlElementWrapper;
+import javax.xml.bind.annotation.XmlTransient;
 import javax.xml.bind.annotation.XmlType;
 
 /**
@@ -87,6 +88,9 @@ public class EppResponse extends ImmutableObject implements ResponseOrGreeting {
   /** The command result. The RFC allows multiple failure results, but we always return one. */
   Result result;
 
+  /** Indicates if this response is for a login request. */
+  @XmlTransient boolean isLoginResponse = false;
+
   /**
    * Information about messages queued for retrieval. This may appear in response to any EPP message
    * (if messages are queued), but in practice this will only be set in response to a poll request.
@@ -178,6 +182,10 @@ public class EppResponse extends ImmutableObject implements ResponseOrGreeting {
     return result;
   }
 
+  public boolean isLoginResponse() {
+    return isLoginResponse;
+  }
+
   /** Marker interface for types that can go in the {@link #resData} field. */
   public interface ResponseData {}
 
@@ -222,5 +230,10 @@ public class EppResponse extends ImmutableObject implements ResponseOrGreeting {
       getInstance().extensions = forceEmptyToNull(extensions);
       return this;
     }
+
+    public Builder setIsLoginResponse() {
+      getInstance().isLoginResponse = true;
+      return this;
+    }
   }
 }

core/src/test/java/google/registry/flows/EppLifecycleDomainTest.java

@@ -501,7 +501,7 @@ class EppLifecycleDomainTest extends EppTestCase {
 
   @Test
   void testEapDomainDeletion_withinAddGracePeriod_eapFeeIsNotRefunded() throws Exception {
-    assertThatCommand("login_valid_fee_extension.xml").hasResponse("generic_success_response.xml");
+    assertThatCommand("login_valid_fee_extension.xml").hasSuccessfulLogin();
     createContacts(DateTime.parse("2000-06-01T00:00:00Z"));
 
     // Set the EAP schedule.
@@ -718,7 +718,7 @@ class EppLifecycleDomainTest extends EppTestCase {
             START_OF_TIME, PREDELEGATION,
             gaDate, GENERAL_AVAILABILITY));
 
-    assertThatCommand("login_valid_fee_extension.xml").hasResponse("generic_success_response.xml");
+    assertThatCommand("login_valid_fee_extension.xml").hasSuccessfulLogin();
 
     assertThatCommand("domain_check_fee_premium.xml")
         .atTime(gaDate.plusDays(1))
@@ -1196,7 +1196,7 @@ class EppLifecycleDomainTest extends EppTestCase {
 
     assertThatLogin("NewRegistrar", "foo-BAR2")
         .atTime(sunriseDate.minusDays(3))
-        .hasResponse("generic_success_response.xml");
+        .hasSuccessfulLogin();
 
     createContactsAndHosts();
 
@@ -1292,7 +1292,7 @@ class EppLifecycleDomainTest extends EppTestCase {
 
     assertThatLogin("NewRegistrar", "foo-BAR2")
         .atTime(sunriseDate.minusDays(3))
-        .hasResponse("generic_success_response.xml");
+        .hasSuccessfulLogin();
 
     createContactsAndHosts();
 

core/src/test/java/google/registry/flows/EppTestCase.java

@@ -121,6 +121,10 @@ public class EppTestCase {
       return assertCommandAndResponse(
           inputFilename, inputSubstitutions, outputFilename, outputSubstitutions, now);
     }
+
+    public String hasSuccessfulLogin() throws Exception {
+      return assertLoginCommandAndResponse(inputFilename, inputSubstitutions, null, now);
+    }
   }
 
   protected CommandAsserter assertThatCommand(String inputFilename) {
@@ -137,13 +141,33 @@ public class EppTestCase {
   }
 
   protected void assertThatLoginSucceeds(String clientId, String password) throws Exception {
-    assertThatLogin(clientId, password).hasResponse("generic_success_response.xml");
+    assertThatLogin(clientId, password).hasSuccessfulLogin();
   }
 
   protected void assertThatLogoutSucceeds() throws Exception {
     assertThatCommand("logout.xml").hasResponse("logout_response.xml");
   }
 
+  private String assertLoginCommandAndResponse(
+      String inputFilename,
+      @Nullable Map<String, String> inputSubstitutions,
+      @Nullable Map<String, String> outputSubstitutions,
+      DateTime now)
+      throws Exception {
+    String outputFilename = "generic_success_response.xml";
+    clock.setTo(now);
+    String input = loadFile(EppTestCase.class, inputFilename, inputSubstitutions);
+    String expectedOutput = loadFile(EppTestCase.class, outputFilename, outputSubstitutions);
+    setUpSession();
+    FakeResponse response = executeXmlCommand(input);
+
+    // Check that the logged-in header was added to the response
+    assertThat(response.getHeaders()).isEqualTo(ImmutableMap.of("Logged-In", "true"));
+
+    return verifyAndReturnOutput(
+        response.getPayload(), expectedOutput, inputFilename, outputFilename);
+  }
+
   private String assertCommandAndResponse(
       String inputFilename,
       @Nullable Map<String, String> inputSubstitutions,
@@ -154,6 +178,18 @@ public class EppTestCase {
     clock.setTo(now);
     String input = loadFile(EppTestCase.class, inputFilename, inputSubstitutions);
     String expectedOutput = loadFile(EppTestCase.class, outputFilename, outputSubstitutions);
+    setUpSession();
+    FakeResponse response = executeXmlCommand(input);
+
+    // Checks that the Logged-In header is not in the response. If testing the login command, use
+    // assertLoginCommandAndResponse instead of this method.
+    assertThat(response.getHeaders()).doesNotContainEntry("Logged-In", "true");
+
+    return verifyAndReturnOutput(
+        response.getPayload(), expectedOutput, inputFilename, outputFilename);
+  }
+
+  private void setUpSession() {
     if (sessionMetadata == null) {
       sessionMetadata =
           new HttpSessionMetadata(new FakeHttpSession()) {
@@ -165,7 +201,13 @@ public class EppTestCase {
             }
           };
     }
-    String actualOutput = executeXmlCommand(input);
+  }
+
+  private String verifyAndReturnOutput(
+      String actualOutput, String expectedOutput, String inputFilename, String outputFilename)
+      throws Exception {
+    // Run the resulting xml through the unmarshaller to verify that it was valid.
+    EppXmlTransformer.validateOutput(actualOutput);
     assertXmlEqualsWithMessage(
         expectedOutput,
         actualOutput,
@@ -176,7 +218,7 @@ public class EppTestCase {
     return actualOutput;
   }
 
-  private String executeXmlCommand(String inputXml) throws Exception {
+  private FakeResponse executeXmlCommand(String inputXml) throws Exception {
     EppRequestHandler handler = new EppRequestHandler();
     FakeResponse response = new FakeResponse();
     handler.response = response;
@@ -195,10 +237,7 @@ public class EppTestCase {
         inputXml.getBytes(UTF_8));
     assertThat(response.getStatus()).isEqualTo(SC_OK);
     assertThat(response.getContentType()).isEqualTo(APPLICATION_EPP_XML_UTF8);
-    String result = response.getPayload();
-    // Run the resulting xml through the unmarshaller to verify that it was valid.
-    EppXmlTransformer.validateOutput(result);
-    return result;
+    return response;
   }
 
   EppMetric getRecordedEppMetric() {

core/src/test/java/google/registry/flows/session/LoginFlowTestCase.java

@@ -14,6 +14,7 @@
 
 package google.registry.flows.session;
 
+import static com.google.common.truth.Truth.assertThat;
 import static google.registry.testing.DatabaseHelper.deleteResource;
 import static google.registry.testing.DatabaseHelper.loadRegistrar;
 import static google.registry.testing.DatabaseHelper.persistResource;
@@ -32,6 +33,7 @@ import google.registry.flows.session.LoginFlow.PasswordChangesNotSupportedExcept
 import google.registry.flows.session.LoginFlow.RegistrarAccountNotActiveException;
 import google.registry.flows.session.LoginFlow.TooManyFailedLoginsException;
 import google.registry.flows.session.LoginFlow.UnsupportedLanguageException;
+import google.registry.model.eppoutput.EppOutput;
 import google.registry.model.registrar.Registrar;
 import google.registry.model.registrar.Registrar.State;
 import org.junit.jupiter.api.BeforeEach;
@@ -74,6 +76,14 @@ public abstract class LoginFlowTestCase extends FlowTestCase<LoginFlow> {
     doSuccessfulTest("login_valid.xml");
   }
 
+  @Test
+  void testSuccess_setsIsLoginResponse() throws Exception {
+    setEppInput("login_valid.xml");
+    assertTransactionalFlow(false);
+    EppOutput output = runFlow();
+    assertThat(output.getResponse().isLoginResponse()).isTrue();
+  }
+
   @Test
   void testSuccess_suspendedRegistrar() throws Exception {
     persistResource(getRegistrarBuilder().setState(State.SUSPENDED).build());

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

@@ -249,6 +249,20 @@ class EppServiceHandlerTest {
     assertThat(channel.isActive()).isFalse();
   }
 
+  @Test
+  void sendResponseToNextHandler_unrecognizedHeader() throws Exception {
+    setHandshakeSuccess();
+    String content = "<epp>stuff</epp>";
+    HttpResponse response = makeEppHttpResponse(content, HttpResponseStatus.OK);
+    response.headers().set("unrecognized-header", "test");
+    channel.writeOutbound(response);
+    ByteBuf expectedResponse = channel.readOutbound();
+    assertThat(Unpooled.wrappedBuffer(content.getBytes(UTF_8))).isEqualTo(expectedResponse);
+    // Nothing further to pass to the next handler.
+    assertThat((Object) channel.readOutbound()).isNull();
+    assertThat(channel.isActive()).isTrue();
+  }
+
   @Test
   void testFailure_disconnectOnNonOKResponseStatus() throws Exception {
     setHandshakeSuccess();