CrScreenshotDxe.c

Go to the documentation of this file.

/* CrScreenshotDxe.c
 
Copyright (c) 2016, Nikolaj Schlej, All rights reserved.
 
Redistribution and use in source and binary forms,
with or without modification, are permitted provided that the following conditions are met:
- Redistributions of source code must retain the above copyright notice,
  this list of conditions and the following disclaimer.
- Redistributions in binary form must reproduce the above copyright notice,
  this list of conditions and the following disclaimer in the documentation
  and/or other materials provided with the distribution.
 
THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO,
THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED.
IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT,
INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION)
HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE,
EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 
*/
 
#include <Uefi.h>
#include <Library/BaseLib.h>
#include <Library/DebugLib.h>
#include <Library/OcBootManagementLib.h>
#include <Library/OcPngLib.h>
#include <Library/OcFileLib.h>
#include <Library/OcMiscLib.h>
#include <Library/PrintLib.h>
#include <Library/TimerLib.h>
#include <Library/UefiDriverEntryPoint.h>
#include <Library/UefiBootServicesTableLib.h>
#include <Library/UefiRuntimeServicesTableLib.h>
 
#include <Protocol/AppleEvent.h>
#include <Protocol/GraphicsOutput.h>
#include <Protocol/OcBootstrap.h>
#include <Protocol/SimpleTextInEx.h>
#include <Protocol/SimpleFileSystem.h>
 
//
// Keyboard protocol arrival event.
//
STATIC EFI_EVENT  mProtocolNotification;
 
STATIC UINT64   mPreviousTime     = 0;
STATIC BOOLEAN  mEnableMouseClick = FALSE;
 
STATIC
EFI_STATUS
EFIAPI
ShowStatus (
  IN UINT8  Red,
  IN UINT8  Green,
  IN UINT8  Blue
  )
{
  //
  // Determines the size of status square.
  //
  #define STATUS_SQUARE_SIDE  5
 
  EFI_STATUS                     Status;
  EFI_GRAPHICS_OUTPUT_PROTOCOL   *GraphicsOutput;
  EFI_GRAPHICS_OUTPUT_BLT_PIXEL  Square[STATUS_SQUARE_SIDE * STATUS_SQUARE_SIDE];
  EFI_GRAPHICS_OUTPUT_BLT_PIXEL  Backup[STATUS_SQUARE_SIDE * STATUS_SQUARE_SIDE];
  UINTN                          Index;
 
  Status = OcHandleProtocolFallback (
             gST->ConsoleOutHandle,
             &gEfiGraphicsOutputProtocolGuid,
             (VOID **)&GraphicsOutput
             );
  if (EFI_ERROR (Status)) {
    DEBUG ((DEBUG_INFO, "OCSCR: Graphics output protocol not found for status - %r\n", Status));
    return EFI_UNSUPPORTED;
  }
 
  //
  // Set square color.
  //
  for (Index = 0; Index < STATUS_SQUARE_SIDE * STATUS_SQUARE_SIDE; ++Index) {
    Square[Index].Blue     = Blue;
    Square[Index].Green    = Green;
    Square[Index].Red      = Red;
    Square[Index].Reserved = 0x00;
  }
 
  //
  // Backup current image.
  //
  GraphicsOutput->Blt (
                    GraphicsOutput,
                    Backup,
                    EfiBltVideoToBltBuffer,
                    0,
                    0,
                    0,
                    0,
                    STATUS_SQUARE_SIDE,
                    STATUS_SQUARE_SIDE,
                    0
                    );
 
  //
  // Draw the status square.
  //
  GraphicsOutput->Blt (
                    GraphicsOutput,
                    Square,
                    EfiBltBufferToVideo,
                    0,
                    0,
                    0,
                    0,
                    STATUS_SQUARE_SIDE,
                    STATUS_SQUARE_SIDE,
                    0
                    );
 
  //
  // Wait 500 ms.
  //
  gBS->Stall (500*1000);
 
  //
  // Restore the backup.
  //
  GraphicsOutput->Blt (
                    GraphicsOutput,
                    Backup,
                    EfiBltBufferToVideo,
                    0,
                    0,
                    0,
                    0,
                    STATUS_SQUARE_SIDE,
                    STATUS_SQUARE_SIDE,
                    0
                    );
 
  return EFI_SUCCESS;
}
 
