Add -a to capture all the physical displays.

when -a is given, it will append a postfix to the FILENAME, such that
FILENAME_0.png, FILENAME_1.png. this won't break the existing behavior
as it doesn't mutate the name if it's taking a single display.
when both -a and -d are given, it will ignore -d and capture all the
displays.
also updated the usage doc to clarify it captures the default display
when the id is not given.

Test cases:
- no argument > prints out in the cmd window
- only -p > prints out in the cmd window as png format
- -h > correctly shows the usage
- -d with and without .png > saves as a file
- -a with and without .png > saves as files
- -d and -a > ignores -d and shows the same results as single -a

Bug: 321278149
Test: adb shell screencap with various flags
Flag: NA
Change-Id: Iecfeec1a1edbc95d7e8931ac3b22ac7f0706c3e7

1 files changed, 184 insertions, 102 deletions

diff --git a/cmds/screencap/screencap.cpp b/cmds/screencap/screencap.cpp
index 917529ec1dcf..7e4f95bc9274 100644
--- a/cmds/screencap/screencap.cpp
+++ b/cmds/screencap/screencap.cpp
@@ -51,11 +51,13 @@ using namespace android;
 
 void usage(const char* pname, ftl::Optional<DisplayId> displayIdOpt) {
     fprintf(stderr, R"(
-usage: %s [-hp] [-d display-id] [FILENAME]
+usage: %s [-ahp] [-d display-id] [FILENAME]
    -h: this message
-   -p: save the file as a png.
+   -a: captures all the active displays. This appends an integer postfix to the FILENAME.
+       e.g., FILENAME_0.png, FILENAME_1.png. If both -a and -d are given, it ignores -d.
    -d: specify the display ID to capture%s
        see "dumpsys SurfaceFlinger --display-id" for valid display IDs.
+   -p: outputs in png format.
    --hint-for-seamless If set will use the hintForSeamless path in SF
 
 If FILENAME ends with .png it will be saved as a png.
@@ -63,11 +65,13 @@ If FILENAME is not given, the results will be printed to stdout.
 )",
             pname,
             displayIdOpt
-                    .transform([](DisplayId id) {
-                        return std::string(ftl::Concat(" (default: ", id.value, ')').str());
-                    })
-                    .value_or(std::string())
-                    .c_str());
+                .transform([](DisplayId id) {
+                    return std::string(ftl::Concat(
+                    " (If the id is not given, it defaults to ", id.value,')'
+                    ).str());
+                })
+                .value_or(std::string())
+                .c_str());
 }
 
 // For options that only exist in long-form. Anything in the
