client/cros/graphics: Introduce DRM-based screenshotting.
Directly take screenshots by mapping the framebuffers without using X11. This
effectively permits Freon-based screenshots during testing.
To test: Try to run something that uses screenshots. ui_SystemTray for
example. The test fails, but the screenshot succeeds.
BUG=chromium:412530
TEST=See above
Change-Id: If9160d3a2781f30be1edd0af5f9ce07ac115bd4a
Signed-off-by: Corbin Simpson <simpsoco@chromium.org>
Reviewed-on: https://chromium-review.googlesource.com/222973
Reviewed-by: Ilja Friedel <ihf@chromium.org>
diff --git a/client/cros/graphics/drm.py b/client/cros/graphics/drm.py
new file mode 100644
index 0000000..d53e5d1
--- /dev/null
+++ b/client/cros/graphics/drm.py
@@ -0,0 +1,367 @@
+"""
+A wrapper around the Direct Rendering Manager (DRM) library, which itself is a
+wrapper around the Direct Rendering Interface (DRI) between the kernel and
+userland.
+
+Since we are masochists, we use ctypes instead of cffi to load libdrm and
+access several symbols within it. We use Python's file descriptor and mmap
+wrappers.
+
+At some point in the future, cffi could be used, for approximately the same
+cost in lines of code.
+"""
+
+from ctypes import *
+import mmap
+import os
+
+from PIL import Image
+
+
+class DrmVersion(Structure):
+ """
+ The version of a DRM node.
+ """
+
+ _fields_ = [
+ ("version_major", c_int),
+ ("version_minor", c_int),
+ ("version_patchlevel", c_int),
+ ("name_len", c_int),
+ ("name", c_char_p),
+ ("date_len", c_int),
+ ("date", c_char_p),
+ ("desc_len", c_int),
+ ("desc", c_char_p),
+ ]
+
+ _l = None
+
+ def __repr__(self):
+ return "%s %d.%d.%d (%s) (%s)" % (
+ self.name,
+ self.version_major,
+ self.version_minor,
+ self.version_patchlevel,
+ self.desc,
+ self.date,
+ )
+
+ def __del__(self):
+ if self._l:
+ self._l.drmFreeVersion(self)
+
+
+class DrmModeResources(Structure):
+ """
+ Resources associated with setting modes on a DRM node.
+ """
+
+ _fields_ = [
+ ("count_fbs", c_int),
+ ("fbs", POINTER(c_uint)),
+ ("count_crtcs", c_int),
+ ("crtcs", POINTER(c_uint)),
+ # XXX incomplete struct!
+ ]
+
+ _fd = None
+ _l = None
+
+ def __repr__(self):
+ return "<DRM mode resources>"
+
+ def __del__(self):
+ if self._l:
+ self._l.drmModeFreeResources(self)
+
+ def getCrtc(self, index):
+ """
+ Obtain the CRTC at a given index.
+
+ @param index: The CRTC to get.
+ """
+
+ if not 0 <= index < self.count_crtcs:
+ raise IndexError("CRTC index out of range")
+
+ crtc = self._l.drmModeGetCrtc(self._fd, self.crtcs[index]).contents
+ crtc._fd = self._fd
+ crtc._l = self._l
+ return crtc
+
+
+class DrmModeCrtc(Structure):
+ """
+ A DRM modesetting CRTC.
+ """
+
+ _fields_ = [
+ ("crtc_id", c_uint),
+ ("buffer_id", c_uint),
+ # XXX incomplete struct!
+ ]
+
+ _fd = None
+ _l = None
+
+ def __repr__(self):
+ return "<CRTC (%d)>" % self.crtc_id
+
+ def __del__(self):
+ if self._l:
+ self._l.drmModeFreeCrtc(self)
+
+ def hasFb(self):
+ """
+ Whether this CRTC has an associated framebuffer.
+ """
+
+ return self.buffer_id != 0
+
+ def fb(self):
+ """
+ Obtain the framebuffer, if one is associated.
+ """
+
+ if self.hasFb():
+ fb = self._l.drmModeGetFB(self._fd, self.buffer_id).contents
+ fb._fd = self._fd
+ fb._l = self._l
+ return fb
+
+
+class drm_mode_map_dumb(Structure):
+ """
+ Request a mapping of a modesetting buffer.
+
+ The map will be "dumb;" it will be accessible via mmap() but very slow.
+ """
+
+ _fields_ = [
+ ("handle", c_uint),
+ ("pad", c_uint),
+ ("offset", c_ulonglong),
+ ]
+
+
+# This constant is not defined in any one header; it is the pieced-together
+# incantation for the ioctl that performs dumb mappings. I would love for this
+# to not have to be here, but it can't be imported from any header easily.
+DRM_IOCTL_MODE_MAP_DUMB = 0xc01064b3
+
+
+class DrmModeFB(Structure):
+ """
+ A DRM modesetting framebuffer.
+ """
+
+ _fields_ = [
+ ("fb_id", c_uint),
+ ("width", c_uint),
+ ("height", c_uint),
+ ("pitch", c_uint),
+ ("bpp", c_uint),
+ ("depth", c_uint),
+ ("handle", c_uint),
+ ]
+
+ _l = None
+ _map = None
+
+ def __repr__(self):
+ s = "<Framebuffer (%dx%d (pitch %d bytes), %d bits/pixel, depth %d)"
+ vitals = s % (
+ self.width,
+ self.height,
+ self.pitch,
+ self.bpp,
+ self.depth,
+ )
+ if self._map:
+ tail = " (mapped)>"
+ else:
+ tail = ">"
+ return vitals + tail
+
+ def __del__(self):
+ if self._l:
+ self._l.drmModeFreeFB(self)
+
+ def map(self):
+ """
+ Map the framebuffer.
+ """
+
+ if self._map:
+ return
+
+ mapDumb = drm_mode_map_dumb()
+ mapDumb.handle = self.handle
+
+ rv = self._l.drmIoctl(self._fd, DRM_IOCTL_MODE_MAP_DUMB,
+ pointer(mapDumb))
+ if rv:
+ raise IOError(rv, os.strerror(rv))
+
+ size = self.pitch * self.height
+
+ # mmap.mmap() has a totally different order of arguments in Python
+ # compared to C; check the documentation before altering this
+ # incantation.
+ self._map = mmap.mmap(self._fd, size, flags=mmap.MAP_SHARED,
+ prot=mmap.PROT_READ, offset=mapDumb.offset)
+
+ def unmap(self):
+ """
+ Unmap the framebuffer.
+ """
+
+ if self._map:
+ self._map.close()
+ self._map = None
+
+
+def loadDRM():
+ """
+ Load a handle to libdrm.
+
+ In addition to loading, this function also configures the argument and
+ return types of functions.
+ """
+
+ l = cdll.LoadLibrary("libdrm.so")
+
+ l.drmGetVersion.argtypes = [c_int]
+ l.drmGetVersion.restype = POINTER(DrmVersion)
+
+ l.drmFreeVersion.argtypes = [POINTER(DrmVersion)]
+ l.drmFreeVersion.restype = None
+
+ l.drmModeGetResources.argtypes = [c_int]
+ l.drmModeGetResources.restype = POINTER(DrmModeResources)
+
+ l.drmModeFreeResources.argtypes = [POINTER(DrmModeResources)]
+ l.drmModeFreeResources.restype = None
+
+ l.drmModeGetCrtc.argtypes = [c_int, c_uint]
+ l.drmModeGetCrtc.restype = POINTER(DrmModeCrtc)
+
+ l.drmModeFreeCrtc.argtypes = [POINTER(DrmModeCrtc)]
+ l.drmModeFreeCrtc.restype = None
+
+ l.drmModeGetFB.argtypes = [c_int, c_uint]
+ l.drmModeGetFB.restype = POINTER(DrmModeFB)
+
+ l.drmModeFreeFB.argtypes = [POINTER(DrmModeFB)]
+ l.drmModeFreeFB.restype = None
+
+ l.drmIoctl.argtypes = [c_int, c_ulong, c_voidp]
+ l.drmIoctl.restype = c_int
+
+ return l
+
+
+class DRM(object):
+ """
+ A DRM node.
+ """
+
+ def __init__(self, library, fd):
+ self._l = library
+ self._fd = fd
+
+ def __repr__(self):
+ return "<DRM (FD %d)>" % self._fd
+
+ @classmethod
+ def fromHandle(cls, handle):
+ """
+ Create a node from a file handle.
+
+ @param handle: A file-like object backed by a file descriptor.
+ """
+
+ self = cls(loadDRM(), handle.fileno())
+ # We must keep the handle alive, and we cannot trust the caller to
+ # keep it alive for us.
+ self._handle = handle
+ return self
+
+ def version(self):
+ """
+ Obtain the version.
+ """
+
+ v = self._l.drmGetVersion(self._fd).contents
+ v._l = self._l
+ return v
+
+ def resources(self):
+ """
+ Obtain the modesetting resources.
+ """
+
+ r = self._l.drmModeGetResources(self._fd).contents
+ r._fd = self._fd
+ r._l = self._l
+ return r
+
+
+def drmFromMinor(minor):
+ """
+ Given a DRM node number, open the corresponding node.
+
+ @param minor: The number of the minor node to open.
+ """
+
+ path = "/dev/dri/card%d" % minor
+ handle = open(path)
+ return DRM.fromHandle(handle)
+
+
+def _bgrx24(i):
+ b = ord(next(i))
+ g = ord(next(i))
+ r = ord(next(i))
+ next(i)
+ return r, g, b
+
+
+def _screenshot(image, fb):
+ fb.map()
+ m = fb._map
+ lineLength = fb.width * fb.bpp // 8
+ pitch = fb.pitch
+ pixels = []
+
+ if fb.depth == 24:
+ unformat = _bgrx24
+ else:
+ raise RuntimeError("Couldn't unformat FB: %r" % fb)
+
+ for y in range(fb.height):
+ offset = y * pitch
+ m.seek(offset)
+ channels = m.read(lineLength)
+ ichannels = iter(channels)
+ for x in range(fb.width):
+ rgb = unformat(ichannels)
+ image.putpixel((x, y), rgb)
+
+ fb.unmap()
+
+ return pixels
+
+
+def screenshot():
+ """
+ Take a screenshot, returning an image object.
+ """
+
+ d = drmFromMinor(0)
+ fb = d.resources().getCrtc(0).fb()
+ image = Image.new("RGB", (fb.width, fb.height))
+ pixels = _screenshot(image, fb)
+
+ return image