STATIC
EFI_STATUS
EFIAPI
TakeScreenshot (
  IN EFI_KEY_DATA  *KeyData
  )
{
  EFI_FILE_PROTOCOL              *Fs;
  EFI_GRAPHICS_OUTPUT_PROTOCOL   *GraphicsOutput;
  EFI_GRAPHICS_OUTPUT_BLT_PIXEL  *Image;
  UINTN                          ImageSize;        
  VOID                           *PngFile;
  UINTN                          PngFileSize;      
  EFI_STATUS                     Status;
  UINT32                         ScreenWidth;
  UINT32                         ScreenHeight;
  CHAR16                         FileName[16];
  EFI_TIME                       Time;
  UINTN                          Index;
  UINT8                          Temp;
 
  //
  // This is required to avoid assert (only noticeable on firmware compiled to
  // assert) from gBS->RaiseTPL(TPL_CALLBACK) within the file system accesses
  // below. Makes explicit what was happening anyway, which is that we're
  // effectively lowering the TPL for a long running task, during a keyboard
  // interrupt.
  // REF: https://github.com/acidanthera/audk/blob/bcdcc4160d7460c46c08c9395aae81be44ef23a9/FatPkg/EnhancedFatDxe/Misc.c#L399
  //
  gBS->RestoreTPL (TPL_CALLBACK);
 
  Status = OcFindWritableOcFileSystem (&Fs);
  if (EFI_ERROR (Status)) {
    DEBUG ((DEBUG_INFO, "OCSCR: Can't find writable FS - %r\n", Status));
    ShowStatus (0xFF, 0xFF, 0x00); 
    return EFI_SUCCESS;
  }
 
  Status = OcHandleProtocolFallback (
             gST->ConsoleOutHandle,
             &gEfiGraphicsOutputProtocolGuid,
             (VOID **)&GraphicsOutput
             );
  if (EFI_ERROR (Status)) {
    DEBUG ((DEBUG_INFO, "OCSCR: Graphics output protocol not found for screen - %r\n", Status));
    Fs->Close (Fs);
    return EFI_SUCCESS;
  }
 
  //
  // Set screen width, height and image size in pixels.
  //
  ScreenWidth  = GraphicsOutput->Mode->Info->HorizontalResolution;
  ScreenHeight = GraphicsOutput->Mode->Info->VerticalResolution;
  ImageSize    = (UINTN)ScreenWidth * ScreenHeight;
 
  if (ImageSize == 0) {
    DEBUG ((DEBUG_INFO, "OCSCR: Empty screen size\n"));
    Fs->Close (Fs);
    return EFI_SUCCESS;
  }
 
  //
  // Get current time.
  //
  Status = gRT->GetTime (
                  &Time,
                  NULL
                  );
  if (!EFI_ERROR (Status)) {
    //
    // Set file name to current day and time
    //
    UnicodeSPrint (
      FileName,
      sizeof (FileName),
      L"%02d%02d%02d%02d.png",
      Time.Day,
      Time.Hour,
      Time.Minute,
      Time.Second
      );
  } else {
    //
    // Set file name to scrnshot.png
    //
    StrCpyS (FileName, sizeof (FileName), L"scrnshot.png");
  }
 
  //
  // Allocate memory for screenshot.
  //
  Status = gBS->AllocatePool (
                  EfiBootServicesData,
                  ImageSize * sizeof (EFI_GRAPHICS_OUTPUT_BLT_PIXEL),
                  (VOID **)&Image
                  );
  if (EFI_ERROR (Status)) {
    DEBUG ((DEBUG_INFO, "CRSCR: gBS->AllocatePool returned %r\n", Status));
    ShowStatus (0xFF, 0x00, 0x00); 
    Fs->Close (Fs);
    return EFI_SUCCESS;
  }
 
  //
  // Take screenshot.
  //
  Status = GraphicsOutput->Blt (
                             GraphicsOutput,
                             Image,
                             EfiBltVideoToBltBuffer,
                             0,
                             0,
                             0,
                             0,
                             ScreenWidth,
                             ScreenHeight,
                             0
                             );
  if (EFI_ERROR (Status)) {
    DEBUG ((DEBUG_INFO, "CRSCR: GraphicsOutput->Blt returned %r\n", Status));
    gBS->FreePool (Image);
    ShowStatus (0xFF, 0x00, 0x00); 
    Fs->Close (Fs);
    return EFI_SUCCESS;
  }
 
  //
  // Convert BGR to RGBA with Alpha set to 0xFF.
  //
  for (Index = 0; Index < ImageSize; ++Index) {
    Temp                  = Image[Index].Blue;
    Image[Index].Blue     = Image[Index].Red;
    Image[Index].Red      = Temp;
    Image[Index].Reserved = 0xFF;
  }
 
  Status = OcEncodePng (
             Image,
             ScreenWidth,
             ScreenHeight,
             &PngFile,
             &PngFileSize
             );
  gBS->FreePool (Image);
  if (EFI_ERROR (Status)) {
    DEBUG ((DEBUG_INFO, "CRSCR: OcEncodePng returned %r\n", Status));
    ShowStatus (0xFF, 0x00, 0x00); 
    Fs->Close (Fs);
    return EFI_SUCCESS;
  }
 
  //
  // Write PNG image into the file.
  //
  Status = OcSetFileData (Fs, FileName, PngFile, (UINT32)PngFileSize);
  gBS->FreePool (PngFile);
  Fs->Close (Fs);
  if (EFI_ERROR (Status)) {
    DEBUG ((DEBUG_INFO, "CRSCR: OcEncodePng returned %r\n", Status));
    ShowStatus (0xFF, 0x00, 0x00); 
    return EFI_SUCCESS;
  }
 
  //
  // Show success.
  //
  ShowStatus (0x00, 0xFF, 0x00); 
 
  return EFI_SUCCESS;
}
 
