MediaSession2Service: Initial commit

Bug: 122563346
Test: Build
Change-Id: I250ee493837bfa7964fa7baf3d11f1673c879010
diff --git a/Android.bp b/Android.bp
index c202408..2808c35 100644
--- a/Android.bp
+++ b/Android.bp
@@ -476,6 +476,7 @@
         "media/java/android/media/IMediaRouterClient.aidl",
         "media/java/android/media/IMediaRouterService.aidl",
         "media/java/android/media/IMediaSession2.aidl",
+        "media/java/android/media/IMediaSession2Service.aidl",
         "media/java/android/media/IMediaScannerListener.aidl",
         "media/java/android/media/IMediaScannerService.aidl",
         "media/java/android/media/IPlaybackConfigDispatcher.aidl",
diff --git a/media/java/android/media/IMediaSession2Service.aidl b/media/java/android/media/IMediaSession2Service.aidl
new file mode 100644
index 0000000..10ac1be
--- /dev/null
+++ b/media/java/android/media/IMediaSession2Service.aidl
@@ -0,0 +1,32 @@
+/*
+ * Copyright 2019 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package android.media;
+
+import android.os.Bundle;
+import android.media.Controller2Link;
+
+/**
+ * Interface from MediaController2 to MediaSession2Service.
+ * <p>
+ * Keep this interface oneway. Otherwise a malicious app may implement fake version of this,
+ * and holds calls from controller to make controller owner(s) frozen.
+ * @hide
+ */
+oneway interface IMediaSession2Service {
+    void connect(in Controller2Link caller, int seq, in Bundle connectionRequest) = 0;
+    // Next Id : 1
+}
diff --git a/media/java/android/media/MediaController2.java b/media/java/android/media/MediaController2.java
index 774ea18..165ea41 100644
--- a/media/java/android/media/MediaController2.java
+++ b/media/java/android/media/MediaController2.java
@@ -26,12 +26,16 @@
 
 import android.annotation.NonNull;
 import android.annotation.Nullable;
+import android.content.ComponentName;
 import android.content.Context;
+import android.content.Intent;
+import android.content.ServiceConnection;
 import android.os.Binder;
 import android.os.Bundle;
 import android.os.Handler;
 import android.os.IBinder;
 import android.os.Process;
+import android.os.RemoteException;
 import android.os.ResultReceiver;
 import android.util.ArrayMap;
 import android.util.ArraySet;
@@ -63,6 +67,7 @@
     private final Executor mCallbackExecutor;
     private final Controller2Link mControllerStub;
     private final Handler mResultHandler;
+    private final SessionServiceConnection mServiceConnection;
 
     private final Object mLock = new Object();
     //@GuardedBy("mLock")
@@ -118,16 +123,25 @@
         mPendingCommands = new ArrayMap<>();
         mRequestedCommandSeqNumbers = new ArraySet<>();
 
