Merge "Harden NsdManager against null-dereference crashes"
diff --git a/core/java/android/net/nsd/NsdManager.java b/core/java/android/net/nsd/NsdManager.java
index 9c3e405..7b2c623 100644
--- a/core/java/android/net/nsd/NsdManager.java
+++ b/core/java/android/net/nsd/NsdManager.java
@@ -301,27 +301,36 @@
 
         @Override
         public void handleMessage(Message message) {
-            Object listener = getListener(message.arg2);
-            boolean listenerRemove = true;
             switch (message.what) {
                 case AsyncChannel.CMD_CHANNEL_HALF_CONNECTED:
                     mAsyncChannel.sendMessage(AsyncChannel.CMD_CHANNEL_FULL_CONNECTION);
-                    break;
+                    return;
                 case AsyncChannel.CMD_CHANNEL_FULLY_CONNECTED:
                     mConnected.countDown();
-                    break;
+                    return;
                 case AsyncChannel.CMD_CHANNEL_DISCONNECTED:
                     Log.e(TAG, "Channel lost");
+                    return;
+                default:
                     break;
+            }
+            Object listener = getListener(message.arg2);
+            if (listener == null) {
+                Log.d(TAG, "Stale key " + message.arg2);
+                return;
+            }
+            boolean listenerRemove = true;
+            NsdServiceInfo ns = getNsdService(message.arg2);
+            switch (message.what) {
                 case DISCOVER_SERVICES_STARTED:
-                    String s = ((NsdServiceInfo) message.obj).getServiceType();
+                    String s = getNsdServiceInfoType((NsdServiceInfo) message.obj);
                     ((DiscoveryListener) listener).onDiscoveryStarted(s);
                     // Keep listener until stop discovery
                     listenerRemove = false;
                     break;
                 case DISCOVER_SERVICES_FAILED:
-                    ((DiscoveryListener) listener).onStartDiscoveryFailed(
-                            getNsdService(message.arg2).getServiceType(), message.arg1);
+                    ((DiscoveryListener) listener).onStartDiscoveryFailed(getNsdServiceInfoType(ns),
+                            message.arg1);
                     break;
                 case SERVICE_FOUND:
                     ((DiscoveryListener) listener).onServiceFound((NsdServiceInfo) message.obj);
@@ -334,16 +343,14 @@
                     listenerRemove = false;
                     break;
                 case STOP_DISCOVERY_FAILED:
-                    ((DiscoveryListener) listener).onStopDiscoveryFailed(
-                            getNsdService(message.arg2).getServiceType(), message.arg1);
+                    ((DiscoveryListener) listener).onStopDiscoveryFailed(getNsdServiceInfoType(ns),
+                            message.arg1);
                     break;
                 case STOP_DISCOVERY_SUCCEEDED:
-                    ((DiscoveryListener) listener).onDiscoveryStopped(
-                            getNsdService(message.arg2).getServiceType());
+                    ((DiscoveryListener) listener).onDiscoveryStopped(getNsdServiceInfoType(ns));
                     break;
                 case REGISTER_SERVICE_FAILED:
-                    ((RegistrationListener) listener).onRegistrationFailed(
-                            getNsdService(message.arg2), message.arg1);
+                    ((RegistrationListener) listener).onRegistrationFailed(ns, message.arg1);
                     break;
                 case REGISTER_SERVICE_SUCCEEDED:
                     ((RegistrationListener) listener).onServiceRegistered(
@@ -352,16 +359,13 @@
                     listenerRemove = false;
                     break;
                 case UNREGISTER_SERVICE_FAILED:
-                    ((RegistrationListener) listener).onUnregistrationFailed(
-                            getNsdService(message.arg2), message.arg1);
+                    ((RegistrationListener) listener).onUnregistrationFailed(ns, message.arg1);
                     break;
                 case UNREGISTER_SERVICE_SUCCEEDED:
-                    ((RegistrationListener) listener).onServiceUnregistered(
-                            getNsdService(message.arg2));
+                    ((RegistrationListener) listener).onServiceUnregistered(ns);
                     break;
                 case RESOLVE_SERVICE_FAILED:
-                    ((ResolveListener) listener).onResolveFailed(
-                            getNsdService(message.arg2), message.arg1);
+                    ((ResolveListener) listener).onResolveFailed(ns, message.arg1);
                     break;
                 case RESOLVE_SERVICE_SUCCEEDED:
                     ((ResolveListener) listener).onServiceResolved((NsdServiceInfo) message.obj);
@@ -421,6 +425,11 @@
     }
 
 
+    private String getNsdServiceInfoType(NsdServiceInfo s) {
+        if (s == null) return "?";
+        return s.getServiceType();
+    }
+
     /**
      * Initialize AsyncChannel
      */