clients: rewrite screenshot.c for new protocol

The functionality of this screenshooting helper client is kept exactly
the same as before: if you have multiple outputs, some transformed, some
scale, in any layout, this will create a "multi-image" where the
framebuffer (the physical image) of each output is pasted into a row of
images in the order the outputs were advertised thrugh wl_registry.
Output transform or scale are not accounted for. If you have a monitor
rotated sideways, the screenshot will have the image of that monitor
reverse-sideways.

Otherwise the client is almost completely re-written, so trying to read
the diff is not that useful.

The old screenshooting protocol is replaced with the new
weston-output-capture protocol. This makes it unnecessary to listen for
wl_output information (since we do not handle output transform or scale
anyway).

The buffer sizes and formats are dictated by the compositor, which also
means we cannot hardcode the format. Hence, use Pixman for the blitting,
in case it needs to do format conversion. It is good to get rid of
hand-crafted pixel data manipulation code too.

For that reason we also need a pixel format database to convert between
DRM fourcc, wl_shm and Pixman codes. We link to libweston to borrow its
database instead of inventing another partial copy of it. It's weird to
use compositor library private API in a client, but better than the
alternative.

The original code had no tear-down code at all. Now, if everything
succeeds, the program ends with no unfreed memory according to ASan. If
something fails, it still YOLO's it (doesn't free stuff). That's how far
my pedantry carried.

I also did not bother taking output transform or scale into account,
since the old code did not either. It would be nice to create a seamless
image of the desktop with shots rotated and scaled to align, in the max
scale over all outputs. Meh.

Signed-off-by: Pekka Paalanen <pekka.paalanen@collabora.com>

clients/meson.build

@@ -364,10 +364,14 @@ if get_option('shell-desktop')
 	exe_shooter = executable(
 		'weston-screenshooter',
 		'screenshot.c',
-		weston_screenshooter_client_protocol_h,
+		weston_output_capture_client_protocol_h,
-		weston_screenshooter_protocol_c,
+		weston_output_capture_protocol_c,
 		include_directories: common_inc,
-		dependencies: dep_toytoolkit,
+		dependencies: [
+			dep_toytoolkit,
+			dep_libweston_private, # for pixel-formats.h
+			dep_pixman,
+		],
 		install_dir: get_option('bindir'),
 		install: true
 	)

clients/screenshot.c

@@ -1,5 +1,6 @@
 /*
  * Copyright © 2008 Kristian Høgsberg
+ * Copyright 2022 Collabora, Ltd.
  *
  * Permission is hereby granted, free of charge, to any person obtaining a
  * copy of this software and associated documentation files (the "Software"),
@@ -24,6 +25,7 @@
 #include "config.h"
 
 #include <stdint.h>
+#include <stdbool.h>
 #include <errno.h>
 #include <stdlib.h>
 #include <stdio.h>
@@ -33,25 +35,49 @@
 #include <limits.h>
 #include <sys/param.h>
 #include <sys/mman.h>
+#include <pixman.h>
 #include <cairo.h>
+#include <assert.h>
 
 #include <wayland-client.h>
-#include "weston-screenshooter-client-protocol.h"
+#include "weston-output-capture-client-protocol.h"
 #include "shared/os-compatibility.h"
 #include "shared/xalloc.h"
 #include "shared/file-util.h"
+#include "pixel-formats.h"
 
-/* The screenshooter is a good example of a custom object exposed by
+struct screenshooter_app {
- * the compositor and serves as a test bed for implementing client
+	struct wl_registry *registry;
- * side marshalling outside libwayland.so */
+	struct wl_shm *shm;
+	struct weston_capture_v1 *capture_factory;
 