+        boolean connectRequested;
         if (token.getType() == TYPE_SESSION) {
-            connectToSession();
+            mServiceConnection = null;
+            connectRequested = requestConnectToSession();
         } else {
-            // TODO: Handle connect to session service.
+            mServiceConnection = new SessionServiceConnection();
+            connectRequested = requestConnectToService();
+        }
+        if (!connectRequested) {
+            close();
         }
     }
 
     @Override
     public void close() {
         synchronized (mLock) {
+            if (mServiceConnection != null) {
+                mContext.unbindService(mServiceConnection);
+            }
             if (mSessionBinder != null) {
                 try {
                     mSessionBinder.disconnect(mControllerStub, getNextSeqNumber());
@@ -299,18 +313,55 @@
         }
     }
 
-    private void connectToSession() {
-        Session2Link sessionBinder = mSessionToken.getSessionLink();
+    private Bundle createConnectionRequest() {
         Bundle connectionRequest = new Bundle();
         connectionRequest.putString(KEY_PACKAGE_NAME, mContext.getPackageName());
         connectionRequest.putInt(KEY_PID, Process.myPid());
+        return connectionRequest;
+    }
 
+    private boolean requestConnectToSession() {
+        Session2Link sessionBinder = mSessionToken.getSessionLink();
+        Bundle connectionRequest = createConnectionRequest();
         try {
             sessionBinder.connect(mControllerStub, getNextSeqNumber(), connectionRequest);
         } catch (RuntimeException e) {
-            Log.w(TAG, "Failed to call connection request. Framework will retry"
-                    + " automatically");
+            Log.w(TAG, "Failed to call connection request", e);
+            return false;
         }
+        return true;
+    }
+
+    private boolean requestConnectToService() {
+        // Service. Needs to get fresh binder whenever connection is needed.
+        final Intent intent = new Intent(MediaSession2Service.SERVICE_INTERFACE);
+        intent.setClassName(mSessionToken.getPackageName(), mSessionToken.getServiceName());
+
+        // Use bindService() instead of startForegroundService() to start session service for three
+        // reasons.
+        // 1. Prevent session service owner's stopSelf() from destroying service.
+        //    With the startForegroundService(), service's call of stopSelf() will trigger immediate
+        //    onDestroy() calls on the main thread even when onConnect() is running in another
+        //    thread.
+        // 2. Minimize APIs for developers to take care about.
+        //    With bindService(), developers only need to take care about Service.onBind()
+        //    but Service.onStartCommand() should be also taken care about with the
+        //    startForegroundService().
+        // 3. Future support for UI-less playback
+        //    If a service wants to keep running, it should be either foreground service or
+        //    bound service. But there had been request for the feature for system apps
+        //    and using bindService() will be better fit with it.
+        synchronized (mLock) {
+            boolean result = mContext.bindService(
+                    intent, mServiceConnection, Context.BIND_AUTO_CREATE);
+            if (!result) {
+                Log.w(TAG, "bind to " + mSessionToken + " failed");
+                return false;
+            } else if (DEBUG) {
+                Log.d(TAG, "bind to " + mSessionToken + " succeeded");
+            }
+        }
+        return true;
     }
 
     /**
@@ -367,4 +418,59 @@
         public void onCommandResult(@NonNull MediaController2 controller, @NonNull Object token,
                 @NonNull Session2Command command, @NonNull Session2Command.Result result) {}
     }
+
+    // This will be called on the main thread.
+    private class SessionServiceConnection implements ServiceConnection {
+        SessionServiceConnection() {
+        }
+
+        @Override
+        public void onServiceConnected(ComponentName name, IBinder service) {
+            // Note that it's always main-thread.
+            boolean connectRequested = false;
+            try {
+                if (DEBUG) {
+                    Log.d(TAG, "onServiceConnected " + name + " " + this);
+                }
+                // Sanity check
+                if (!mSessionToken.getPackageName().equals(name.getPackageName())) {
+                    Log.wtf(TAG, "Expected connection to " + mSessionToken.getPackageName()
+                            + " but is connected to " + name);
+                    return;
+                }
+                IMediaSession2Service iService = IMediaSession2Service.Stub.asInterface(service);
+                if (iService == null) {
+                    Log.wtf(TAG, "Service interface is missing.");
+                    return;
+                }
+                Bundle connectionRequest = createConnectionRequest();
+                iService.connect(mControllerStub, getNextSeqNumber(), connectionRequest);
+                connectRequested = true;
+            } catch (RemoteException e) {
+                Log.w(TAG, "Service " + name + " has died prematurely", e);
+            } finally {
+                if (!connectRequested) {
+                    close();
+                }
+            }
+        }
+
+        @Override
+        public void onServiceDisconnected(ComponentName name) {
+            // Temporal lose of the binding because of the service crash. System will automatically
+            // rebind, so just no-op.
+            if (DEBUG) {
+                Log.w(TAG, "Session service " + name + " is disconnected.");
+            }
+            close();
+        }
+
+        @Override
+        public void onBindingDied(ComponentName name) {
+            // Permanent lose of the binding because of the service package update or removed.
+            // This SessionServiceRecord will be removed accordingly, but forget session binder here
+            // for sure.
+            close();
+        }
+    }
 }
diff --git a/media/java/android/media/MediaSession2.java b/media/java/android/media/MediaSession2.java
index e008adf..dceef34 100644
--- a/media/java/android/media/MediaSession2.java
+++ b/media/java/android/media/MediaSession2.java
@@ -31,7 +31,6 @@
 import android.content.Intent;
 import android.media.session.MediaSessionManager;
 import android.media.session.MediaSessionManager.RemoteUserInfo;
-import android.os.Binder;
 import android.os.Bundle;
 import android.os.Handler;
 import android.os.Process;
@@ -41,7 +40,6 @@
 import android.util.Log;
 
 import java.util.ArrayList;
-import java.util.Collection;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
@@ -117,15 +115,19 @@
     @Override
     public void close() {
         try {
+            List<ControllerInfo> controllerInfos;
+            synchronized (mLock) {
+                if (mClosed) {
+                    return;
+                }
+                mClosed = true;
+                controllerInfos = getConnectedControllers();
+                mConnectedControllers.clear();
+                mCallback.onSessionClosed(this);
+            }
             synchronized (MediaSession2.class) {
                 SESSION_ID_LIST.remove(mSessionId);
             }
-            Collection<ControllerInfo> controllerInfos;
-            synchronized (mLock) {
-                controllerInfos = mConnectedControllers.values();
-                mConnectedControllers.clear();
-                mClosed = true;
-            }
             for (ControllerInfo info : controllerInfos) {
                 info.notifyDisconnected();
             }
@@ -160,10 +162,7 @@
         if (command == null) {
             throw new IllegalArgumentException("command shouldn't be null");
         }
-        Collection<ControllerInfo> controllerInfos;
-        synchronized (mLock) {
-            controllerInfos = mConnectedControllers.values();
-        }
+        List<ControllerInfo> controllerInfos = getConnectedControllers();
         for (ControllerInfo controller : controllerInfos) {
             controller.sendSessionCommand(command, args, null);
         }
@@ -222,23 +221,26 @@
         }
     }
 
-    // Called by Session2Link.onConnect
-    void onConnect(final Controller2Link controller, int seq, Bundle connectionRequest) {
-        if (controller == null || connectionRequest == null) {
-            return;
+    SessionCallback getCallback() {
+        return mCallback;
+    }
+
+    // Called by Session2Link.onConnect and MediaSession2Service.MediaSession2ServiceStub.connect
+    void onConnect(final Controller2Link controller, int callingPid, int callingUid, int seq,
+            Bundle connectionRequest) {
+        if (callingPid == 0) {
+            // The pid here is from Binder.getCallingPid(), which can be 0 for an oneway call from
+            // the remote process. If it's the case, use PID from the connectionRequest.
+            callingPid = connectionRequest.getInt(KEY_PID);
         }
-        final int uid = Binder.getCallingUid();
-        final int callingPid = Binder.getCallingPid();
-        final long token = Binder.clearCallingIdentity();
-        // Binder.getCallingPid() can be 0 for an oneway call from the remote process.
-        // If it's the case, use PID from the ConnectionRequest.
-        final int pid = (callingPid != 0) ? callingPid : connectionRequest.getInt(KEY_PID);
-        final String pkg = connectionRequest.getString(KEY_PACKAGE_NAME);
-        try {
-            RemoteUserInfo remoteUserInfo = new RemoteUserInfo(pkg, pid, uid);
-            final ControllerInfo controllerInfo = new ControllerInfo(remoteUserInfo,
-                    mSessionManager.isTrustedForMediaControl(remoteUserInfo), controller);
-            mCallbackExecutor.execute(() -> {
+        String callingPkg = connectionRequest.getString(KEY_PACKAGE_NAME);
+
+        RemoteUserInfo remoteUserInfo = new RemoteUserInfo(callingPkg, callingPid, callingUid);
+        final ControllerInfo controllerInfo = new ControllerInfo(remoteUserInfo,
+                mSessionManager.isTrustedForMediaControl(remoteUserInfo), controller);
+        mCallbackExecutor.execute(() -> {
+            boolean accept = false;
+            try {
                 if (isClosed()) {
                     return;
                 }
@@ -247,77 +249,67 @@
                 // Don't reject connection for the request from trusted app.
                 // Otherwise server will fail to retrieve session's information to dispatch
                 // media keys to.
-                boolean accept =
-                        controllerInfo.mAllowedCommands != null || controllerInfo.isTrusted();
-                if (accept) {
-                    if (controllerInfo.mAllowedCommands == null) {
-                        // For trusted apps, send non-null allowed commands to keep
-                        // connection.
-                        controllerInfo.mAllowedCommands =
-                                new Session2CommandGroup.Builder().build();
+                accept = controllerInfo.mAllowedCommands != null || controllerInfo.isTrusted();
+                if (!accept) {
+                    return;
+                }
+                if (controllerInfo.mAllowedCommands == null) {
+                    // For trusted apps, send non-null allowed commands to keep
+                    // connection.
+                    controllerInfo.mAllowedCommands =
+                            new Session2CommandGroup.Builder().build();
+                }
+                if (DEBUG) {
+                    Log.d(TAG, "Accepting connection: " + controllerInfo);
+                }
+                synchronized (mLock) {
+                    if (mConnectedControllers.containsKey(controller)) {
+                        Log.w(TAG, "Controller " + controllerInfo + " has sent connection"
+                                + " request multiple times");
                     }
-                    if (DEBUG) {
-                        Log.d(TAG, "Accepting connection: " + controllerInfo);
-                    }
-                    synchronized (mLock) {
-                        if (mConnectedControllers.containsKey(controller)) {
-                            Log.w(TAG, "Controller " + controllerInfo + " has sent connection"
-                                    + " request multiple times");
-                        }
-                        mConnectedControllers.put(controller, controllerInfo);
-                    }
-                    // If connection is accepted, notify the current state to the controller.
-                    // It's needed because we cannot call synchronous calls between
-                    // session/controller.
-                    Bundle connectionResult = new Bundle();
-                    connectionResult.putParcelable(KEY_SESSION2LINK, mSessionStub);
-                    connectionResult.putParcelable(KEY_ALLOWED_COMMANDS,
-                            controllerInfo.mAllowedCommands);
+                    mConnectedControllers.put(controller, controllerInfo);
+                }
+                // If connection is accepted, notify the current state to the controller.
+                // It's needed because we cannot call synchronous calls between
+                // session/controller.
+                Bundle connectionResult = new Bundle();
+                connectionResult.putParcelable(KEY_SESSION2LINK, mSessionStub);
+                connectionResult.putParcelable(KEY_ALLOWED_COMMANDS,
+                        controllerInfo.mAllowedCommands);
 
-                    // Double check if session is still there, because close() can be called in
-                    // another thread.
-                    if (isClosed()) {
-                        return;
-                    }
-                    controllerInfo.notifyConnected(connectionResult);
-                } else {
+                // Double check if session is still there, because close() can be called in
+                // another thread.
+                if (isClosed()) {
+                    return;
+                }
+                controllerInfo.notifyConnected(connectionResult);
+            } finally {
+                if (!accept) {
                     if (DEBUG) {
                         Log.d(TAG, "Rejecting connection, controllerInfo=" + controllerInfo);
                     }
-                    controllerInfo.notifyDisconnected();
                 }
-            });
-        } finally {
-            Binder.restoreCallingIdentity(token);
-        }
+                controllerInfo.notifyDisconnected();
+            }
+        });
     }
 
     // Called by Session2Link.onDisconnect
-    void onDisconnect(final Controller2Link controller, int seq) {
-        if (controller == null) {
-            return;
-        }
+    void onDisconnect(@NonNull final Controller2Link controller, int seq) {
         final ControllerInfo controllerInfo;
         synchronized (mLock) {
-            controllerInfo = mConnectedControllers.get(controller);
+            controllerInfo = mConnectedControllers.remove(controller);
         }
         if (controllerInfo == null) {
             return;
         }
-
-        final long token = Binder.clearCallingIdentity();
-        try {
-            mCallbackExecutor.execute(() -> {
-                mCallback.onDisconnected(MediaSession2.this, controllerInfo);
-            });
-            mConnectedControllers.remove(controller);
-        } finally {
-            Binder.restoreCallingIdentity(token);
-        }
+        mCallbackExecutor.execute(() -> {
+            mCallback.onDisconnected(MediaSession2.this, controllerInfo);
+        });
     }
 
     // Called by Session2Link.onSessionCommand
-    void onSessionCommand(final Controller2Link controller, final int seq,
+    void onSessionCommand(@NonNull final Controller2Link controller, final int seq,
             final Session2Command command, final Bundle args,
             @Nullable ResultReceiver resultReceiver) {
         if (controller == null) {
@@ -332,34 +324,28 @@
         }
 
         // TODO: check allowed commands.
-        final long token = Binder.clearCallingIdentity();
-        try {
-            synchronized (mLock) {
-                controllerInfo.addRequestedCommandSeqNumber(seq);
-            }
-
-            mCallbackExecutor.execute(() -> {
-                if (!controllerInfo.removeRequestedCommandSeqNumber(seq)) {
-                    resultReceiver.send(RESULT_INFO_SKIPPED, null);
-                    return;
-                }
-                Session2Command.Result result = mCallback.onSessionCommand(
-                        MediaSession2.this, controllerInfo, command, args);
-                if (resultReceiver != null) {
-                    if (result == null) {
-                        throw new RuntimeException("onSessionCommand shouldn't return null");
-                    } else {
-                        resultReceiver.send(result.getResultCode(), result.getResultData());
-                    }
-                }
-            });
-        } finally {
-            Binder.restoreCallingIdentity(token);
+        synchronized (mLock) {
+            controllerInfo.addRequestedCommandSeqNumber(seq);
         }
+        mCallbackExecutor.execute(() -> {
+            if (!controllerInfo.removeRequestedCommandSeqNumber(seq)) {
+                resultReceiver.send(RESULT_INFO_SKIPPED, null);
+                return;
+            }
+            Session2Command.Result result = mCallback.onSessionCommand(
+                    MediaSession2.this, controllerInfo, command, args);
+            if (resultReceiver != null) {
+                if (result == null) {
+                    throw new RuntimeException("onSessionCommand shouldn't return null");
+                } else {
+                    resultReceiver.send(result.getResultCode(), result.getResultData());
+                }
+            }
+        });
     }
 
     // Called by Session2Link.onCancelCommand
-    void onCancelCommand(final Controller2Link controller, final int seq) {
+    void onCancelCommand(@NonNull final Controller2Link controller, final int seq) {
         final ControllerInfo controllerInfo;
         synchronized (mLock) {
             controllerInfo = mConnectedControllers.get(controller);
@@ -367,13 +353,15 @@
         if (controllerInfo == null) {
             return;
         }
+        controllerInfo.removeRequestedCommandSeqNumber(seq);
+    }
 
-        final long token = Binder.clearCallingIdentity();
-        try {
-            controllerInfo.removeRequestedCommandSeqNumber(seq);
-        } finally {
-            Binder.restoreCallingIdentity(token);
+    private List<ControllerInfo> getConnectedControllers() {
+        List<ControllerInfo> controllers = new ArrayList<>();
+        synchronized (mLock) {
+            controllers.addAll(mConnectedControllers.values());
         }
+        return controllers;
     }
 
     /**
@@ -660,6 +648,8 @@
      * This API is not generally intended for third party application developers.
      */
     public abstract static class SessionCallback {
+        ForegroundServiceEventCallback mForegroundServiceEventCallback;
+
         /**
          * Called when a controller is created for this session. Return allowed commands for
          * controller. By default it returns {@code null}.
@@ -716,5 +706,19 @@
         public void onCommandResult(@NonNull MediaSession2 session,
                 @NonNull ControllerInfo controller, @NonNull Object token,
                 @NonNull Session2Command command, @NonNull Session2Command.Result result) {}
+
+        final void onSessionClosed(MediaSession2 session) {
+            if (mForegroundServiceEventCallback != null) {
+                mForegroundServiceEventCallback.onSessionClosed(session);
+            }
+        }
+
+        void setForegroundServiceEventCallback(ForegroundServiceEventCallback callback) {
+            mForegroundServiceEventCallback = callback;
+        }
+
+        abstract static class ForegroundServiceEventCallback {
+            public void onSessionClosed(MediaSession2 session) {}
+        }
     }
 }
diff --git a/media/java/android/media/MediaSession2Service.java b/media/java/android/media/MediaSession2Service.java
new file mode 100644
index 0000000..8fb00fe
--- /dev/null
+++ b/media/java/android/media/MediaSession2Service.java
@@ -0,0 +1,288 @@
+/*
+ * Copyright 2019 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package android.media;
+
+import android.annotation.CallSuper;
+import android.annotation.NonNull;
+import android.annotation.Nullable;
+import android.app.Service;
+import android.content.Intent;
+import android.os.Binder;
+import android.os.Bundle;
+import android.os.Handler;
+import android.os.IBinder;
+import android.util.ArrayMap;
+import android.util.Log;
+
+import com.android.internal.annotations.GuardedBy;
+
+import java.lang.ref.WeakReference;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Map;
+
+/**
+ * Service containing {@link MediaSession2}.
+ * <p>
+ * This API is not generally intended for third party application developers.
+ * Use the <a href="{@docRoot}jetpack/androidx.html">AndroidX</a>
+ * <a href="{@docRoot}reference/androidx/media2/package-summary.html">Media2 Library</a>
+ * for consistent behavior across all devices.
+ * @hide
+ */
+// TODO: Unhide
+// TODO: Add onUpdateNotification(), and calls it to get Notification for startForegroundService()
+//       when a session's player state becomes playing.
+public abstract class MediaSession2Service extends Service {
+    /**
+     * The {@link Intent} that must be declared as handled by the service.
+     */
+    public static final String SERVICE_INTERFACE = "android.media.MediaSession2Service";
+
+    private static final String TAG = "MediaSession2Service";
+    private static final boolean DEBUG = Log.isLoggable(TAG, Log.DEBUG);
+
+    private final Object mLock = new Object();
+    @GuardedBy("mLock")
+    private Map<String, MediaSession2> mSessions = new ArrayMap<>();
+
+    private MediaSession2ServiceStub mStub;
+
+    /**
+     * Called by the system when the service is first created. Do not call this method directly.
+     * <p>
+     * Override this method if you need your own initialization. Derived classes MUST call through
+     * to the super class's implementation of this method.
+     */
+    @CallSuper
+    @Override
+    public void onCreate() {
+        super.onCreate();
+        mStub = new MediaSession2ServiceStub(this);
+    }
+
+    @CallSuper
+    @Override
+    @Nullable
+    public IBinder onBind(@NonNull Intent intent) {
+        if (SERVICE_INTERFACE.equals(intent.getAction())) {
+            return mStub;
+        }
+        return null;
+    }
+
+    @CallSuper
+    @Override
+    public int onStartCommand(Intent intent, int flags, int startId) {
+        // TODO: Dispatch media key events to the primary session.
+        return START_STICKY;
+    }
+
+    /**
+     * Called by the system to notify that it is no longer used and is being removed. Do not call
+     * this method directly.
+     * <p>
+     * Override this method if you need your own clean up. Derived classes MUST call through
+     * to the super class's implementation of this method.
+     */
+    @CallSuper
+    @Override
+    public void onDestroy() {
+        super.onDestroy();
+        synchronized (mLock) {
+            for (MediaSession2 session : mSessions.values()) {
+                session.getCallback().setForegroundServiceEventCallback(null);
+            }
+            mSessions.clear();
+        }
+        mStub.close();
+    }
+
+    /**
+     * Called when a {@link MediaController2} is created with the this service's
+     * {@link Session2Token}. Return the primary session for telling the controller which session to
+     * connect.
+     * <p>
+     * Primary session is the highest priority session that this service manages. Here are some
+     * recommendations of the primary session.
+     * <ol>
+     * <li>When there's no {@link MediaSession2}, create and return a new session. Resume the
+     * playback that the app has the lastly played with the new session. The behavior is what
+     * framework expects when the framework sends key events to the service.</li>
+     * <li>When there's multiple {@link MediaSession2}s, pick the session that has the lastly
+     * started the playback. This is the same way as the framework prioritize sessions to receive
+     * media key events.</li>
+     * </ol>
+     * <p>
+     * Session returned here will be added to this service automatically. You don't need to call
+     * {@link #addSession(MediaSession2)} for that.
+     * <p>
+     * Session service will accept or reject the connection with the
+     * {@link MediaSession2.SessionCallback} in the session returned here.
+     * <p>
+     * This method is always called on the main thread.
+     *
+     * @return a new session
+     * @see MediaSession2.Builder
+     * @see #getSessions()
+     */
+    @NonNull
+    public abstract MediaSession2 onGetPrimarySession();
+
+    /**
+     * Adds a session to this service.
+     * <p>
+     * Added session will be removed automatically when it's closed, or removed when
+     * {@link #removeSession} is called.
+     *
+     * @param session a session to be added.
+     * @see #removeSession(MediaSession2)
+     */
+    public final void addSession(@NonNull MediaSession2 session) {
+        if (session == null) {
+            throw new IllegalArgumentException("session shouldn't be null");
+        }
+        if (session.isClosed()) {
+            throw new IllegalArgumentException("session is already closed");
+        }
+        synchronized (mLock) {
+            MediaSession2 previousSession = mSessions.get(session.getSessionId());
+            if (previousSession != session) {
+                if (previousSession != null) {
+                    Log.w(TAG, "Session ID should be unique, ID=" + session.getSessionId()
+                            + ", previous=" + previousSession + ", session=" + session);
+                }
+                return;
+            }
+            mSessions.put(session.getSessionId(), session);
+            session.getCallback().setForegroundServiceEventCallback(
+                    new MediaSession2.SessionCallback.ForegroundServiceEventCallback() {
+                        @Override
+                        public void onSessionClosed(MediaSession2 session) {
+                            removeSession(session);
+                        }
+                    });
+        }
+    }
+
+    /**
+     * Removes a session from this service.
+     *
+     * @param session a session to be removed.
+     * @see #addSession(MediaSession2)
+     */
+    public final void removeSession(@NonNull MediaSession2 session) {
+        if (session == null) {
+            throw new IllegalArgumentException("session shouldn't be null");
+        }
+        synchronized (mLock) {
+            mSessions.remove(session.getSessionId());
+        }
+    }
+
+    /**
+     * Gets the list of {@link MediaSession2}s that you've added to this service.
+     *
+     * @return sessions
+     */
+    public final @NonNull List<MediaSession2> getSessions() {
+        List<MediaSession2> list = new ArrayList<>();
+        synchronized (mLock) {
+            list.addAll(mSessions.values());
+        }
+        return list;
+    }
+
+    private static final class MediaSession2ServiceStub extends IMediaSession2Service.Stub
+            implements AutoCloseable {
+        final WeakReference<MediaSession2Service> mService;
+        final Handler mHandler;
+
+        MediaSession2ServiceStub(MediaSession2Service service) {
+            mService = new WeakReference<>(service);
+            mHandler = new Handler(service.getMainLooper());
+        }
+
+        @Override
+        public void connect(Controller2Link caller, int seq, Bundle connectionRequest) {
+            if (mService.get() == null) {
+                if (DEBUG) {
+                    Log.d(TAG, "Service is already destroyed");
+                }
+                return;
+            }
+            if (caller == null || connectionRequest == null) {
+                if (DEBUG) {
+                    Log.d(TAG, "Ignoring calls with illegal arguments, caller=" + caller
+                            + ", connectionRequest=" + connectionRequest);
+                }
+                return;
+            }
+            final int pid = Binder.getCallingPid();
+            final int uid = Binder.getCallingUid();
+            final long token = Binder.clearCallingIdentity();
+            try {
+                mHandler.post(() -> {
+                    boolean shouldNotifyDisconnected = true;
+                    try {
+                        final MediaSession2Service service = mService.get();
+                        if (service == null) {
+                            if (DEBUG) {
+                                Log.d(TAG, "Service isn't available");
+                            }
+                            return;
+                        }
+                        if (DEBUG) {
+                            Log.d(TAG, "Handling incoming connection request from the"
+                                    + " controller, controller=" + caller + ", uid=" + uid);
+                        }
+                        final MediaSession2 session;
+                        session = service.onGetPrimarySession();
+                        service.addSession(session);
+                        shouldNotifyDisconnected = false;
+                        session.onConnect(caller, pid, uid, seq, connectionRequest);
+                    } catch (Exception e) {
+                        // Don't propagate exception in service to the controller.
+                        Log.w(TAG, "Failed to add a session to session service", e);
+                    } finally {
+                        // Trick to call onDisconnected() in one place.
+                        if (shouldNotifyDisconnected) {
+                            if (DEBUG) {
+                                Log.d(TAG, "Service has destroyed prematurely."
+                                        + " Rejecting connection");
+                            }
+                            try {
+                                caller.notifyDisconnected(0);
+                            } catch (RuntimeException e) {
+                                // Controller may be died prematurely.
+                                // Not an issue because we'll ignore it anyway.
+                            }
+                        }
+                    }
+                });
+            } finally {
+                Binder.restoreCallingIdentity(token);
+            }
+        }
+
+        @Override
+        public void close() {
+            mHandler.removeCallbacksAndMessages(null);
+            mService.clear();
+        }
+    }
+}
diff --git a/media/java/android/media/Session2Link.java b/media/java/android/media/Session2Link.java
index 5fe558d..08664aa 100644
--- a/media/java/android/media/Session2Link.java
+++ b/media/java/android/media/Session2Link.java
@@ -17,6 +17,7 @@
 package android.media;
 
 import android.annotation.NonNull;
+import android.os.Binder;
 import android.os.Bundle;
 import android.os.IBinder;
 import android.os.Parcel;
@@ -145,8 +146,9 @@
     }
 
     /** Stub implementation for IMediaSession2.connect */
-    public void onConnect(final Controller2Link caller, int seq, Bundle connectionRequest) {
-        mSession.onConnect(caller, seq, connectionRequest);
+    public void onConnect(final Controller2Link caller, int pid, int uid, int seq,
+            Bundle connectionRequest) {
+        mSession.onConnect(caller, pid, uid, seq, connectionRequest);
     }
 
     /** Stub implementation for IMediaSession2.disconnect */
@@ -168,23 +170,57 @@
     private class Session2Stub extends IMediaSession2.Stub {
         @Override
         public void connect(final Controller2Link caller, int seq, Bundle connectionRequest) {
-            Session2Link.this.onConnect(caller, seq, connectionRequest);
+            if (caller == null || connectionRequest == null) {
+                return;
+            }
+            final int pid = Binder.getCallingPid();
+            final int uid = Binder.getCallingUid();
+            final long token = Binder.clearCallingIdentity();
+            try {
+                Session2Link.this.onConnect(caller, pid, uid, seq, connectionRequest);
+            } finally {
+                Binder.restoreCallingIdentity(token);
+            }
         }
 
         @Override
         public void disconnect(final Controller2Link caller, int seq) {
-            Session2Link.this.onDisconnect(caller, seq);
+            if (caller == null) {
+                return;
+            }
+            final long token = Binder.clearCallingIdentity();
+            try {
+                Session2Link.this.onDisconnect(caller, seq);
+            } finally {
+                Binder.restoreCallingIdentity(token);
+            }
         }
 
         @Override
         public void sendSessionCommand(final Controller2Link caller, final int seq,
                 final Session2Command command, final Bundle args, ResultReceiver resultReceiver) {
-            Session2Link.this.onSessionCommand(caller, seq, command, args, resultReceiver);
+            if (caller == null) {
+                return;
+            }
+            final long token = Binder.clearCallingIdentity();
+            try {
+                Session2Link.this.onSessionCommand(caller, seq, command, args, resultReceiver);
+            } finally {
+                Binder.restoreCallingIdentity(token);
+            }
         }
 
         @Override
         public void cancelSessionCommand(final Controller2Link caller, final int seq) {
-            Session2Link.this.onCancelCommand(caller, seq);
+            if (caller == null) {
+                return;
+            }
+            final long token = Binder.clearCallingIdentity();
+            try {
+                Session2Link.this.onCancelCommand(caller, seq);
+            } finally {
+                Binder.restoreCallingIdentity(token);
+            }
         }
     }
 }