Add callbacks for service offload

Components that can provide offload like IpClient (packet
filter offloading) can use the API to register a callback to be notified
when offload is necessary.

Bug: 269240366
Test: atest CtsNetTestCases
Change-Id: I8080702f5b530001b88e79e504f4722ac01bc576
diff --git a/service-t/Android.bp b/service-t/Android.bp
index 7de749c..c277cf6 100644
--- a/service-t/Android.bp
+++ b/service-t/Android.bp
@@ -88,8 +88,9 @@
     name: "service-connectivity-mdns-standalone-build-test",
     sdk_version: "core_platform",
     srcs: [
-        ":service-mdns-droidstubs",
         "src/com/android/server/connectivity/mdns/**/*.java",
+        ":framework-connectivity-t-mdns-standalone-build-sources",
+        ":service-mdns-droidstubs"
     ],
     exclude_srcs: [
         "src/com/android/server/connectivity/mdns/internal/SocketNetlinkMonitor.java",
diff --git a/service-t/src/com/android/server/NsdService.java b/service-t/src/com/android/server/NsdService.java
index 745c5bc..1a05d46 100644
--- a/service-t/src/com/android/server/NsdService.java
+++ b/service-t/src/com/android/server/NsdService.java
@@ -16,6 +16,7 @@
 
 package com.android.server;
 
+import static android.Manifest.permission.NETWORK_SETTINGS;
 import static android.net.ConnectivityManager.NETID_UNSET;
 import static android.net.NetworkCapabilities.TRANSPORT_VPN;
 import static android.net.NetworkCapabilities.TRANSPORT_WIFI;
@@ -46,9 +47,12 @@
 import android.net.nsd.INsdManager;
 import android.net.nsd.INsdManagerCallback;
 import android.net.nsd.INsdServiceConnector;
+import android.net.nsd.IOffloadEngine;
 import android.net.nsd.MDnsManager;
 import android.net.nsd.NsdManager;
 import android.net.nsd.NsdServiceInfo;
+import android.net.nsd.OffloadEngine;
+import android.net.nsd.OffloadServiceInfo;
 import android.net.wifi.WifiManager;
 import android.os.Binder;
 import android.os.Handler;
@@ -56,6 +60,7 @@
 import android.os.IBinder;
 import android.os.Looper;
 import android.os.Message;
+import android.os.RemoteCallbackList;
 import android.os.RemoteException;
 import android.os.UserHandle;
 import android.provider.DeviceConfig;
@@ -98,6 +103,7 @@
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.Objects;
 import java.util.regex.Matcher;
 import java.util.regex.Pattern;
 
@@ -215,6 +221,24 @@
     // The number of client that ever connected.
     private int mClientNumberId = 1;
 
+    private final RemoteCallbackList<IOffloadEngine> mOffloadEngines =
+            new RemoteCallbackList<>();
+
+    private static class OffloadEngineInfo {
+        @NonNull final String mInterfaceName;
+        final long mOffloadCapabilities;
+        final long mOffloadType;
+        @NonNull final IOffloadEngine mOffloadEngine;
+
+        OffloadEngineInfo(@NonNull IOffloadEngine offloadEngine,
+                @NonNull String interfaceName, long capabilities, long offloadType) {
+            this.mOffloadEngine = offloadEngine;
+            this.mInterfaceName = interfaceName;
+            this.mOffloadCapabilities = capabilities;
+            this.mOffloadType = offloadType;
+        }
+    }
+
     private static class MdnsListener implements MdnsServiceBrowserListener {
         protected final int mClientRequestId;
         protected final int mTransactionId;
@@ -719,6 +743,7 @@
                 final int transactionId;
                 final int clientRequestId = msg.arg2;
                 final ListenerArgs args;
+                final OffloadEngineInfo offloadEngineInfo;
                 switch (msg.what) {
                     case NsdManager.DISCOVER_SERVICES: {
                         if (DBG) Log.d(TAG, "Discover services");
@@ -1114,6 +1139,16 @@
                             return NOT_HANDLED;
                         }
                         break;
+                    case NsdManager.REGISTER_OFFLOAD_ENGINE:
+                        offloadEngineInfo = (OffloadEngineInfo) msg.obj;
+                        // TODO: Limits the number of registrations created by a given class.
+                        mOffloadEngines.register(offloadEngineInfo.mOffloadEngine,
+                                offloadEngineInfo);
+                        // TODO: Sends all the existing OffloadServiceInfos back.
+                        break;
+                    case NsdManager.UNREGISTER_OFFLOAD_ENGINE:
+                        mOffloadEngines.unregister((IOffloadEngine) msg.obj);
+                        break;
                     default:
                         return NOT_HANDLED;
                 }
@@ -1771,7 +1806,42 @@
         }
     }
 
