Fix doc flag for device having multiple storages

Previously MtpDocumentsProvider#queryDocument returned
FLAG_DIR_SUPPORTS_CREATE for device having multiple storages.

This was wrong because if the device has multiple storages,
MtpDocumentsProvider places the storages under the device document, thus
users cannot create a document just under the device document.

Bug: 35700994
Test: MtpDocumentsProviderTests
Change-Id: Id73a34a2eaf4e10e23be3c2da7488036cea10000
(cherry picked from commit 66fcb4beae9605cd034c9437e510b21260d6f519)

3 files changed, 120 insertions, 8 deletions

diff --git a/packages/MtpDocumentsProvider/src/com/android/mtp/MtpDatabase.java b/packages/MtpDocumentsProvider/src/com/android/mtp/MtpDatabase.java
index 4950af3e14e0..9b5982b96e31 100644
--- a/packages/MtpDocumentsProvider/src/com/android/mtp/MtpDatabase.java
+++ b/packages/MtpDocumentsProvider/src/com/android/mtp/MtpDatabase.java
@@ -200,10 +200,7 @@ class MtpDatabase {
                     storageCursor.close();
                 }
 
-                final RowBuilder row = result.newRow();
-                for (final String key : values.keySet()) {
-                    row.add(key, values.get(key));
-                }
+                putValuesToCursor(values, result);
             }
 
             return result;
@@ -760,7 +757,9 @@ class MtpDatabase {
                 Document.MIME_TYPE_DIR,
                 0,
                 MtpConstants.PROTECTION_STATUS_NONE,
-                DOCUMENT_TYPE_DEVICE));
+                // Storages are placed under device so we cannot create a document just under
+                // device.
+                DOCUMENT_TYPE_DEVICE) & ~Document.FLAG_DIR_SUPPORTS_CREATE);
         values.putNull(Document.COLUMN_SIZE);
 
         extraValues.clear();
@@ -915,6 +914,13 @@ class MtpDatabase {
         return results;
     }
 