@@ -123,8 +127,8 @@ static status_t notifyMediaScanner(const char* fileName) {
     int status;
     int pid = fork();
     if (pid < 0){
-       fprintf(stderr, "Unable to fork in order to send intent for media scanner.\n");
-       return UNKNOWN_ERROR;
+        fprintf(stderr, "Unable to fork in order to send intent for media scanner.\n");
+        return UNKNOWN_ERROR;
     }
     if (pid == 0){
         int fd = open("/dev/null", O_WRONLY);
@@ -146,19 +150,119 @@ static status_t notifyMediaScanner(const char* fileName) {
     return NO_ERROR;
 }
 
+status_t capture(const DisplayId displayId,
+            const gui::CaptureArgs& captureArgs,
+            ScreenCaptureResults& outResult) {
+    sp<SyncScreenCaptureListener> captureListener = new SyncScreenCaptureListener();
+    ScreenshotClient::captureDisplay(displayId, captureArgs, captureListener);
+
+    ScreenCaptureResults captureResults = captureListener->waitForResults();
+    if (!captureResults.fenceResult.ok()) {
+        fprintf(stderr, "Failed to take screenshot. Status: %d\n",
+                fenceStatus(captureResults.fenceResult));
+        return 1;
+    }
+
+    outResult = captureResults;
+
+    return 0;
+}
+
+status_t saveImage(const char* fn, bool png, const ScreenCaptureResults& captureResults) {
+    void* base = nullptr;
+    ui::Dataspace dataspace = captureResults.capturedDataspace;
+    sp<GraphicBuffer> buffer = captureResults.buffer;
+
+    status_t result = buffer->lock(GraphicBuffer::USAGE_SW_READ_OFTEN, &base);
+
+    if (base == nullptr || result != NO_ERROR) {
+        String8 reason;
+        if (result != NO_ERROR) {
+            reason.appendFormat(" Error Code: %d", result);
+        } else {
+            reason = "Failed to write to buffer";
+        }
+        fprintf(stderr, "Failed to take screenshot (%s)\n", reason.c_str());
+        return 1;
+    }
+
+    int fd = -1;
+    if (fn == nullptr) {
+        fd = dup(STDOUT_FILENO);
+        if (fd == -1) {
+            fprintf(stderr, "Error writing to stdout. (%s)\n", strerror(errno));
+            return 1;
+        }
+    } else {
+        fd = open(fn, O_WRONLY | O_CREAT | O_TRUNC, 0664);
+        if (fd == -1) {
+            fprintf(stderr, "Error opening file: %s (%s)\n", fn, strerror(errno));
+            return 1;
+        }
+    }
+
+    if (png) {
+        AndroidBitmapInfo info;
+        info.format = flinger2bitmapFormat(buffer->getPixelFormat());
+        info.flags = ANDROID_BITMAP_FLAGS_ALPHA_PREMUL;
+        info.width = buffer->getWidth();
+        info.height = buffer->getHeight();
+        info.stride = buffer->getStride() * bytesPerPixel(buffer->getPixelFormat());
+
+        int result = AndroidBitmap_compress(&info, static_cast<int32_t>(dataspace), base,
+                                            ANDROID_BITMAP_COMPRESS_FORMAT_PNG, 100, &fd,
+                                            [](void* fdPtr, const void* data, size_t size) -> bool {
+                                                int bytesWritten = write(*static_cast<int*>(fdPtr),
+                                                                         data, size);
+                                                return bytesWritten == size;
+                                            });
+
+        if (result != ANDROID_BITMAP_RESULT_SUCCESS) {
+            fprintf(stderr, "Failed to compress PNG (error code: %d)\n", result);
+        }
+
+        if (fn != NULL) {
+            notifyMediaScanner(fn);
+        }
+    } else {
+        uint32_t w = buffer->getWidth();
+        uint32_t h = buffer->getHeight();
+        uint32_t s = buffer->getStride();
+        uint32_t f = buffer->getPixelFormat();
+        uint32_t c = dataSpaceToInt(dataspace);
+
+        write(fd, &w, 4);
+        write(fd, &h, 4);
+        write(fd, &f, 4);
+        write(fd, &c, 4);
+        size_t Bpp = bytesPerPixel(f);
+        for (size_t y=0 ; y<h ; y++) {
+            write(fd, base, w*Bpp);
+            base = (void *)((char *)base + s*Bpp);
+        }
+    }
+    close(fd);
+
+    return 0;
+}
+
 int main(int argc, char** argv)
 {
-    const std::vector<PhysicalDisplayId> ids = SurfaceComposerClient::getPhysicalDisplayIds();
-    if (ids.empty()) {
+    const std::vector<PhysicalDisplayId> physicalDisplays =
+        SurfaceComposerClient::getPhysicalDisplayIds();
+
+    if (physicalDisplays.empty()) {
         fprintf(stderr, "Failed to get ID for any displays.\n");
         return 1;
     }
     std::optional<DisplayId> displayIdOpt;
+    std::vector<DisplayId> displaysToCapture;
     gui::CaptureArgs captureArgs;
     const char* pname = argv[0];
     bool png = false;
+    bool all = false;
     int c;
-    while ((c = getopt_long(argc, argv, "phd:", LONG_OPTIONS, nullptr)) != -1) {
+    while ((c = getopt_long(argc, argv, "aphd:", LONG_OPTIONS, nullptr)) != -1) {
         switch (c) {
             case 'p':
                 png = true;
@@ -177,12 +281,17 @@ int main(int argc, char** argv)
                     fprintf(stderr, "Invalid display ID: Incorrect encoding.\n");
                     return 1;
                 }
+                displaysToCapture.push_back(displayIdOpt.value());
+                break;
+            }
+            case 'a': {
+                all = true;
                 break;
             }
             case '?':
             case 'h':
-                if (ids.size() == 1) {
-                    displayIdOpt = ids.front();
+                if (physicalDisplays.size() >= 1) {
+                    displayIdOpt = physicalDisplays.front();
                 }
                 usage(pname, displayIdOpt);
                 return 1;
@@ -192,44 +301,52 @@ int main(int argc, char** argv)
         }
     }
 
-    if (!displayIdOpt) {
-        displayIdOpt = ids.front();
-        if (ids.size() > 1) {
-            fprintf(stderr,
-                    "[Warning] Multiple displays were found, but no display id was specified! "
-                    "Defaulting to the first display found, however this default is not guaranteed "
-                    "to be consistent across captures. A display id should be specified.\n");
-            fprintf(stderr, "A display ID can be specified with the [-d display-id] option.\n");
-            fprintf(stderr, "See \"dumpsys SurfaceFlinger --display-id\" for valid display IDs.\n");
-        }
-    }
-
     argc -= optind;
     argv += optind;
 
-    int fd = -1;
-    const char* fn = NULL;
-    if (argc == 0) {
-        fd = dup(STDOUT_FILENO);
-    } else if (argc == 1) {
-        fn = argv[0];
-        fd = open(fn, O_WRONLY | O_CREAT | O_TRUNC, 0664);
-        if (fd == -1) {
-            fprintf(stderr, "Error opening file: %s (%s)\n", fn, strerror(errno));
-            return 1;
+    // We don't expect more than 2 arguments.
+    if (argc >= 2) {
+        if (physicalDisplays.size() >= 1) {
+            usage(pname, physicalDisplays.front());
+        } else {
+            usage(pname, std::nullopt);
         }
-        const int len = strlen(fn);
-        if (len >= 4 && 0 == strcmp(fn+len-4, ".png")) {
+        return 1;
+    }
+
+    std::string baseName;
+    std::string suffix;
+
+    if (argc == 1) {
+        std::string_view filename = { argv[0] };
+        if (filename.ends_with(".png")) {
+            baseName = filename.substr(0, filename.size()-4);
+            suffix = ".png";
             png = true;
+        } else {
+            baseName = filename;
         }
     }
 
-    if (fd == -1) {
-        usage(pname, displayIdOpt);
-        return 1;
+    if (all) {
+        // Ignores -d if -a is given.
+        displaysToCapture.clear();
+        for (int i = 0; i < physicalDisplays.size(); i++) {
+            displaysToCapture.push_back(physicalDisplays[i]);
+        }
     }
 
-    void* base = NULL;
+    if (displaysToCapture.empty()) {
+        displaysToCapture.push_back(physicalDisplays.front());
+        if (physicalDisplays.size() > 1) {
+            fprintf(stderr,
+                    "[Warning] Multiple displays were found, but no display id was specified! "
+                    "Defaulting to the first display found, however this default is not guaranteed "
+                    "to be consistent across captures. A display id should be specified.\n");
+            fprintf(stderr, "A display ID can be specified with the [-d display-id] option.\n");
+            fprintf(stderr, "See \"dumpsys SurfaceFlinger --display-id\" for valid display IDs.\n");
+        }
+    }
 
     // setThreadPoolMaxThreadCount(0) actually tells the kernel it's
     // not allowed to spawn any additional threads, but we still spawn
@@ -238,74 +355,39 @@ int main(int argc, char** argv)
     ProcessState::self()->setThreadPoolMaxThreadCount(0);
     ProcessState::self()->startThreadPool();
 
-    sp<SyncScreenCaptureListener> captureListener = new SyncScreenCaptureListener();
-    ScreenshotClient::captureDisplay(*displayIdOpt, captureArgs, captureListener);
-
-    ScreenCaptureResults captureResults = captureListener->waitForResults();
-    if (!captureResults.fenceResult.ok()) {
-        close(fd);
-        fprintf(stderr, "Failed to take screenshot. Status: %d\n",
-            fenceStatus(captureResults.fenceResult));
-        return 1;
-    }
-    ui::Dataspace dataspace = captureResults.capturedDataspace;
-    sp<GraphicBuffer> buffer = captureResults.buffer;
+    std::vector<ScreenCaptureResults> results;
+    const size_t numDisplays = displaysToCapture.size();
+    for (int i=0; i<numDisplays; i++) {
+        ScreenCaptureResults result;
 
-    status_t result = buffer->lock(GraphicBuffer::USAGE_SW_READ_OFTEN, &base);
+        // 1. Capture the screen
+        if (const status_t captureStatus =
+            capture(displaysToCapture[i], captureArgs, result) != 0) {
 
-    if (base == nullptr || result != NO_ERROR) {
-        String8 reason;
-        if (result != NO_ERROR) {
-            reason.appendFormat(" Error Code: %d", result);
-        } else {
-            reason = "Failed to write to buffer";
+            fprintf(stderr, "Capturing failed.\n");
+            return captureStatus;
         }
-        fprintf(stderr, "Failed to take screenshot (%s)\n", reason.c_str());
-        close(fd);
-        return 1;
-    }
-
-    if (png) {
-        AndroidBitmapInfo info;
-        info.format = flinger2bitmapFormat(buffer->getPixelFormat());
-        info.flags = ANDROID_BITMAP_FLAGS_ALPHA_PREMUL;
-        info.width = buffer->getWidth();
-        info.height = buffer->getHeight();
-        info.stride = buffer->getStride() * bytesPerPixel(buffer->getPixelFormat());
-
-        int result = AndroidBitmap_compress(&info, static_cast<int32_t>(dataspace), base,
-                                            ANDROID_BITMAP_COMPRESS_FORMAT_PNG, 100, &fd,
-                                            [](void* fdPtr, const void* data, size_t size) -> bool {
-                                                int bytesWritten = write(*static_cast<int*>(fdPtr),
-                                                                         data, size);
-                                                return bytesWritten == size;
-                                            });
 
-        if (result != ANDROID_BITMAP_RESULT_SUCCESS) {
-            fprintf(stderr, "Failed to compress PNG (error code: %d)\n", result);
+        // 2. Save the capture result as an image.
+        // When there's more than one file to capture, add the index as postfix.
+        std::string filename;
+        if (!baseName.empty()) {
+            filename = baseName;
+            if (numDisplays > 1) {
+                filename += "_";
+                filename += std::to_string(i);
+            }
+            filename += suffix;
         }
-
-        if (fn != NULL) {
-            notifyMediaScanner(fn);
+        const char* fn = nullptr;
+        if (!filename.empty()) {
+            fn = filename.c_str();
         }
-    } else {
-        uint32_t w = buffer->getWidth();
-        uint32_t h = buffer->getHeight();
-        uint32_t s = buffer->getStride();
-        uint32_t f = buffer->getPixelFormat();
-        uint32_t c = dataSpaceToInt(dataspace);
-
-        write(fd, &w, 4);
-        write(fd, &h, 4);
-        write(fd, &f, 4);
-        write(fd, &c, 4);
-        size_t Bpp = bytesPerPixel(f);
-        for (size_t y=0 ; y<h ; y++) {
-            write(fd, base, w*Bpp);
-            base = (void *)((char *)base + s*Bpp);
+        if (const status_t saveImageStatus = saveImage(fn, png, result) != 0) {
+            fprintf(stderr, "Saving image failed.\n");
+            return saveImageStatus;
         }
     }
-    close(fd);
 
     return 0;
 }