+    private void sendOffloadServiceInfosUpdate(@NonNull String targetInterfaceName,
+            @NonNull OffloadServiceInfo offloadServiceInfo, boolean isRemove) {
+        final int count = mOffloadEngines.beginBroadcast();
+        try {
+            for (int i = 0; i < count; i++) {
+                final OffloadEngineInfo offloadEngineInfo =
+                        (OffloadEngineInfo) mOffloadEngines.getBroadcastCookie(i);
+                final String interfaceName = offloadEngineInfo.mInterfaceName;
+                if (!targetInterfaceName.equals(interfaceName)
+                        || ((offloadEngineInfo.mOffloadType
+                        & offloadServiceInfo.getOffloadType()) == 0)) {
+                    continue;
+                }
+                try {
+                    if (isRemove) {
+                        mOffloadEngines.getBroadcastItem(i).onOffloadServiceRemoved(
+                                offloadServiceInfo);
+                    } else {
+                        mOffloadEngines.getBroadcastItem(i).onOffloadServiceUpdated(
+                                offloadServiceInfo);
+                    }
+                } catch (RemoteException e) {
+                    // Can happen in regular cases, do not log a stacktrace
+                    Log.i(TAG, "Failed to send offload callback, remote died", e);
+                }
+            }
+        } finally {
+            mOffloadEngines.finishBroadcast();
+        }
+    }
+
     private class AdvertiserCallback implements MdnsAdvertiser.AdvertiserCallback {
+        // TODO: add a callback to notify when a service is being added on each interface (as soon
+        // as probing starts), and call mOffloadCallbacks. This callback is for
+        // OFFLOAD_CAPABILITY_FILTER_REPLIES offload type.
+
         @Override
         public void onRegisterServiceSucceeded(int transactionId, NsdServiceInfo registeredInfo) {
             mServiceLogs.log("onRegisterServiceSucceeded: transactionId " + transactionId);
@@ -1801,6 +1871,18 @@
                     request.calculateRequestDurationMs());
         }
 
+        @Override
+        public void onOffloadStartOrUpdate(@NonNull String interfaceName,
+                @NonNull OffloadServiceInfo offloadServiceInfo) {
+            sendOffloadServiceInfosUpdate(interfaceName, offloadServiceInfo, false /* isRemove */);
+        }
+
+        @Override
+        public void onOffloadStop(@NonNull String interfaceName,
+                @NonNull OffloadServiceInfo offloadServiceInfo) {
+            sendOffloadServiceInfosUpdate(interfaceName, offloadServiceInfo, true /* isRemove */);
+        }
+
         private ClientInfo getClientInfoOrLog(int transactionId) {
             final ClientInfo clientInfo = mTransactionIdToClientInfoMap.get(transactionId);
             if (clientInfo == null) {
@@ -1920,6 +2002,32 @@
         public void binderDied() {
             mNsdStateMachine.sendMessage(
                     mNsdStateMachine.obtainMessage(NsdManager.UNREGISTER_CLIENT, this));
+
+        }
+
+        @Override
+        public void registerOffloadEngine(String ifaceName, IOffloadEngine cb,
+                @OffloadEngine.OffloadCapability long offloadCapabilities,
+                @OffloadEngine.OffloadType long offloadTypes) {
+            // TODO: Relax the permission because NETWORK_SETTINGS is a signature permission, and
+            //  it may not be possible for all the callers of this API to have it.
+            PermissionUtils.enforceNetworkStackPermissionOr(mContext, NETWORK_SETTINGS);
+            Objects.requireNonNull(ifaceName);
+            Objects.requireNonNull(cb);
+            mNsdStateMachine.sendMessage(
+                    mNsdStateMachine.obtainMessage(NsdManager.REGISTER_OFFLOAD_ENGINE,
+                            new OffloadEngineInfo(cb, ifaceName, offloadCapabilities,
+                                    offloadTypes)));
+        }
+
+        @Override
+        public void unregisterOffloadEngine(IOffloadEngine cb) {
+            // TODO: Relax the permission because NETWORK_SETTINGS is a signature permission, and
+            //  it may not be possible for all the callers of this API to have it.
+            PermissionUtils.enforceNetworkStackPermissionOr(mContext, NETWORK_SETTINGS);
+            Objects.requireNonNull(cb);
+            mNsdStateMachine.sendMessage(
+                    mNsdStateMachine.obtainMessage(NsdManager.UNREGISTER_OFFLOAD_ENGINE, cb));
         }
     }
 
@@ -2003,25 +2111,41 @@
             return IFACE_IDX_ANY;
         }
 