STATIC
VOID
EFIAPI
AppleEventKeyHandler (
  IN APPLE_EVENT_INFORMATION  *Information,
  IN VOID                     *NotifyContext
  )
{
  UINT64  CurrTime;
 
  //
  // Mark the context argument as used.
  //
  (VOID)NotifyContext;
 
  //
  // Ignore invalid information if it happened to arrive.
  //
  if (Information == NULL) {
    return;
  }
 
  if (  ((Information->EventType & APPLE_EVENT_TYPE_KEY_DOWN) != 0)
     && (Information->EventData.KeyData->AppleKeyCode == AppleHidUsbKbUsageKeyF10))
  {
    //
    // Take a screenshot for F10.
    //
    TakeScreenshot (NULL);
  }
 
  if (mEnableMouseClick && ((Information->EventType & APPLE_EVENT_TYPE_MOUSE_CLICK) != 0)) {
    //
    // Debounce, prevents multiple detections from one click.
    //
    CurrTime = GetTimeInNanoSecond (GetPerformanceCounter ());
    if (CurrTime - mPreviousTime > MS_TO_NANOSECONDS (1000ULL)) {
      //
      // Take a screenshot for mouse click.
      //
      TakeScreenshot (NULL);
      mPreviousTime = GetTimeInNanoSecond (GetPerformanceCounter ());
    }
  }
}
 
STATIC
EFI_STATUS
InstallKeyHandler (
  VOID
  )
{
  EFI_STATUS                         Status;
  UINTN                              HandleCount;
  EFI_HANDLE                         *HandleBuffer;
  UINTN                              Index;
  EFI_KEY_DATA                       SimpleTextInExKeyStroke;
  EFI_HANDLE                         SimpleTextInExHandle;
  EFI_SIMPLE_TEXT_INPUT_EX_PROTOCOL  *SimpleTextInEx;
  APPLE_EVENT_HANDLE                 AppleEventHandle;
  APPLE_EVENT_PROTOCOL               *AppleEvent;
  BOOLEAN                            Installed;
 
  Installed = FALSE;
 
  //
  // Locate compatible protocols, firstly try AppleEvent otherwise try SimpleTextInEx.
  // This is because we want key swap to take precedence.
  //
  Status = gBS->LocateHandleBuffer (
                  ByProtocol,
                  &gAppleEventProtocolGuid,
                  NULL,
                  &HandleCount,
                  &HandleBuffer
                  );
  if (!EFI_ERROR (Status)) {
    for (Index = 0; Index < HandleCount; ++Index) {
      Status = gBS->HandleProtocol (HandleBuffer[Index], &gAppleEventProtocolGuid, (VOID **)&AppleEvent);
 
      if (EFI_ERROR (Status)) {
        DEBUG ((
          DEBUG_INFO,
          "CRSCR: gBS->HandleProtocol[%u] AppleEvent returned %r\n",
          (UINT32)Index,
          Status
          ));
        continue;
      }
 
      if (AppleEvent->Revision < APPLE_EVENT_PROTOCOL_REVISION_MINIMUM) {
        DEBUG ((
          DEBUG_INFO,
          "CRSCR: AppleEvent[%u] has outdated revision %u, expected %u\n",
          (UINT32)Index,
          (UINT32)AppleEvent->Revision,
          (UINT32)APPLE_EVENT_PROTOCOL_REVISION_MINIMUM
          ));
        continue;
      }
 
      //
      // Register key handler, which will later determine the combination.
      //
      Status = AppleEvent->RegisterHandler (
                             APPLE_EVENT_TYPE_KEY_DOWN | (mEnableMouseClick ? APPLE_EVENT_TYPE_MOUSE_CLICK : APPLE_EVENT_TYPE_NONE),
                             AppleEventKeyHandler,
                             &AppleEventHandle,
                             NULL
                             );
      if (!EFI_ERROR (Status)) {
        Installed = TRUE;
      }
 
      DEBUG ((
        DEBUG_INFO,
        "CRSCR: AppleEvent->RegisterHandler[%u] returned %r\n",
        (UINT32)Index,
        Status
        ));
    }
 
    gBS->FreePool (HandleBuffer);
  } else {
    DEBUG ((
      DEBUG_INFO,
      "CRSCR: gBS->LocateHandleBuffer AppleEvent returned %r\n",
      Status
      ));
  }
 
  if (!Installed) {
    Status = gBS->LocateHandleBuffer (
                    ByProtocol,
                    &gEfiSimpleTextInputExProtocolGuid,
                    NULL,
                    &HandleCount,
                    &HandleBuffer
                    );
    if (EFI_ERROR (Status)) {
      DEBUG ((
        DEBUG_INFO,
        "CRSCR: gBS->LocateHandleBuffer SimpleTextInEx returned %r\n",
        Status
        ));
      return EFI_UNSUPPORTED;
    }
 
    //
    // Register key notification function for F10.
    //
    SimpleTextInExKeyStroke.Key.ScanCode            = SCAN_F10;
    SimpleTextInExKeyStroke.Key.UnicodeChar         = 0;
    SimpleTextInExKeyStroke.KeyState.KeyShiftState  = 0;
    SimpleTextInExKeyStroke.KeyState.KeyToggleState = 0;
 
    for (Index = 0; Index < HandleCount; ++Index) {
      Status = gBS->HandleProtocol (
                      HandleBuffer[Index],
                      &gEfiSimpleTextInputExProtocolGuid,
                      (VOID **)&SimpleTextInEx
                      );
 
      if (EFI_ERROR (Status)) {
        DEBUG ((
          DEBUG_INFO,
          "CRSCR: gBS->HandleProtocol[%u] SimpleTextInputEx returned %r\n",
          (UINT32)Index,
          Status
          ));
        continue;
      }
 
      Status = SimpleTextInEx->RegisterKeyNotify (
                                 SimpleTextInEx,
                                 &SimpleTextInExKeyStroke,
                                 TakeScreenshot,
                                 &SimpleTextInExHandle
                                 );
 
      if (!EFI_ERROR (Status)) {
        Installed = TRUE;
      }
 
      DEBUG ((
        DEBUG_INFO,
        "CRSCR: SimpleTextInEx->RegisterKeyNotify[%u] returned %r\n",
        (UINT32)Index,
        Status
        ));
    }
 
    gBS->FreePool (HandleBuffer);
  }
 
  return EFI_SUCCESS;
}
 
