Add DefaultDrmSessionManagerTest

This uses a license server implemented using MockWebServer to test
DefaultDrmSessionManager and DefaultDrmSession.

PiperOrigin-RevId: 318086890

library/core/src/test/java/com/google/android/exoplayer2/drm/DefaultDrmSessionManagerTest.java

+/*
+ * Copyright (C) 2020 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 com.google.android.exoplayer2.drm;
+
+import static com.google.common.truth.Truth.assertThat;
+
+import android.os.Handler;
+import android.os.HandlerThread;
+import androidx.annotation.Nullable;
+import androidx.test.ext.junit.runners.AndroidJUnit4;
+import com.google.android.exoplayer2.source.MediaSource;
+import com.google.android.exoplayer2.testutil.FakeExoMediaDrm;
+import com.google.android.exoplayer2.testutil.TestUtil;
+import com.google.android.exoplayer2.util.ConditionVariable;
+import com.google.android.exoplayer2.util.Function;
+import com.google.android.exoplayer2.util.MediaSourceEventDispatcher;
+import com.google.android.exoplayer2.util.MimeTypes;
+import com.google.common.collect.ImmutableList;
+import java.util.Map;
+import java.util.UUID;
+import java.util.concurrent.atomic.AtomicReference;
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+
+/** Tests for {@link DefaultDrmSessionManager} and {@link DefaultDrmSession}. */
+// TODO: Test more branches:
+// - Different sources for licenseServerUrl.
+// - Multiple acquisitions & releases for same keys -> multiple requests.
+// - Provisioning.
+// - Key denial.
+@RunWith(AndroidJUnit4.class)
+public class DefaultDrmSessionManagerTest {
+
+  private static final int TIMEOUT_MS = 1_000;
+
+  private static final UUID DRM_SCHEME_UUID =
+      UUID.nameUUIDFromBytes(TestUtil.createByteArray(7, 8, 9));
+  private static final ImmutableList<DrmInitData.SchemeData> DRM_SCHEME_DATAS =
+      ImmutableList.of(
+          new DrmInitData.SchemeData(
+              DRM_SCHEME_UUID, MimeTypes.VIDEO_MP4, /* data= */ TestUtil.createByteArray(1, 2, 3)));
+  private static final DrmInitData DRM_INIT_DATA = new DrmInitData(DRM_SCHEME_DATAS);
+
+  private HandlerThread playbackThread;
+  private Handler playbackThreadHandler;
+  private MediaSourceEventDispatcher eventDispatcher;
+  private ConditionVariable keysLoaded;
+
+  @Before
+  public void setUp() {
+    playbackThread = new HandlerThread("Test playback thread");
+    playbackThread.start();
+    playbackThreadHandler = new Handler(playbackThread.getLooper());
+    eventDispatcher = new MediaSourceEventDispatcher();
+    keysLoaded = TestUtil.createRobolectricConditionVariable();
+    eventDispatcher.addEventListener(
+        playbackThreadHandler,
+        new DrmSessionEventListener() {
+          @Override
+          public void onDrmKeysLoaded(
+              int windowIndex, @Nullable MediaSource.MediaPeriodId mediaPeriodId) {
+            keysLoaded.open();
+          }
+        },
+        DrmSessionEventListener.class);
+  }
+
+  @After
+  public void tearDown() {
+    playbackThread.quitSafely();
+  }
+
+  @Test
+  public void acquireSessionTriggersKeyLoadAndSessionIsOpened() throws Exception {
+    FakeExoMediaDrm.LicenseServer licenseServer =
+        FakeExoMediaDrm.LicenseServer.allowingSchemeDatas(DRM_SCHEME_DATAS);
+
+    keysLoaded.close();
+    AtomicReference<DrmSession> drmSession = new AtomicReference<>();
+    playbackThreadHandler.post(
+        () -> {
+          DefaultDrmSessionManager drmSessionManager =
+              new DefaultDrmSessionManager.Builder()
+                  .setUuidAndExoMediaDrmProvider(DRM_SCHEME_UUID, uuid -> new FakeExoMediaDrm())
+                  .build(/* mediaDrmCallback= */ licenseServer);
+
+          drmSessionManager.prepare();
+          drmSession.set(
+              drmSessionManager.acquireSession(
+                  playbackThread.getLooper(), eventDispatcher, DRM_INIT_DATA));
+        });
+
+    keysLoaded.block(TIMEOUT_MS);
+
+    @DrmSession.State int state = post(drmSession.get(), DrmSession::getState);
+    assertThat(state).isEqualTo(DrmSession.STATE_OPENED_WITH_KEYS);
+    Map<String, String> keyStatus = post(drmSession.get(), DrmSession::queryKeyStatus);
+    assertThat(keyStatus)
+        .containsExactly(FakeExoMediaDrm.KEY_STATUS_KEY, FakeExoMediaDrm.KEY_STATUS_AVAILABLE);
+  }
+
+  /** Call a function on {@code drmSession} on the playback thread and return the result. */
+  private <T> T post(DrmSession drmSession, Function<DrmSession, T> fn)
+      throws InterruptedException {
+    AtomicReference<T> result = new AtomicReference<>();
+    ConditionVariable resultReady = TestUtil.createRobolectricConditionVariable();
+    resultReady.close();
+    playbackThreadHandler.post(
+        () -> {
+          result.set(fn.apply(drmSession));
+          resultReady.open();
+        });
+    resultReady.block(TIMEOUT_MS);
+    return result.get();
+  }
+}