-        final ConnectivityManager cm = mContext.getSystemService(ConnectivityManager.class);
-        if (cm == null) {
-            Log.wtf(TAG, "No ConnectivityManager for resolveService");
+        String interfaceName = getNetworkInterfaceName(network);
+        if (interfaceName == null) {
             return IFACE_IDX_ANY;
         }
-        final LinkProperties lp = cm.getLinkProperties(network);
-        if (lp == null) return IFACE_IDX_ANY;
+        return getNetworkInterfaceIndexByName(interfaceName);
+    }
 
+    private String getNetworkInterfaceName(@Nullable Network network) {
+        if (network == null) {
+            return null;
+        }
+        final ConnectivityManager cm = mContext.getSystemService(ConnectivityManager.class);
+        if (cm == null) {
+            Log.wtf(TAG, "No ConnectivityManager");
+            return null;
+        }
+        final LinkProperties lp = cm.getLinkProperties(network);
+        if (lp == null) {
+            return null;
+        }
         // Only resolve on non-stacked interfaces
+        return lp.getInterfaceName();
+    }
+
+    private int getNetworkInterfaceIndexByName(final String ifaceName) {
         final NetworkInterface iface;
         try {
-            iface = NetworkInterface.getByName(lp.getInterfaceName());
+            iface = NetworkInterface.getByName(ifaceName);
         } catch (SocketException e) {
             Log.e(TAG, "Error querying interface", e);
             return IFACE_IDX_ANY;
         }
 
         if (iface == null) {
-            Log.e(TAG, "Interface not found: " + lp.getInterfaceName());
+            Log.e(TAG, "Interface not found: " + ifaceName);
             return IFACE_IDX_ANY;
         }
 
diff --git a/service-t/src/com/android/server/connectivity/mdns/MdnsAdvertiser.java b/service-t/src/com/android/server/connectivity/mdns/MdnsAdvertiser.java
index 158d7a3..1bc059d 100644
--- a/service-t/src/com/android/server/connectivity/mdns/MdnsAdvertiser.java
+++ b/service-t/src/com/android/server/connectivity/mdns/MdnsAdvertiser.java
@@ -24,15 +24,19 @@
 import android.net.Network;
 import android.net.nsd.NsdManager;
 import android.net.nsd.NsdServiceInfo;
+import android.net.nsd.OffloadEngine;
+import android.net.nsd.OffloadServiceInfo;
 import android.os.Looper;
 import android.util.ArrayMap;
 import android.util.Log;
 import android.util.SparseArray;
 
 import com.android.internal.annotations.VisibleForTesting;
+import com.android.net.module.util.CollectionUtils;
 import com.android.net.module.util.SharedLog;
 import com.android.server.connectivity.mdns.util.MdnsUtils;
 
+import java.util.ArrayList;
 import java.util.List;
 import java.util.Map;
 import java.util.UUID;
@@ -68,9 +72,10 @@
             new ArrayMap<>();
     private final SparseArray<Registration> mRegistrations = new SparseArray<>();
     private final Dependencies mDeps;
-
     private String[] mDeviceHostName;
     @NonNull private final SharedLog mSharedLog;
+    private final Map<String, List<OffloadServiceInfoWrapper>> mInterfaceOffloadServices =
+            new ArrayMap<>();
 
     /**
      * Dependencies for {@link MdnsAdvertiser}, useful for testing.
@@ -115,18 +120,32 @@
     private final MdnsInterfaceAdvertiser.Callback mInterfaceAdvertiserCb =
             new MdnsInterfaceAdvertiser.Callback() {
         @Override
-        public void onRegisterServiceSucceeded(
+        public void onServiceProbingSucceeded(
                 @NonNull MdnsInterfaceAdvertiser advertiser, int serviceId) {
-            // Wait for all current interfaces to be done probing before notifying of success.
-            if (any(mAllAdvertisers, (k, a) -> a.isProbing(serviceId))) return;
-            // The service may still be unregistered/renamed if a conflict is found on a later added
-            // interface, or if a conflicting announcement/reply is detected (RFC6762 9.)
-
             final Registration registration = mRegistrations.get(serviceId);
             if (registration == null) {
                 Log.wtf(TAG, "Register succeeded for unknown registration");
                 return;
             }
+
+            final String interfaceName = advertiser.getSocketInterfaceName();
+            final List<OffloadServiceInfoWrapper> existingOffloadServiceInfoWrappers =
+                    mInterfaceOffloadServices.computeIfAbsent(
+                            interfaceName, k -> new ArrayList<>());
+            // Remove existing offload services from cache for update.
+            existingOffloadServiceInfoWrappers.removeIf(item -> item.mServiceId == serviceId);
+            final OffloadServiceInfoWrapper newOffloadServiceInfoWrapper = createOffloadService(
+                    serviceId,
+                    registration);
+            existingOffloadServiceInfoWrappers.add(newOffloadServiceInfoWrapper);
+            mCb.onOffloadStartOrUpdate(interfaceName,
+                    newOffloadServiceInfoWrapper.mOffloadServiceInfo);
+
+            // Wait for all current interfaces to be done probing before notifying of success.
+            if (any(mAllAdvertisers, (k, a) -> a.isProbing(serviceId))) return;
+            // The service may still be unregistered/renamed if a conflict is found on a later added
+            // interface, or if a conflicting announcement/reply is detected (RFC6762 9.)
+
             if (!registration.mNotifiedRegistrationSuccess) {
                 mCb.onRegisterServiceSucceeded(serviceId, registration.getServiceInfo());
                 registration.mNotifiedRegistrationSuccess = true;
@@ -148,7 +167,12 @@
                 registration.mNotifiedRegistrationSuccess = false;
 
                 // The service was done probing, just reset it to probing state (RFC6762 9.)
-                forAllAdvertisers(a -> a.restartProbingForConflict(serviceId));
+                forAllAdvertisers(a -> {
+                    if (!a.maybeRestartProbingForConflict(serviceId)) {
+                        return;
+                    }
+                    maybeSendOffloadStop(a.getSocketInterfaceName(), serviceId);
+                });
                 return;
             }
 
@@ -196,6 +220,22 @@
         registration.updateForConflict(newInfo, renameCount);
     }
 
+    private void maybeSendOffloadStop(final String interfaceName, int serviceId) {
+        final List<OffloadServiceInfoWrapper> existingOffloadServiceInfoWrappers =
+                mInterfaceOffloadServices.get(interfaceName);
+        if (existingOffloadServiceInfoWrappers == null) {
+            return;
+        }
+        // Stop the offloaded service by matching the service id
+        int idx = CollectionUtils.indexOf(existingOffloadServiceInfoWrappers,
+                item -> item.mServiceId == serviceId);
+        if (idx >= 0) {
+            mCb.onOffloadStop(interfaceName,
+                    existingOffloadServiceInfoWrappers.get(idx).mOffloadServiceInfo);
+            existingOffloadServiceInfoWrappers.remove(idx);
+        }
+    }
+
     /**
      * A request for a {@link MdnsInterfaceAdvertiser}.
      *
@@ -221,7 +261,22 @@
          * @return true if this {@link InterfaceAdvertiserRequest} should now be deleted.
          */
         boolean onAdvertiserDestroyed(@NonNull MdnsInterfaceSocket socket) {
-            mAdvertisers.remove(socket);
+            final MdnsInterfaceAdvertiser removedAdvertiser = mAdvertisers.remove(socket);
+            if (removedAdvertiser != null) {
+                final String interfaceName = removedAdvertiser.getSocketInterfaceName();
+                // If the interface is destroyed, stop all hardware offloading on that interface.
+                final List<OffloadServiceInfoWrapper> offloadServiceInfoWrappers =
+                        mInterfaceOffloadServices.remove(
+                                interfaceName);
+                if (offloadServiceInfoWrappers != null) {
+                    for (OffloadServiceInfoWrapper offloadServiceInfoWrapper :
+                            offloadServiceInfoWrappers) {
+                        mCb.onOffloadStop(interfaceName,
+                                offloadServiceInfoWrapper.mOffloadServiceInfo);
+                    }
+                }
+            }
+
             if (mAdvertisers.size() == 0 && mPendingRegistrations.size() == 0) {
                 // No advertiser is using sockets from this request anymore (in particular for exit
                 // announcements), and there is no registration so newer sockets will not be
@@ -282,7 +337,10 @@
         void removeService(int id) {
             mPendingRegistrations.remove(id);
             for (int i = 0; i < mAdvertisers.size(); i++) {
-                mAdvertisers.valueAt(i).removeService(id);
+                final MdnsInterfaceAdvertiser advertiser = mAdvertisers.valueAt(i);
+                advertiser.removeService(id);
+
+                maybeSendOffloadStop(advertiser.getSocketInterfaceName(), id);
             }
         }
 
@@ -325,6 +383,16 @@
         }
     }
 
+    private static class OffloadServiceInfoWrapper {
+        private final @NonNull OffloadServiceInfo mOffloadServiceInfo;
+        private final int mServiceId;
+
+        OffloadServiceInfoWrapper(int serviceId, OffloadServiceInfo offloadServiceInfo) {
+            mOffloadServiceInfo = offloadServiceInfo;
+            mServiceId = serviceId;
+        }
+    }
+
     private static class Registration {
         @NonNull
         final String mOriginalName;
@@ -425,6 +493,24 @@
 
         // Unregistration is notified immediately as success in NsdService so no callback is needed
         // here.
+
+        /**
+         * Called when a service is ready to be sent for hardware offloading.
+         *
+         * @param interfaceName the interface for sending the update to.
+         * @param offloadServiceInfo the offloading content.
+         */
+        void onOffloadStartOrUpdate(@NonNull String interfaceName,
+                @NonNull OffloadServiceInfo offloadServiceInfo);
+
+        /**
+         * Called when a service is removed or the MdnsInterfaceAdvertiser is destroyed.
+         *
+         * @param interfaceName the interface for sending the update to.
+         * @param offloadServiceInfo the offloading content.
+         */
+        void onOffloadStop(@NonNull String interfaceName,
+                @NonNull OffloadServiceInfo offloadServiceInfo);
     }
 
     public MdnsAdvertiser(@NonNull Looper looper, @NonNull MdnsSocketProvider socketProvider,
@@ -525,4 +611,28 @@
             return false;
         });
     }