VOID
EFIAPI
InstallKeyHandlerWrapper (
  IN EFI_EVENT  Event,
  IN VOID       *Context
  )
{
  InstallKeyHandler ();
  gBS->CloseEvent (mProtocolNotification);
}
 
EFI_STATUS
EFIAPI
CrScreenshotDxeEntry (
  IN EFI_HANDLE        ImageHandle,
  IN EFI_SYSTEM_TABLE  *SystemTable
  )
{
  EFI_STATUS                 Status;
  VOID                       *Registration;
  EFI_LOADED_IMAGE_PROTOCOL  *LoadedImage;
  OC_FLEX_ARRAY              *ParsedLoadOptions;
 
  //
  // Parse optional driver params.
  //
  Status = gBS->HandleProtocol (
                  ImageHandle,
                  &gEfiLoadedImageProtocolGuid,
                  (VOID **)&LoadedImage
                  );
 
  if (EFI_ERROR (Status)) {
    return Status;
  }
 
  Status = OcParseLoadOptions (LoadedImage, &ParsedLoadOptions);
  if (!EFI_ERROR (Status)) {
    mEnableMouseClick = OcHasParsedVar (ParsedLoadOptions, L"--enable-mouse-click", OcStringFormatUnicode);
 
    OcFlexArrayFree (&ParsedLoadOptions);
  } else if (Status != EFI_NOT_FOUND) {
    return Status;
  }
 
  Status = InstallKeyHandler ();
  if (EFI_ERROR (Status)) {
    Status = gBS->CreateEvent (
                    EVT_NOTIFY_SIGNAL,
                    TPL_CALLBACK,
                    InstallKeyHandlerWrapper,
                    NULL,
                    &mProtocolNotification
                    );
 
    if (!EFI_ERROR (Status)) {
      gBS->RegisterProtocolNotify (
             &gEfiSimpleTextInputExProtocolGuid,
             mProtocolNotification,
             &Registration
             );
 
      gBS->RegisterProtocolNotify (
             &gAppleEventProtocolGuid,
             mProtocolNotification,
             &Registration
             );
    } else {
      DEBUG ((
        DEBUG_INFO,
        "CRSCR: Cannot create event for keyboard protocol arrivals %r\n",
        Status
        ));
    }
  }
 
  return EFI_SUCCESS;
}