testutils/src/main/java/com/google/android/exoplayer2/testutil/FakeExoMediaDrm.java

@@ -20,36 +20,62 @@ import android.media.DeniedByServerException;
 import android.media.MediaCryptoException;
 import android.media.MediaDrmException;
 import android.media.NotProvisionedException;
+import android.os.Parcel;
+import android.os.Parcelable;
 import android.os.PersistableBundle;
 import androidx.annotation.Nullable;
 import androidx.annotation.RequiresApi;
 import com.google.android.exoplayer2.drm.DrmInitData;
 import com.google.android.exoplayer2.drm.ExoMediaCrypto;
 import com.google.android.exoplayer2.drm.ExoMediaDrm;
+import com.google.android.exoplayer2.drm.MediaDrmCallback;
+import com.google.android.exoplayer2.drm.MediaDrmCallbackException;
 import com.google.android.exoplayer2.util.Assertions;
 import com.google.android.exoplayer2.util.Util;
+import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableMap;
+import com.google.common.collect.ImmutableSet;
+import com.google.common.primitives.Bytes;
 import java.util.ArrayList;
-import java.util.Collections;
 import java.util.HashMap;
 import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
+import java.util.Objects;
 import java.util.Set;
+import java.util.UUID;
 import java.util.concurrent.atomic.AtomicInteger;
 
