From 362c3fbc32a042125f1aae64b31cd55657beb396 Mon Sep 17 00:00:00 2001
From: Pierre Precourt <pprecourt@google.com>
Date: Mon, 2 Dec 2024 16:32:08 -0800
Subject: [PATCH] Provide getters for the callback address and port in the
 Tsunami callback client.

PiperOrigin-RevId: 702120836
Change-Id: I7efa6684b5cca5f734ee2aac8cb99717a6a82161
---
 .../com/google/tsunami/plugin/TcsClient.java  | 21 ++++++++++-
 .../google/tsunami/plugin/TcsClientTest.java  | 36 +++++++++++++++++++
 2 files changed, 56 insertions(+), 1 deletion(-)

diff --git a/plugin/src/main/java/com/google/tsunami/plugin/TcsClient.java b/plugin/src/main/java/com/google/tsunami/plugin/TcsClient.java
index bcf8e79c..9a748c4e 100644
--- a/plugin/src/main/java/com/google/tsunami/plugin/TcsClient.java
+++ b/plugin/src/main/java/com/google/tsunami/plugin/TcsClient.java
@@ -89,7 +89,26 @@ public String getCallbackUri(String secretString) {
       return CbidProcessor.addCbidToSubdomain(cbid, hostAndPort);
     }
     // Should never reach here
-    throw new AssertionError("Unrecognized address format, should be Ip address or valid domain");
+    throw new AssertionError("Unrecognized address format, should be IP address or valid domain");
+  }
+
+  public String getCallbackAddress() {
+    if (InetAddresses.isInetAddress(callbackAddress)) {
+      return callbackAddress;
+    } else if (InternetDomainName.isValid(callbackAddress)) {
+      return callbackAddress;
+    }
+
+    // Should never reach here
+    throw new AssertionError("Unrecognized address format, should be IP address or valid domain");
+  }
+
+  public int getCallbackPort() {
+    if (!isValidPortNumber(callbackPort)) {
+      throw new AssertionError("Invalid callbackPort number specified");
+    }
+
+    return callbackPort;
   }
 
   public boolean hasOobLog(String secretString) {
diff --git a/plugin/src/test/java/com/google/tsunami/plugin/TcsClientTest.java b/plugin/src/test/java/com/google/tsunami/plugin/TcsClientTest.java
index bbfc415c..c9d1d6ba 100644
--- a/plugin/src/test/java/com/google/tsunami/plugin/TcsClientTest.java
+++ b/plugin/src/test/java/com/google/tsunami/plugin/TcsClientTest.java
@@ -110,6 +110,42 @@ public void getCallbackUri_invalidCallbackPort_throwsError() {
     assertThrows(AssertionError.class, () -> client.getCallbackUri(SECRET));
   }
 
+  @Test
+  public void getCallbackAddress_validIpv4Address_returnsAddress() {
+    client = new TcsClient(VALID_IPV4_ADDRESS, VALID_PORT, VALID_URL, httpClient);
+    assertThat(client.getCallbackAddress()).isEqualTo(VALID_IPV4_ADDRESS);
+  }
+
+  @Test
+  public void getCallbackAddress_validIpv6Address_returnsAddress() {
+    client = new TcsClient(VALID_IPV6_ADDRESS, VALID_PORT, VALID_URL, httpClient);
+    assertThat(client.getCallbackAddress()).isEqualTo(VALID_IPV6_ADDRESS);
+  }
+
+  @Test
+  public void getCallbackAddress_validDomainAddress_returnsAddress() {
+    client = new TcsClient(VALID_DOMAIN, VALID_PORT, VALID_URL, httpClient);
+    assertThat(client.getCallbackAddress()).isEqualTo(VALID_DOMAIN);
+  }
+
+  @Test
+  public void getCallbackAddress_invalidAddress_throwsError() {
+    client = new TcsClient(INVALID_ADDRESS, 100000, VALID_URL, httpClient);
+    assertThrows(AssertionError.class, () -> client.getCallbackAddress());
+  }
+
+  @Test
+  public void getCallbackPort_validPort_returnsPort() {
+    client = new TcsClient(VALID_DOMAIN, VALID_PORT, VALID_URL, httpClient);
+    assertThat(client.getCallbackPort()).isEqualTo(VALID_PORT);
+  }
+
+  @Test
+  public void getCallbackPort_invalidPort_throwsError() {
+    client = new TcsClient(VALID_DOMAIN, 100000, VALID_URL, httpClient);
+    assertThrows(AssertionError.class, () -> client.getCallbackPort());
+  }
+
   @Test
   public void isVulnerable_sendsValidPollingRequest() throws IOException, InterruptedException {
     MockWebServer mockWebServer = new MockWebServer();