+
+    private OffloadServiceInfoWrapper createOffloadService(int serviceId,
+            @NonNull Registration registration) {
+        final NsdServiceInfo nsdServiceInfo = registration.getServiceInfo();
+        List<String> subTypes = new ArrayList<>();
+        String subType = registration.getSubtype();
+        if (subType != null) {
+            subTypes.add(subType);
+        }
+        final OffloadServiceInfo offloadServiceInfo = new OffloadServiceInfo(
+                new OffloadServiceInfo.Key(nsdServiceInfo.getServiceName(),
+                        nsdServiceInfo.getServiceType()),
+                subTypes,
+                String.join(".", mDeviceHostName),
+                null /* rawOffloadPacket */,
+                // TODO: define overlayable resources in
+                // ServiceConnectivityResources that set the priority based on
+                // service type.
+                0 /* priority */,
+                // TODO: set the offloadType based on the callback timing.
+                OffloadEngine.OFFLOAD_TYPE_REPLY);
+        return new OffloadServiceInfoWrapper(serviceId, offloadServiceInfo);
+    }
+
 }
diff --git a/service-t/src/com/android/server/connectivity/mdns/MdnsInterfaceAdvertiser.java b/service-t/src/com/android/server/connectivity/mdns/MdnsInterfaceAdvertiser.java
index 724a704..c5177b7 100644
--- a/service-t/src/com/android/server/connectivity/mdns/MdnsInterfaceAdvertiser.java
+++ b/service-t/src/com/android/server/connectivity/mdns/MdnsInterfaceAdvertiser.java
@@ -73,7 +73,7 @@
         /**
          * Called by the advertiser after it successfully registered a service, after probing.
          */