+    static void putValuesToCursor(ContentValues values, MatrixCursor cursor) {
+        final RowBuilder row = cursor.newRow();
+        for (final String name : cursor.getColumnNames()) {
+            row.add(values.get(name));
+        }
+    }
+
     private static String getIdList(Set<String> ids) {
         String result = "(";
         for (final String id : ids) {
diff --git a/packages/MtpDocumentsProvider/src/com/android/mtp/MtpDocumentsProvider.java b/packages/MtpDocumentsProvider/src/com/android/mtp/MtpDocumentsProvider.java
index db88f2c038ec..eb2d8aa71955 100644
--- a/packages/MtpDocumentsProvider/src/com/android/mtp/MtpDocumentsProvider.java
+++ b/packages/MtpDocumentsProvider/src/com/android/mtp/MtpDocumentsProvider.java
@@ -17,11 +17,13 @@
 package com.android.mtp;
 
 import android.content.ContentResolver;
+import android.content.ContentValues;
 import android.content.Context;
 import android.content.UriPermission;
 import android.content.res.AssetFileDescriptor;
 import android.content.res.Resources;
 import android.database.Cursor;
+import android.database.DatabaseUtils;
 import android.database.MatrixCursor;
 import android.database.sqlite.SQLiteDiskIOException;
 import android.graphics.Point;
@@ -55,7 +57,6 @@ import java.util.LinkedList;
 import java.util.List;
 import java.util.Map;
 import java.util.concurrent.TimeoutException;
-
 import libcore.io.IoUtils;
 
 /**
@@ -177,7 +178,57 @@ public class MtpDocumentsProvider extends DocumentsProvider {
         if (projection == null) {
             projection = MtpDocumentsProvider.DEFAULT_DOCUMENT_PROJECTION;
         }
-        return mDatabase.queryDocument(documentId, projection);
+        final Cursor cursor = mDatabase.queryDocument(documentId, projection);
+        final int cursorCount = cursor.getCount();
+        if (cursorCount == 0) {
+            cursor.close();
+            throw new FileNotFoundException();
+        } else if (cursorCount != 1) {
+            cursor.close();
+            Log.wtf(TAG, "Unexpected cursor size: " + cursorCount);
+            return null;
+        }
+
+        final Identifier identifier = mDatabase.createIdentifier(documentId);
+        if (identifier.mDocumentType != MtpDatabaseConstants.DOCUMENT_TYPE_DEVICE) {
+            return cursor;
+        }
+        final String[] storageDocIds = mDatabase.getStorageDocumentIds(documentId);
+        if (storageDocIds.length != 1) {
+            return mDatabase.queryDocument(documentId, projection);
+        }
+
+        // If the documentId specifies a device having exact one storage, we repalce some device
+        // attributes with the storage attributes.
+        try {
+            final String storageName;
+            final int storageFlags;
+            try (final Cursor storageCursor = mDatabase.queryDocument(
+                    storageDocIds[0],
+                    MtpDatabase.strings(Document.COLUMN_DISPLAY_NAME, Document.COLUMN_FLAGS))) {
+                if (!storageCursor.moveToNext()) {
+                    throw new FileNotFoundException();
+                }
+                storageName = storageCursor.getString(0);
+                storageFlags = storageCursor.getInt(1);
+            }
+
+            cursor.moveToNext();
+            final ContentValues values = new ContentValues();
+            DatabaseUtils.cursorRowToContentValues(cursor, values);
+            if (values.containsKey(Document.COLUMN_DISPLAY_NAME)) {
+                values.put(Document.COLUMN_DISPLAY_NAME, mResources.getString(
+                        R.string.root_name,
+                        values.getAsString(Document.COLUMN_DISPLAY_NAME),
+                        storageName));
+            }
+            values.put(Document.COLUMN_FLAGS, storageFlags);
+            final MatrixCursor output = new MatrixCursor(projection, 1);
+            MtpDatabase.putValuesToCursor(values, output);
+            return output;
+        } finally {
+            cursor.close();
+        }
     }
 
     @Override
diff --git a/packages/MtpDocumentsProvider/tests/src/com/android/mtp/MtpDocumentsProviderTest.java b/packages/MtpDocumentsProvider/tests/src/com/android/mtp/MtpDocumentsProviderTest.java
index c9420d13a6e5..3fa5eb581ef7 100644
--- a/packages/MtpDocumentsProvider/tests/src/com/android/mtp/MtpDocumentsProviderTest.java
+++ b/packages/MtpDocumentsProvider/tests/src/com/android/mtp/MtpDocumentsProviderTest.java
@@ -368,7 +368,7 @@ public class MtpDocumentsProviderTest extends AndroidTestCase {
         assertEquals(0, cursor.getInt(5));
     }
 
-    public void testQueryDocument_forRoot()
+    public void testQueryDocument_forStorage()
             throws IOException, InterruptedException, TimeoutException {
         setupProvider(MtpDatabaseConstants.FLAG_DATABASE_IN_MEMORY);
         setupRoots(0, new MtpRoot[] {
@@ -392,6 +392,61 @@ public class MtpDocumentsProviderTest extends AndroidTestCase {
         assertEquals(3072, cursor.getInt(5));
     }
 
+    public void testQueryDocument_forDeviceWithSingleStorage()
+            throws IOException, InterruptedException, TimeoutException {
+        setupProvider(MtpDatabaseConstants.FLAG_DATABASE_IN_MEMORY);
+        setupRoots(0, new MtpRoot[] {
+                new MtpRoot(
+                        0 /* deviceId */,
+                        1 /* storageId */,
+                        "Storage A" /* volume description */,
+                        1024 /* free space */,
+                        4096 /* total space */,
+                        "" /* no volume identifier */)
+        });
+        final Cursor cursor = mProvider.queryDocument("1", null);
+        assertEquals(1, cursor.getCount());
+
+        cursor.moveToNext();
+        assertEquals("1", cursor.getString(0));
+        assertEquals(DocumentsContract.Document.MIME_TYPE_DIR, cursor.getString(1));
+        assertEquals("Device Storage A", cursor.getString(2));
+        assertTrue(cursor.isNull(3));
+        assertEquals(DocumentsContract.Document.FLAG_DIR_SUPPORTS_CREATE, cursor.getInt(4));
+        assertTrue(cursor.isNull(5));
+    }
+
+    public void testQueryDocument_forDeviceWithTwoStorages()
+            throws IOException, InterruptedException, TimeoutException {
+        setupProvider(MtpDatabaseConstants.FLAG_DATABASE_IN_MEMORY);
+        setupRoots(0, new MtpRoot[] {
+                new MtpRoot(
+                        0 /* deviceId */,
+                        1 /* storageId */,
+                        "Storage A" /* volume description */,
+                        1024 /* free space */,
+                        4096 /* total space */,
+                        "" /* no volume identifier */),
+                new MtpRoot(
+                        0 /* deviceId */,
+                        2 /* storageId */,
+                        "Storage B" /* volume description */,
+                        1024 /* free space */,
+                        4096 /* total space */,
+                        "" /* no volume identifier */)
+        });
+        final Cursor cursor = mProvider.queryDocument("1", null);
+        assertEquals(1, cursor.getCount());
+
+        cursor.moveToNext();
+        assertEquals("1", cursor.getString(0));
+        assertEquals(DocumentsContract.Document.MIME_TYPE_DIR, cursor.getString(1));
+        assertEquals("Device", cursor.getString(2));
+        assertTrue(cursor.isNull(3));
+        assertEquals(0, cursor.getInt(4));
+        assertTrue(cursor.isNull(5));
+    }
+
     public void testQueryChildDocuments() throws Exception {
         setupProvider(MtpDatabaseConstants.FLAG_DATABASE_IN_MEMORY);
         setupRoots(0, new MtpRoot[] { new MtpRoot(0, 0, "Storage", 1000, 1000, "") });