1 /* 2 * Copyright (C) 2024 The Android Open Source Project 3 * 4 * Licensed under the Apache License, Version 2.0 (the "License"); 5 * you may not use this file except in compliance with the License. 6 * You may obtain a copy of the License at 7 * 8 * http://www.apache.org/licenses/LICENSE-2.0 9 * 10 * Unless required by applicable law or agreed to in writing, software 11 * distributed under the License is distributed on an "AS IS" BASIS, 12 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 * See the License for the specific language governing permissions and 14 * limitations under the License. 15 */ 16 17 package android.test.wifi.nsd; 18 19 import static org.junit.Assert.assertEquals; 20 import static org.junit.Assert.assertTrue; 21 import static org.junit.Assert.fail; 22 23 import android.content.Context; 24 import android.net.nsd.NsdManager; 25 import android.net.nsd.NsdServiceInfo; 26 import android.net.wifi.WifiManager; 27 import android.util.Log; 28 import java.io.BufferedReader; 29 import java.io.IOException; 30 import java.io.InputStreamReader; 31 import java.io.PrintWriter; 32 import java.net.ServerSocket; 33 import java.net.Socket; 34 import java.util.concurrent.CountDownLatch; 35 import java.util.concurrent.Executors; 36 import java.util.concurrent.TimeUnit; 37 38 /** Helper class for NsdManager. */ 39 public final class NsdHelper { 40 private final NsdManager nsdManager; 41 private final WifiManager wifi; 42 private WifiManager.MulticastLock multicastLock; 43 44 private String serviceName = "WifiNsdTest"; 45 private static final String SERVICE_TYPE = "_wifi_nsd_test._tcp."; 46 private static final String TAG = "WifiTest-NsdInstrumentationTest"; 47 private static final String PING = "Hello"; 48 private static final String PONG = "World"; 49 NsdHelper(Context context, String testId)50 public NsdHelper(Context context, String testId) { 51 this.wifi = context.getSystemService(WifiManager.class); 52 // Previously unregistered services may still be discoverable. 53 // Use unique service name to prevent the discovery of unregistered services. 54 this.serviceName += "-" + testId; 55 this.multicastLock = this.wifi.createMulticastLock("multicastLock"); 56 this.nsdManager = context.getSystemService(NsdManager.class); 57 } 58 await(CountDownLatch latch)59 private void await(CountDownLatch latch) throws InterruptedException { 60 assertTrue(latch.await(60, TimeUnit.SECONDS)); 61 } 62 serviceTest()63 public void serviceTest() throws InterruptedException, IOException { 64 ServerSocket serverSocket = new ServerSocket(0); 65 RegistrationListener listener = registerService(serverSocket.getLocalPort()); 66 await(listener.serviceRegistered); 67 68 Socket clientSocket = serverSocket.accept(); 69 70 // Wait for ping message. 71 BufferedReader in = new BufferedReader(new InputStreamReader(clientSocket.getInputStream())); 72 String msg = in.readLine(); 73 assertEquals(PING, msg); 74 75 // Send pong message. 76 PrintWriter out = new PrintWriter(clientSocket.getOutputStream(), true); 77 out.println(PONG); 78 79 serverSocket.close(); 80 clientSocket.close(); 81 82 unregisterService(listener); 83 await(listener.serviceUnregistered); 84 } 85 86 private static class RegistrationListener implements NsdManager.RegistrationListener { 87 CountDownLatch serviceRegistered; 88 CountDownLatch serviceUnregistered; 89 RegistrationListener()90 RegistrationListener() { 91 serviceRegistered = new CountDownLatch(1); 92 serviceUnregistered = new CountDownLatch(1); 93 } 94 95 @Override onServiceRegistered(NsdServiceInfo serviceInfo)96 public void onServiceRegistered(NsdServiceInfo serviceInfo) { 97 Log.d(TAG, "Service registered. NsdServiceInfo:" + serviceInfo); 98 serviceRegistered.countDown(); 99 } 100 101 @Override onServiceUnregistered(NsdServiceInfo serviceInfo)102 public void onServiceUnregistered(NsdServiceInfo serviceInfo) { 103 Log.d(TAG, "Service unregistered"); 104 serviceUnregistered.countDown(); 105 } 106 107 @Override onRegistrationFailed(NsdServiceInfo serviceInfo, int errorCode)108 public void onRegistrationFailed(NsdServiceInfo serviceInfo, int errorCode) { 109 fail("Registration failed"); 110 } 111 112 @Override onUnregistrationFailed(NsdServiceInfo serviceInfo, int errorCode)113 public void onUnregistrationFailed(NsdServiceInfo serviceInfo, int errorCode) { 114 fail("Unregistration failed"); 115 } 116 } 117 registerService(int port)118 private RegistrationListener registerService(int port) { 119 RegistrationListener listener = new RegistrationListener(); 120 NsdServiceInfo serviceInfo = new NsdServiceInfo(); 121 serviceInfo.setPort(port); 122 serviceInfo.setServiceName(serviceName); 123 serviceInfo.setServiceType(SERVICE_TYPE); 124 125 nsdManager.registerService(serviceInfo, NsdManager.PROTOCOL_DNS_SD, listener); 126 return listener; 127 } 128 unregisterService(RegistrationListener listener)129 private CountDownLatch unregisterService(RegistrationListener listener) { 130 nsdManager.unregisterService(listener); 131 return listener.serviceUnregistered; 132 } 133 discoverTest()134 public void discoverTest() throws InterruptedException, IOException { 135 DiscoveryListener listener = discoverServices(); 136 await(listener.serviceFound); 137 ServiceInfoCallback callback = resolveServices(listener.service); 138 await(callback.serviceResolved); 139 CountDownLatch discoveryStopped = stopDiscovery(listener); 140 await(discoveryStopped); 141 142 // Set up connection. 143 NsdServiceInfo service = callback.service; 144 Socket clientSocket = new Socket(service.getHost(), service.getPort()); 145 146 // Send ping message. 147 PrintWriter out = new PrintWriter(clientSocket.getOutputStream(), true); 148 out.println(PING); 149 150 // Wait for pong message. 151 BufferedReader in = new BufferedReader(new InputStreamReader(clientSocket.getInputStream())); 152 String msg = in.readLine(); 153 assertEquals(PONG, msg); 154 clientSocket.close(); 155 } 156 157 private static class DiscoveryListener implements NsdManager.DiscoveryListener { 158 CountDownLatch serviceFound; 159 CountDownLatch discoveryStopped; 160 NsdServiceInfo service; 161 162 private String serviceName; 163 DiscoveryListener(String serviceName)164 DiscoveryListener(String serviceName) { 165 serviceFound = new CountDownLatch(1); 166 discoveryStopped = new CountDownLatch(1); 167 this.serviceName = serviceName; 168 } 169 170 @Override onDiscoveryStarted(String regType)171 public void onDiscoveryStarted(String regType) {} 172 173 @Override onServiceFound(NsdServiceInfo serviceInfo)174 public void onServiceFound(NsdServiceInfo serviceInfo) { 175 if (serviceInfo.getServiceType().equals(SERVICE_TYPE) 176 && serviceInfo.getServiceName().equals(serviceName)) { 177 service = serviceInfo; 178 serviceFound.countDown(); 179 } 180 } 181 182 @Override onServiceLost(NsdServiceInfo nsdServiceInfo)183 public void onServiceLost(NsdServiceInfo nsdServiceInfo) {} 184 185 @Override onDiscoveryStopped(String serviceType)186 public void onDiscoveryStopped(String serviceType) { 187 discoveryStopped.countDown(); 188 } 189 190 @Override onStartDiscoveryFailed(String serviceType, int errorCode)191 public void onStartDiscoveryFailed(String serviceType, int errorCode) { 192 fail("Failed to start discovery"); 193 } 194 195 @Override onStopDiscoveryFailed(String serviceType, int errorCode)196 public void onStopDiscoveryFailed(String serviceType, int errorCode) { 197 fail("Failed to stop discovery"); 198 } 199 } 200 discoverServices()201 private DiscoveryListener discoverServices() { 202 DiscoveryListener discoveryListener = new DiscoveryListener(serviceName); 203 204 multicastLock.setReferenceCounted(true); 205 multicastLock.acquire(); 206 207 nsdManager.discoverServices(SERVICE_TYPE, NsdManager.PROTOCOL_DNS_SD, discoveryListener); 208 209 return discoveryListener; 210 } 211 resolveServices(NsdServiceInfo service)212 private ServiceInfoCallback resolveServices(NsdServiceInfo service) { 213 ServiceInfoCallback callback = new ServiceInfoCallback(serviceName); 214 nsdManager.registerServiceInfoCallback(service, Executors.newSingleThreadExecutor(), callback); 215 216 return callback; 217 } 218 219 private static class ServiceInfoCallback implements NsdManager.ServiceInfoCallback { 220 CountDownLatch serviceResolved; 221 NsdServiceInfo service; 222 223 private String serviceName; 224 ServiceInfoCallback(String serviceName)225 ServiceInfoCallback(String serviceName) { 226 serviceResolved = new CountDownLatch(1); 227 this.serviceName = serviceName; 228 } 229 230 @Override onServiceUpdated(NsdServiceInfo serviceInfo)231 public void onServiceUpdated(NsdServiceInfo serviceInfo) { 232 if (serviceInfo.getServiceName().equals(serviceName)) { 233 service = serviceInfo; 234 serviceResolved.countDown(); 235 } 236 } 237 238 @Override onServiceInfoCallbackRegistrationFailed(int errorCode)239 public void onServiceInfoCallbackRegistrationFailed(int errorCode) { 240 fail("Callback registration failed"); 241 } 242 243 @Override onServiceInfoCallbackUnregistered()244 public void onServiceInfoCallbackUnregistered() {} 245 246 @Override onServiceLost()247 public void onServiceLost() {} 248 } 249 stopDiscovery(DiscoveryListener listener)250 private CountDownLatch stopDiscovery(DiscoveryListener listener) { 251 nsdManager.stopServiceDiscovery(listener); 252 if (multicastLock.isHeld()) { 253 multicastLock.release(); // release after browsing 254 } 255 return listener.discoveryStopped; 256 } 257 } 258