2 files changed, 44 insertions, 10 deletions

diff --git a/services/core/java/com/android/server/media/MediaKeyDispatcher.java b/services/core/java/com/android/server/media/MediaKeyDispatcher.java
index 3a29622556cc..176ec3f003ea 100644
--- a/services/core/java/com/android/server/media/MediaKeyDispatcher.java
+++ b/services/core/java/com/android/server/media/MediaKeyDispatcher.java
@@ -19,6 +19,7 @@ package com.android.server.media;
 import android.annotation.IntDef;
 import android.annotation.NonNull;
 import android.annotation.Nullable;
+import android.app.PendingIntent;
 import android.media.session.ISessionManager;
 import android.media.session.MediaSession;
 import android.os.Binder;
@@ -76,6 +77,9 @@ public abstract class MediaKeyDispatcher {
     /**
      * Implement this to customize the logic for which MediaSession should consume which key event.
      *
+     * Note: This session will have greater priority over the {@link PendingIntent} returned from
+     * {@link #getMediaButtonReceiver(KeyEvent, int, boolean)}.
+     *
      * @param keyEvent a non-null KeyEvent whose key code is one of the supported media buttons.
      * @param uid the uid value retrieved by calling {@link Binder#getCallingUid()} from
      *         {@link ISessionManager#dispatchMediaKeyEvent(String, boolean, KeyEvent, boolean)}
@@ -84,7 +88,27 @@ public abstract class MediaKeyDispatcher {
      * @return a {@link MediaSession.Token} instance that should consume the given key event.
      */
     @Nullable
-    MediaSession.Token getSessionForKeyEvent(@NonNull KeyEvent keyEvent, int uid,
+    MediaSession.Token getMediaSession(@NonNull KeyEvent keyEvent, int uid,
+            boolean asSystemService) {
+        return null;
+    }
+
+    /**
+     * Implement this to customize the logic for which MediaButtonReceiver should consume a
+     * dispatched key event.
+     *
+     * Note: This pending intent will have lower priority over the {@link MediaSession.Token}
+     * returned from {@link #getMediaSession(KeyEvent, int, boolean)}.
+     *
+     * @param keyEvent a non-null KeyEvent whose key code is one of the supported media buttons.
+     * @param uid the uid value retrieved by calling {@link Binder#getCallingUid()} from
+     *         {@link ISessionManager#dispatchMediaKeyEvent(String, boolean, KeyEvent, boolean)}
+     * @param asSystemService {@code true} if the event came from the system service via hardware
+     *         devices. {@code false} if the event came from the app process through key injection.
+     * @return a {@link PendingIntent} instance that should receive the dispatched key event.
+     */
+    @Nullable
+    PendingIntent getMediaButtonReceiver(@NonNull KeyEvent keyEvent, int uid,
             boolean asSystemService) {
         return null;
     }
diff --git a/services/core/java/com/android/server/media/MediaSessionService.java b/services/core/java/com/android/server/media/MediaSessionService.java
index f360a4ae7ad8..afae20dcf525 100644
--- a/services/core/java/com/android/server/media/MediaSessionService.java
+++ b/services/core/java/com/android/server/media/MediaSessionService.java
@@ -2404,18 +2404,32 @@ public class MediaSessionService extends SystemService implements Monitor {
                 return;
             }
             MediaSessionRecord session = null;
+            MediaButtonReceiverHolder mediaButtonReceiverHolder = null;
 
-            // Retrieve custom session for key event if it exists.
             if (mCustomMediaKeyDispatcher != null) {
-                MediaSession.Token token = mCustomMediaKeyDispatcher.getSessionForKeyEvent(
-                        keyEvent, uid, asSystemService);
+                MediaSession.Token token = mCustomMediaKeyDispatcher.getMediaSession(keyEvent, uid,
+                        asSystemService);
                 if (token != null) {
                     session = getMediaSessionRecordLocked(token);
                 }
+
+                if (session == null) {
+                    PendingIntent pi = mCustomMediaKeyDispatcher.getMediaButtonReceiver(keyEvent,
+                            uid, asSystemService);
+                    if (pi != null) {
+                        mediaButtonReceiverHolder = MediaButtonReceiverHolder.create(mContext,
+                                mCurrentFullUserRecord.mFullUserId, pi);
+                    }
+                }
             }
 
-            if (session == null) {
+            if (session == null && mediaButtonReceiverHolder == null) {
                 session = (MediaSessionRecord) mCurrentFullUserRecord.getMediaButtonSessionLocked();
+
+                if (session == null) {
+                    mediaButtonReceiverHolder =
+                            mCurrentFullUserRecord.mLastMediaButtonReceiverHolder;
+                }
             }
 
             if (session != null) {
@@ -2438,16 +2452,12 @@ public class MediaSessionService extends SystemService implements Monitor {
                 } catch (RemoteException e) {
                     Log.w(TAG, "Failed to send callback", e);
                 }
-            } else if (mCurrentFullUserRecord.mLastMediaButtonReceiverHolder != null) {
+            } else if (mediaButtonReceiverHolder != null) {
                 if (needWakeLock) {
                     mKeyEventReceiver.acquireWakeLockLocked();
                 }
                 String callingPackageName =
                         (asSystemService) ? mContext.getPackageName() : packageName;
-
-                MediaButtonReceiverHolder mediaButtonReceiverHolder =
-                        mCurrentFullUserRecord.mLastMediaButtonReceiverHolder;
-
                 boolean sent = mediaButtonReceiverHolder.send(
                         mContext, keyEvent, callingPackageName,
                         needWakeLock ? mKeyEventReceiver.mLastTimeoutId : -1, mKeyEventReceiver,