+	struct wl_list output_list; /* struct screenshooter_output::link */
+
+	bool retry;
+	bool failed;
+	int waitcount;
+};
+
+struct screenshooter_buffer {
+	size_t len;
+	void *data;
+	struct wl_buffer *wl_buffer;
+	pixman_image_t *image;
+};
 
 struct screenshooter_output {
-	struct wl_output *output;
+	struct screenshooter_app *app;
-	struct wl_buffer *buffer;
+	struct wl_list link; /* struct screenshooter_app::output_list */
-	int width, height, offset_x, offset_y;
+
-	void *data;
+	struct wl_output *wl_output;
-	struct wl_list link;
+	int offset_x, offset_y;
+
+	struct weston_capture_source_v1 *source;
+
+	int buffer_width;
+	int buffer_height;
+	const struct pixel_format_info *fmt;
+	struct screenshooter_buffer *buffer;
 };
 
 struct buffer_size {
@@ -61,97 +87,204 @@ struct buffer_size {
 	int max_x, max_y;
 };
 
-struct screenshooter_data {
+static struct screenshooter_buffer *
-	struct wl_shm *shm;
+screenshot_create_shm_buffer(struct screenshooter_app *app,
-	struct wl_list output_list;
+			     size_t width, size_t height,
+			     const struct pixel_format_info *fmt)
+{
+	struct screenshooter_buffer *buffer;
+	struct wl_shm_pool *pool;
+	int fd;
+	size_t bytes_pp;
+	size_t stride;
 
-	struct weston_screenshooter *screenshooter;
+	assert(width > 0);
-	int buffer_copy_done;
+	assert(height > 0);
-};
+	assert(fmt && fmt->bpp > 0);
+	assert(fmt->pixman_format);
 
+	buffer = xzalloc(sizeof *buffer);
+
+	bytes_pp = fmt->bpp / 8;
+	stride = width * bytes_pp;
+	buffer->len = stride * height;
+
+	assert(width == stride / bytes_pp);
+	assert(height == buffer->len / stride);
+
+	fd = os_create_anonymous_file(buffer->len);
+	if (fd < 0) {
+		fprintf(stderr, "creating a buffer file for %zd B failed: %s\n",
+			buffer->len, strerror(errno));
+		free(buffer);
+		return NULL;
+	}
+
+	buffer->data = mmap(NULL, buffer->len, PROT_READ | PROT_WRITE,
+			    MAP_SHARED, fd, 0);
+	if (buffer->data == MAP_FAILED) {
+		fprintf(stderr, "mmap failed: %s\n", strerror(errno));
+		close(fd);
+		free(buffer);
+		return NULL;
+	}
+
+	pool = wl_shm_create_pool(app->shm, fd, buffer->len);
+	close(fd);
+	buffer->wl_buffer =
+		wl_shm_pool_create_buffer(pool, 0, width, height, stride,
+					  pixel_format_get_shm_format(fmt));
+	wl_shm_pool_destroy(pool);
+
+	buffer->image = pixman_image_create_bits(fmt->pixman_format,
+						 width, height,
+						 buffer->data, stride);
+	abort_oom_if_null(buffer->image);
+
+	return buffer;
+}
 
 static void
-display_handle_geometry(void *data,
+screenshooter_buffer_destroy(struct screenshooter_buffer *buffer)
-			struct wl_output *wl_output,
+{
-			int x,
+	if (!buffer)
-			int y,
+		return;
-			int physical_width,
+
-			int physical_height,
+	pixman_image_unref(buffer->image);
-			int subpixel,
+	munmap(buffer->data, buffer->len);
-			const char *make,
+	wl_buffer_destroy(buffer->wl_buffer);
-			const char *model,
+	free(buffer);
-			int transform)
+}
+
+static void
+capture_source_handle_format(void *data,
+			     struct weston_capture_source_v1 *proxy,
+			     uint32_t drm_format)
+{
+	struct screenshooter_output *output = data;
+
+	assert(output->source == proxy);
+
+	output->fmt = pixel_format_get_info(drm_format);
+}
+
+static void
+capture_source_handle_size(void *data,
+			   struct weston_capture_source_v1 *proxy,
+			   int32_t width, int32_t height)
+{
+	struct screenshooter_output *output = data;
+
+	assert(width > 0);
+	assert(height > 0);
+
+	output->buffer_width = width;
+	output->buffer_height = height;
+}
+
+static void
+capture_source_handle_complete(void *data,
+			       struct weston_capture_source_v1 *proxy)
+{
+	struct screenshooter_output *output = data;
+
+	output->app->waitcount--;
+}
+
+static void
+capture_source_handle_retry(void *data,
+			    struct weston_capture_source_v1 *proxy)
+{
+	struct screenshooter_output *output = data;
+
+	output->app->waitcount--;
+	output->app->retry = true;
+}
+
+static void
+capture_source_handle_failed(void *data,
+			     struct weston_capture_source_v1 *proxy,
+			     const char *msg)
+{
+	struct screenshooter_output *output = data;
+
+	output->app->waitcount--;
+	output->app->failed = true;
+
+	if (msg)
+		fprintf(stderr, "Output capture error: %s\n", msg);
+}
+
+static const struct weston_capture_source_v1_listener capture_source_handlers = {
+	.format = capture_source_handle_format,
+	.size = capture_source_handle_size,
+	.complete = capture_source_handle_complete,
+	.retry = capture_source_handle_retry,
+	.failed = capture_source_handle_failed,
+};
+
+static void
+create_output(struct screenshooter_app *app, uint32_t output_name, uint32_t version)
 {
 	struct screenshooter_output *output;
 
-	output = wl_output_get_user_data(wl_output);
+	version = MIN(version, 4);
+	output = xzalloc(sizeof *output);
+	output->app = app;
+	output->wl_output = wl_registry_bind(app->registry, output_name,
+					     &wl_output_interface, version);
+	abort_oom_if_null(output->wl_output);
 
-	if (wl_output == output->output) {
+	output->source = weston_capture_v1_create(app->capture_factory,
-		output->offset_x = x;
+						  output->wl_output,
-		output->offset_y = y;
+						  WESTON_CAPTURE_V1_SOURCE_FRAMEBUFFER);
-	}
+	abort_oom_if_null(output->source);
+	weston_capture_source_v1_add_listener(output->source,
+					      &capture_source_handlers, output);
+
+	wl_list_insert(&app->output_list, &output->link);
 }
 
 static void
-display_handle_mode(void *data,
+destroy_output(struct screenshooter_output *output)
-		    struct wl_output *wl_output,
-		    uint32_t flags,
-		    int width,
-		    int height,
-		    int refresh)
 {
-	struct screenshooter_output *output;
+	weston_capture_source_v1_destroy(output->source);
 
-	output = wl_output_get_user_data(wl_output);
+	if (wl_output_get_version(output->wl_output) >= WL_OUTPUT_RELEASE_SINCE_VERSION)
+		wl_output_release(output->wl_output);
+	else
+		wl_output_destroy(output->wl_output);
 
-	if (wl_output == output->output && (flags & WL_OUTPUT_MODE_CURRENT)) {
+	screenshooter_buffer_destroy(output->buffer);
-		output->width = width;
+	wl_list_remove(&output->link);
-		output->height = height;
+	free(output);
-	}
 }
 
-static const struct wl_output_listener output_listener = {
-	display_handle_geometry,
-	display_handle_mode
-};
-
-static void
-screenshot_done(void *data, struct weston_screenshooter *screenshooter)
-{
-	struct screenshooter_data *sh_data = data;
-	sh_data->buffer_copy_done = 1;
-}
-
-static const struct weston_screenshooter_listener screenshooter_listener = {
-	screenshot_done
-};
-
 static void
 handle_global(void *data, struct wl_registry *registry,
 	      uint32_t name, const char *interface, uint32_t version)
 {
-	static struct screenshooter_output *output;
+	struct screenshooter_app *app = data;
-	struct screenshooter_data *sh_data = data;
 
-	if (strcmp(interface, "wl_output") == 0) {
+	if (strcmp(interface, wl_output_interface.name) == 0) {
-		output = xmalloc(sizeof *output);
+		create_output(app, name, version);
-		output->output = wl_registry_bind(registry, name,
+	} else if (strcmp(interface, wl_shm_interface.name) == 0) {
-						  &wl_output_interface, 1);
+		app->shm = wl_registry_bind(registry, name, &wl_shm_interface, 1);
-		wl_list_insert(&sh_data->output_list, &output->link);
+		/*
-		wl_output_add_listener(output->output, &output_listener, output);
+		 * Not listening for format advertisements,
-	} else if (strcmp(interface, "wl_shm") == 0) {
+		 * weston_capture_source_v1.format event tells us what to use.
-		sh_data->shm = wl_registry_bind(registry, name, &wl_shm_interface, 1);
+		 */
-	} else if (strcmp(interface, "weston_screenshooter") == 0) {
+	} else if (strcmp(interface, weston_capture_v1_interface.name) == 0) {
-		sh_data->screenshooter = wl_registry_bind(registry, name,
+		app->capture_factory = wl_registry_bind(registry, name,
-							  &weston_screenshooter_interface,
+							&weston_capture_v1_interface,
-							  1);
+							1);
 	}
 }
 
 static void
 handle_global_remove(void *data, struct wl_registry *registry, uint32_t name)
 {
-	/* XXX: unimplemented */
+	/* Dynamic output removals will just fail the respective shot. */
 }
 
 static const struct wl_registry_listener registry_listener = {
@@ -159,80 +292,52 @@ static const struct wl_registry_listener registry_listener = {
 	handle_global_remove
 };
 
-static struct wl_buffer *
+static void
-screenshot_create_shm_buffer(int width, int height, void **data_out,
+screenshooter_output_capture(struct screenshooter_output *output)
-			     struct wl_shm *shm)
 {
-	struct wl_shm_pool *pool;
+	screenshooter_buffer_destroy(output->buffer);
-	struct wl_buffer *buffer;
+	output->buffer = screenshot_create_shm_buffer(output->app,
-	int fd, size, stride;
+						      output->buffer_width,
-	void *data;
+						      output->buffer_height,
+						      output->fmt);
+	abort_oom_if_null(output->buffer);
 
-	stride = width * 4;
+	weston_capture_source_v1_capture(output->source,
-	size = stride * height;
+					 output->buffer->wl_buffer);
-
+	output->app->waitcount++;
-	fd = os_create_anonymous_file(size);
-	if (fd < 0) {
-		fprintf(stderr, "creating a buffer file for %d B failed: %s\n",
-			size, strerror(errno));
-		return NULL;
-	}
-
-	data = mmap(NULL, size, PROT_READ | PROT_WRITE, MAP_SHARED, fd, 0);
-	if (data == MAP_FAILED) {
-		fprintf(stderr, "mmap failed: %s\n", strerror(errno));
-		close(fd);
-		return NULL;
-	}
-
-	pool = wl_shm_create_pool(shm, fd, size);
-	close(fd);
-	buffer = wl_shm_pool_create_buffer(pool, 0, width, height, stride,
-					   WL_SHM_FORMAT_XRGB8888);
-	wl_shm_pool_destroy(pool);
-
-	*data_out = data;
-
-	return buffer;
 }
 
 static void
 screenshot_write_png(const struct buffer_size *buff_size,
 		     struct wl_list *output_list)
 {
-	int output_stride, buffer_stride, i;
+	pixman_image_t *shot;
 	cairo_surface_t *surface;
-	void *data, *d, *s;
+	struct screenshooter_output *output;
-	struct screenshooter_output *output, *next;
 	FILE *fp;
 	char filepath[PATH_MAX];
 
-	buffer_stride = buff_size->width * 4;
+	shot = pixman_image_create_bits(PIXMAN_a8r8g8b8,
+					buff_size->width, buff_size->height,
+					NULL, 0);
+	abort_oom_if_null(shot);
 
-	data = xmalloc(buffer_stride * buff_size->height);
+	wl_list_for_each(output, output_list, link) {
-	if (!data)
+		pixman_image_composite32(PIXMAN_OP_SRC,
-		return;
+					 output->buffer->image, /* src */
-
+					 NULL, /* mask */
-	wl_list_for_each_safe(output, next, output_list, link) {
+					 shot, /* dest */
-		output_stride = output->width * 4;
+					 0, 0, /* src x,y */
-		s = output->data;
+					 0, 0, /* mask x,y */
-		d = data + (output->offset_y - buff_size->min_y) * buffer_stride +
+					 output->offset_x, output->offset_y, /* dst x,y */
-			   (output->offset_x - buff_size->min_x) * 4;
+					 output->buffer_width, output->buffer_height);
-
-		for (i = 0; i < output->height; i++) {
-			memcpy(d, s, output_stride);
-			d += buffer_stride;
-			s += output_stride;
-		}
-
-		free(output);
 	}
 
-	surface = cairo_image_surface_create_for_data(data,
+	surface = cairo_image_surface_create_for_data((void *)pixman_image_get_data(shot),
 						      CAIRO_FORMAT_ARGB32,
-						      buff_size->width,
+						      pixman_image_get_width(shot),
-						      buff_size->height,
+						      pixman_image_get_height(shot),
-						      buffer_stride);
+						      pixman_image_get_stride(shot));
 
 	fp = file_create_dated(getenv("XDG_PICTURES_DIR"), "wayland-screenshot-",
 			       ".png", filepath, sizeof(filepath));
@@ -241,11 +346,12 @@ screenshot_write_png(const struct buffer_size *buff_size,
 		cairo_surface_write_to_png(surface, filepath);
 	}
 	cairo_surface_destroy(surface);
-	free(data);
+	pixman_image_unref(shot);
 }
 
 static int
-screenshot_set_buffer_size(struct buffer_size *buff_size, struct wl_list *output_list)
+screenshot_set_buffer_size(struct buffer_size *buff_size,
+			   struct wl_list *output_list)
 {
 	struct screenshooter_output *output;
 	buff_size->min_x = buff_size->min_y = INT_MAX;
@@ -254,16 +360,16 @@ screenshot_set_buffer_size(struct buffer_size *buff_size, struct wl_list *output
 
 	wl_list_for_each_reverse(output, output_list, link) {
 		output->offset_x = position;
-		position += output->width;
+		position += output->buffer_width;
 	}
 
 	wl_list_for_each(output, output_list, link) {
 		buff_size->min_x = MIN(buff_size->min_x, output->offset_x);
 		buff_size->min_y = MIN(buff_size->min_y, output->offset_y);
 		buff_size->max_x =
-			MAX(buff_size->max_x, output->offset_x + output->width);
+			MAX(buff_size->max_x, output->offset_x + output->buffer_width);
 		buff_size->max_y =
-			MAX(buff_size->max_y, output->offset_y + output->height);
+			MAX(buff_size->max_y, output->offset_y + output->buffer_height);
 	}
 
 	if (buff_size->max_x <= buff_size->min_x ||
@@ -276,13 +382,16 @@ screenshot_set_buffer_size(struct buffer_size *buff_size, struct wl_list *output
 	return 0;
 }
 
-int main(int argc, char *argv[])
+int
+main(int argc, char *argv[])
 {
 	struct wl_display *display;
-	struct wl_registry *registry;
 	struct screenshooter_output *output;
+	struct screenshooter_output *tmp_output;
 	struct buffer_size buff_size = {};
-	struct screenshooter_data sh_data = {};
+	struct screenshooter_app app = {};
+
+	wl_list_init(&app.output_list);
 
 	display = wl_display_connect(NULL);
 	if (display == NULL) {
@@ -291,39 +400,52 @@ int main(int argc, char *argv[])
 		return -1;
 	}
 
-	wl_list_init(&sh_data.output_list);
+	app.registry = wl_display_get_registry(display);
-	registry = wl_display_get_registry(display);
+	wl_registry_add_listener(app.registry, &registry_listener, &app);
-	wl_registry_add_listener(registry, &registry_listener, &sh_data);
+
-	wl_display_dispatch(display);
+	/* Process wl_registry advertisements */
 	wl_display_roundtrip(display);
-	if (sh_data.screenshooter == NULL) {
+
-		fprintf(stderr, "display doesn't support screenshooter\n");
+	if (!app.shm) {
+		fprintf(stderr, "Error: display does not support wl_shm\n");
+		return -1;
+	}
+	if (!app.capture_factory) {
+		fprintf(stderr, "Error: display does not support weston_capture_v1\n");
 		return -1;
 	}
 
-	weston_screenshooter_add_listener(sh_data.screenshooter,
+	/* Process initial events for wl_output and weston_capture_source_v1 */
-					  &screenshooter_listener,
+	wl_display_roundtrip(display);
-					  &sh_data);
 
-	if (screenshot_set_buffer_size(&buff_size, &sh_data.output_list))
+	do {
-		return -1;
+		app.retry = false;
 
+		wl_list_for_each(output, &app.output_list, link)
+			screenshooter_output_capture(output);
 
-	wl_list_for_each(output, &sh_data.output_list, link) {
+		while (app.waitcount > 0 && !app.failed) {
-		output->buffer =
+			if (wl_display_dispatch(display) < 0)
-			screenshot_create_shm_buffer(output->width,
+				app.failed = true;
-						     output->height,
+			assert(app.waitcount >= 0);
-						     &output->data,
+		}
-						     sh_data.shm);
+	} while (app.retry && !app.failed);
-		weston_screenshooter_take_shot(sh_data.screenshooter,
+
-					       output->output,
+	if (!app.failed) {
-					       output->buffer);
+		if (screenshot_set_buffer_size(&buff_size, &app.output_list) < 0)
-		sh_data.buffer_copy_done = 0;
+			return -1;
-		while (!sh_data.buffer_copy_done)
+		screenshot_write_png(&buff_size, &app.output_list);
-			wl_display_roundtrip(display);
+	} else {
+		fprintf(stderr, "Error: screenshot or protocol failure\n");
 	}
 
-	screenshot_write_png(&buff_size, &sh_data.output_list);
+	wl_list_for_each_safe(output, tmp_output, &app.output_list, link)
+		destroy_output(output);
+
+	weston_capture_v1_destroy(app.capture_factory);
+	wl_shm_destroy(app.shm);
+	wl_registry_destroy(app.registry);
+	wl_display_disconnect(display);
 
 	return 0;
 }