-        void onRegisterServiceSucceeded(@NonNull MdnsInterfaceAdvertiser advertiser, int serviceId);
+        void onServiceProbingSucceeded(@NonNull MdnsInterfaceAdvertiser advertiser, int serviceId);
 
         /**
          * Called by the advertiser when a conflict was found, during or after probing.
@@ -101,7 +101,7 @@
         public void onFinished(MdnsProber.ProbingInfo info) {
             final MdnsAnnouncer.AnnouncementInfo announcementInfo;
             mSharedLog.i("Probing finished for service " + info.getServiceId());
-            mCbHandler.post(() -> mCb.onRegisterServiceSucceeded(
+            mCbHandler.post(() -> mCb.onServiceProbingSucceeded(
                     MdnsInterfaceAdvertiser.this, info.getServiceId()));
             try {
                 announcementInfo = mRecordRepository.onProbingSucceeded(info);
@@ -282,11 +282,12 @@
     /**
      * Reset a service to the probing state due to a conflict found on the network.
      */
-    public void restartProbingForConflict(int serviceId) {
+    public boolean maybeRestartProbingForConflict(int serviceId) {
         final MdnsProber.ProbingInfo probingInfo = mRecordRepository.setServiceProbing(serviceId);
-        if (probingInfo == null) return;
+        if (probingInfo == null) return false;
 
         mProber.restartForConflict(probingInfo);
+        return true;
     }
 
     /**
@@ -346,4 +347,8 @@
         if (answers == null) return;
         mReplySender.queueReply(answers);
     }
+
+    public String getSocketInterfaceName() {
+        return mSocket.getInterface().getName();
+    }
 }