-/** A fake implementation of {@link ExoMediaDrm} for use in tests. */
-@RequiresApi(18)
+/**
+ * A fake implementation of {@link ExoMediaDrm} for use in tests.
+ *
+ * <p>{@link LicenseServer} can be used to respond to interactions stemming from {@link
+ * #getKeyRequest(byte[], List, int, HashMap)} and {@link #provideKeyResponse(byte[], byte[])}.
+ *
+ * <p>Currently only supports streaming key requests.
+ */
+// TODO: Consider replacing this with a Robolectric ShadowMediaDrm so we can use a real
+//  FrameworkMediaDrm.
+@RequiresApi(29)
 public class FakeExoMediaDrm implements ExoMediaDrm {
 
-  private static final KeyRequest DUMMY_KEY_REQUEST =
-      new KeyRequest(TestUtil.createByteArray(4, 5, 6), "foo.test");
-
-  private static final ProvisionRequest DUMMY_PROVISION_REQUEST =
+  public static final ProvisionRequest DUMMY_PROVISION_REQUEST =
       new ProvisionRequest(TestUtil.createByteArray(7, 8, 9), "bar.test");
 
+  /** Key for use with the Map returned from {@link FakeExoMediaDrm#queryKeyStatus(byte[])}. */
+  public static final String KEY_STATUS_KEY = "KEY_STATUS";
+  /** Value for use with the Map returned from {@link FakeExoMediaDrm#queryKeyStatus(byte[])}. */
+  public static final String KEY_STATUS_AVAILABLE = "AVAILABLE";
+  /** Value for use with the Map returned from {@link FakeExoMediaDrm#queryKeyStatus(byte[])}. */
+  public static final String KEY_STATUS_UNAVAILABLE = "UNAVAILABLE";
+
+  private static final ImmutableList<Byte> VALID_KEY_RESPONSE = TestUtil.createByteList(1, 2, 3);
+  private static final ImmutableList<Byte> KEY_DENIED_RESPONSE = TestUtil.createByteList(9, 8, 7);
+
   private final Map<String, byte[]> byteProperties;
   private final Map<String, String> stringProperties;
   private final Set<List<Byte>> openSessionIds;
+  private final Set<List<Byte>> sessionIdsWithValidKeys;
   private final AtomicInteger sessionIdGenerator;
 
   private int referenceCount;
@@ -62,11 +88,14 @@ public class FakeExoMediaDrm implements ExoMediaDrm {
     byteProperties = new HashMap<>();
     stringProperties = new HashMap<>();
     openSessionIds = new HashSet<>();
+    sessionIdsWithValidKeys = new HashSet<>();
     sessionIdGenerator = new AtomicInteger();
 
     referenceCount = 1;
   }
 
+  // ExoMediaCrypto implementation
+
   @Override
   public void setOnEventListener(@Nullable OnEventListener listener) {
     // Do nothing.
@@ -99,6 +128,7 @@ public class FakeExoMediaDrm implements ExoMediaDrm {
   @Override
   public void closeSession(byte[] sessionId) {
     Assertions.checkState(referenceCount > 0);
+    // TODO: Store closed session IDs too?
     Assertions.checkState(openSessionIds.remove(toByteList(sessionId)));
   }
 
@@ -110,7 +140,18 @@ public class FakeExoMediaDrm implements ExoMediaDrm {
       @Nullable HashMap<String, String> optionalParameters)
       throws NotProvisionedException {
     Assertions.checkState(referenceCount > 0);
-    return DUMMY_KEY_REQUEST;
+    if (keyType == KEY_TYPE_OFFLINE || keyType == KEY_TYPE_RELEASE) {
+      throw new UnsupportedOperationException("Offline key requests are not supported.");
+    }
+    Assertions.checkArgument(keyType == KEY_TYPE_STREAMING, "Unrecognised keyType: " + keyType);
+    Assertions.checkState(openSessionIds.contains(toByteList(scope)));
+    Assertions.checkNotNull(schemeDatas);
+    KeyRequestData requestData =
+        new KeyRequestData(
+            schemeDatas,
+            keyType,
+            optionalParameters != null ? optionalParameters : ImmutableMap.of());
+    return new KeyRequest(requestData.toByteArray(), /* licenseServerUrl= */ "");
   }
 
   @Nullable
@@ -118,7 +159,13 @@ public class FakeExoMediaDrm implements ExoMediaDrm {
   public byte[] provideKeyResponse(byte[] scope, byte[] response)
       throws NotProvisionedException, DeniedByServerException {
     Assertions.checkState(referenceCount > 0);
-    return null;
+    List<Byte> responseAsList = Bytes.asList(response);
+    if (responseAsList.equals(VALID_KEY_RESPONSE)) {
+      sessionIdsWithValidKeys.add(Bytes.asList(scope));
+    } else if (responseAsList.equals(KEY_DENIED_RESPONSE)) {
+      throw new DeniedByServerException("Key request denied");
+    }
+    return Util.EMPTY_BYTE_ARRAY;
   }
 
   @Override
@@ -135,7 +182,12 @@ public class FakeExoMediaDrm implements ExoMediaDrm {
   @Override
   public Map<String, String> queryKeyStatus(byte[] sessionId) {
     Assertions.checkState(referenceCount > 0);
-    return Collections.emptyMap();
+    Assertions.checkState(openSessionIds.contains(toByteList(sessionId)));
+    return ImmutableMap.of(
+        KEY_STATUS_KEY,
+        sessionIdsWithValidKeys.contains(toByteList(sessionId))
+            ? KEY_STATUS_AVAILABLE
+            : KEY_STATUS_UNAVAILABLE);
   }
 
   @Override
@@ -207,13 +259,156 @@ public class FakeExoMediaDrm implements ExoMediaDrm {
     return FakeExoMediaCrypto.class;
   }
 
-  private static List<Byte> toByteList(byte[] byteArray) {
-    List<Byte> result = new ArrayList<>(byteArray.length);
-    for (byte b : byteArray) {
-      result.add(b);
-    }
-    return result;
+  private static ImmutableList<Byte> toByteList(byte[] byteArray) {
+    return ImmutableList.copyOf(Bytes.asList(byteArray));
   }
 
   private static class FakeExoMediaCrypto implements ExoMediaCrypto {}
+
+  /** An license server implementation to interact with {@link FakeExoMediaDrm}. */
+  public static class LicenseServer implements MediaDrmCallback {
+
+    private final List<ImmutableList<DrmInitData.SchemeData>> receivedSchemeDatas;
+    private final ImmutableSet<ImmutableList<DrmInitData.SchemeData>> allowedSchemeDatas;
+
+    @SafeVarargs
+    public static LicenseServer allowingSchemeDatas(List<DrmInitData.SchemeData>... schemeDatas) {
+      ImmutableSet.Builder<ImmutableList<DrmInitData.SchemeData>> schemeDatasBuilder =
+          ImmutableSet.builder();
+      for (List<DrmInitData.SchemeData> schemeData : schemeDatas) {
+        schemeDatasBuilder.add(ImmutableList.copyOf(schemeData));
+      }
+      return new LicenseServer(schemeDatasBuilder.build());
+    }
+
+    private LicenseServer(ImmutableSet<ImmutableList<DrmInitData.SchemeData>> allowedSchemeDatas) {
+      receivedSchemeDatas = new ArrayList<>();
+      this.allowedSchemeDatas = allowedSchemeDatas;
+    }
+
+    public ImmutableList<ImmutableList<DrmInitData.SchemeData>> getReceivedSchemeDatas() {
+      return ImmutableList.copyOf(receivedSchemeDatas);
+    }
+
+    @Override
+    public byte[] executeProvisionRequest(UUID uuid, ProvisionRequest request)
+        throws MediaDrmCallbackException {
+      return new byte[0];
+    }
+
+    @Override
+    public byte[] executeKeyRequest(UUID uuid, KeyRequest request)
+        throws MediaDrmCallbackException {
+      ImmutableList<DrmInitData.SchemeData> schemeDatas =
+          KeyRequestData.fromByteArray(request.getData()).schemeDatas;
+      receivedSchemeDatas.add(schemeDatas);
+      return Bytes.toArray(
+          allowedSchemeDatas.contains(schemeDatas) ? VALID_KEY_RESPONSE : KEY_DENIED_RESPONSE);
+    }
+  }
+
+  /**
+   * A structured set of key request fields that can be serialized into bytes by {@link
+   * #getKeyRequest(byte[], List, int, HashMap)} and then deserialized by {@link
+   * LicenseServer#executeKeyRequest(UUID, KeyRequest)}.
+   */
+  private static class KeyRequestData implements Parcelable {
+    public final ImmutableList<DrmInitData.SchemeData> schemeDatas;
+    public final int type;
+    public final ImmutableMap<String, String> optionalParameters;
+
+    public KeyRequestData(
+        List<DrmInitData.SchemeData> schemeDatas,
+        int type,
+        Map<String, String> optionalParameters) {
+      this.schemeDatas = ImmutableList.copyOf(schemeDatas);
+      this.type = type;
+      this.optionalParameters = ImmutableMap.copyOf(optionalParameters);
+    }
+
+    public KeyRequestData(Parcel in) {
+      this.schemeDatas =
+          ImmutableList.copyOf(
+              in.readParcelableList(
+                  new ArrayList<>(), DrmInitData.SchemeData.class.getClassLoader()));
+      this.type = in.readInt();
+
+      ImmutableMap.Builder<String, String> optionalParameters = new ImmutableMap.Builder<>();
+      List<String> optionalParameterKeys = Assertions.checkNotNull(in.createStringArrayList());
+      List<String> optionalParameterValues = Assertions.checkNotNull(in.createStringArrayList());
+      Assertions.checkArgument(optionalParameterKeys.size() == optionalParameterValues.size());
+      for (int i = 0; i < optionalParameterKeys.size(); i++) {
+        optionalParameters.put(optionalParameterKeys.get(i), optionalParameterValues.get(i));
+      }
+
+      this.optionalParameters = optionalParameters.build();
+    }
+
+    public byte[] toByteArray() {
+      Parcel parcel = Parcel.obtain();
+      try {
+        writeToParcel(parcel, /* flags= */ 0);
+        return parcel.marshall();
+      } finally {
+        parcel.recycle();
+      }
+    }
+
+    public static KeyRequestData fromByteArray(byte[] bytes) {
+      Parcel parcel = Parcel.obtain();
+      try {
+        parcel.unmarshall(bytes, 0, bytes.length);
+        parcel.setDataPosition(0);
+        return CREATOR.createFromParcel(parcel);
+      } finally {
+        parcel.recycle();
+      }
+    }
+
+    @Override
+    public boolean equals(@Nullable Object obj) {
+      if (!(obj instanceof KeyRequestData)) {
+        return false;
+      }
+
+      KeyRequestData that = (KeyRequestData) obj;
+      return Objects.equals(this.schemeDatas, that.schemeDatas)
+          && this.type == that.type
+          && Objects.equals(this.optionalParameters, that.optionalParameters);
+    }
+
+    @Override
+    public int hashCode() {
+      return Objects.hash(schemeDatas, type, optionalParameters);
+    }
+
+    // Parcelable implementation.
+
+    @Override
+    public int describeContents() {
+      return 0;
+    }
+
+    @Override
+    public void writeToParcel(Parcel dest, int flags) {
+      dest.writeParcelableList(schemeDatas, flags);
+      dest.writeInt(type);
+      dest.writeStringList(optionalParameters.keySet().asList());
+      dest.writeStringList(optionalParameters.values().asList());
+    }
+
+    public static final Parcelable.Creator<KeyRequestData> CREATOR =
+        new Parcelable.Creator<KeyRequestData>() {
+
+          @Override
+          public KeyRequestData createFromParcel(Parcel in) {
+            return new KeyRequestData(in);
+          }
+
+          @Override
+          public KeyRequestData[] newArray(int size) {
+            return new KeyRequestData[size];
+          }
+        };
+  }
 }

testutils/src/main/java/com/google/android/exoplayer2/testutil/TestUtil.java

@@ -44,6 +44,8 @@ import com.google.android.exoplayer2.util.ConditionVariable;
 import com.google.android.exoplayer2.util.Supplier;
 import com.google.android.exoplayer2.util.SystemClock;
 import com.google.android.exoplayer2.util.Util;
+import com.google.common.collect.ImmutableList;
+import com.google.common.primitives.Bytes;
 import java.io.File;
 import java.io.FileOutputStream;
 import java.io.IOException;
@@ -174,18 +176,28 @@ public class TestUtil {
   /**
    * Converts an array of integers in the range [0, 255] into an equivalent byte array.
    *
-   * @param intArray An array of integers, all of which must be in the range [0, 255].
+   * @param bytes An array of integers, all of which must be in the range [0, 255].
    * @return The equivalent byte array.
    */
-  public static byte[] createByteArray(int... intArray) {
-    byte[] byteArray = new byte[intArray.length];
+  public static byte[] createByteArray(int... bytes) {
+    byte[] byteArray = new byte[bytes.length];
     for (int i = 0; i < byteArray.length; i++) {
-      Assertions.checkState(0x00 <= intArray[i] && intArray[i] <= 0xFF);
-      byteArray[i] = (byte) intArray[i];
+      Assertions.checkState(0x00 <= bytes[i] && bytes[i] <= 0xFF);
+      byteArray[i] = (byte) bytes[i];
     }
     return byteArray;
   }
 
+  /**
+   * Converts an array of integers in the range [0, 255] into an equivalent byte list.
+   *
+   * @param bytes An array of integers, all of which must be in the range [0, 255].
+   * @return The equivalent byte list.
+   */
+  public static ImmutableList<Byte> createByteList(int... bytes) {
+    return ImmutableList.copyOf(Bytes.asList(createByteArray(bytes)));
+  }
+
   /** Writes one byte long dummy test data to the file and returns it. */
   public static File createTestFile(File directory, String name) throws IOException {
     return createTestFile(directory, name, /* length= */ 1);