1#define _POSIX_C_SOURCE 200809L
2
3#include <errno.h>
4#include <fcntl.h>
5#include <stdbool.h>
6#include <stdint.h>
7#include <stdio.h>
8#include <stdlib.h>
9#include <string.h>
10#include <sys/mman.h>
11#include <time.h>
12#include <unistd.h>
13
14#include <wayland-client.h>
15
16#include "swc_snap-client-protocol.h"
17
18enum { CAPTURE_FLAGS = SWC_SNAP_FLAGS_OVERLAY_CURSOR };
19
20struct shm_buffer {
21 struct wl_shm_pool *pool;
22 struct wl_buffer *buffer;
23 void *data;
24 size_t size;
25 int fd;
26 int32_t width;
27 int32_t height;
28 int32_t stride;
29};
30
31struct app {
32 struct wl_display *display;
33 struct wl_registry *registry;
34 struct wl_shm *shm;
35 struct swc_snap *snap;
36 uint32_t width;
37 uint32_t height;
38 bool have_buffer_info;
39 bool capture_ready;
40 bool capture_failed;
41 uint32_t failure_reason;
42};
43
44static void
45usage(FILE *stream, const char *argv0)
46{
47 fprintf(stream,
48 "usage:\n"
49 " %s image [output.ppm]\n"
50 " %s record [-o output.raw|-] [-n frames] [-r fps]\n",
51 argv0, argv0);
52}
53
54static ssize_t
55write_all(int fd, const void *data, size_t size)
56{
57 const uint8_t *p = data;
58 size_t written = 0;
59
60 while (written < size) {
61 ssize_t rc = write(fd, p + written, size - written);
62 if (rc < 0) {
63 if (errno == EINTR) {
64 continue;
65 }
66 return -1;
67 }
68 written += (size_t)rc;
69 }
70
71 return (ssize_t)written;
72}
73
74static int
75create_shm_file(size_t size)
76{
77 char path[] = "/tmp/swcsnap-XXXXXX";
78 int fd = mkstemp(path);
79
80 if (fd < 0) {
81 return -1;
82 }
83
84 unlink(path);
85
86 if (ftruncate(fd, (off_t)size) < 0) {
87 close(fd);
88 return -1;
89 }
90
91 return fd;
92}
93
94static bool
95shm_buffer_create(struct shm_buffer *buffer, struct wl_shm *shm,
96 int32_t width, int32_t height, uint32_t format)
97{
98 size_t size;
99
100 memset(buffer, 0, sizeof(*buffer));
101 buffer->fd = -1;
102 buffer->width = width;
103 buffer->height = height;
104 buffer->stride = width * 4;
105 size = (size_t)buffer->stride * buffer->height;
106 buffer->size = size;
107
108 buffer->fd = create_shm_file(size);
109 if (buffer->fd < 0) {
110 return false;
111 }
112
113 buffer->data = mmap(NULL, size, PROT_READ | PROT_WRITE, MAP_SHARED,
114 buffer->fd, 0);
115 if (buffer->data == MAP_FAILED) {
116 close(buffer->fd);
117 buffer->fd = -1;
118 buffer->data = NULL;
119 return false;
120 }
121
122 buffer->pool = wl_shm_create_pool(shm, buffer->fd, (int32_t)size);
123 if (!buffer->pool) {
124 munmap(buffer->data, size);
125 close(buffer->fd);
126 buffer->fd = -1;
127 buffer->data = NULL;
128 return false;
129 }
130
131 buffer->buffer = wl_shm_pool_create_buffer(
132 buffer->pool, 0, width, height, buffer->stride, format);
133 if (!buffer->buffer) {
134 wl_shm_pool_destroy(buffer->pool);
135 munmap(buffer->data, size);
136 close(buffer->fd);
137 buffer->pool = NULL;
138 buffer->fd = -1;
139 buffer->data = NULL;
140 return false;
141 }
142
143 return true;
144}
145
146static void
147shm_buffer_destroy(struct shm_buffer *buffer)
148{
149 if (buffer->buffer) {
150 wl_buffer_destroy(buffer->buffer);
151 }
152 if (buffer->pool) {
153 wl_shm_pool_destroy(buffer->pool);
154 }
155 if (buffer->data) {
156 munmap(buffer->data, buffer->size);
157 }
158 if (buffer->fd >= 0) {
159 close(buffer->fd);
160 }
161 memset(buffer, 0, sizeof(*buffer));
162 buffer->fd = -1;
163}
164
165static bool
166write_ppm(const char *path, const struct shm_buffer *buffer)
167{
168 FILE *f;
169 uint8_t rgb[3];
170
171 f = fopen(path, "wb");
172 if (!f) {
173 perror("fopen");
174 return false;
175 }
176
177 if (fprintf(f, "P6\n%d %d\n255\n", buffer->width, buffer->height) < 0) {
178 fclose(f);
179 return false;
180 }
181
182 for (int32_t y = 0; y < buffer->height; y++) {
183 const uint8_t *row =
184 (const uint8_t *)buffer->data + (size_t)y * buffer->stride;
185 for (int32_t x = 0; x < buffer->width; x++) {
186 const uint8_t *px = row + (size_t)x * 4;
187
188 rgb[0] = px[2];
189 rgb[1] = px[1];
190 rgb[2] = px[0];
191 if (fwrite(rgb, 1, sizeof(rgb), f) != sizeof(rgb)) {
192 fclose(f);
193 return false;
194 }
195 }
196 }
197
198 return fclose(f) == 0;
199}
200
201static void
202handle_snap_buffer(void *data, struct swc_snap *snap, uint32_t width,
203 uint32_t height, uint32_t format)
204{
205 struct app *app = data;
206 (void)snap;
207
208 app->width = width;
209 app->height = height;
210 app->have_buffer_info = true;
211}
212
213static void
214handle_snap_ready(void *data, struct swc_snap *snap, uint32_t tv_sec,
215 uint32_t tv_nsec)
216{
217 struct app *app = data;
218 (void)snap;
219 (void)tv_sec;
220 (void)tv_nsec;
221
222 app->capture_ready = true;
223}
224
225static void
226handle_snap_failed(void *data, struct swc_snap *snap, uint32_t reason)
227{
228 struct app *app = data;
229 (void)snap;
230
231 app->capture_failed = true;
232 app->failure_reason = reason;
233}
234
235static const struct swc_snap_listener snap_listener = {
236 .buffer = handle_snap_buffer,
237 .ready = handle_snap_ready,
238 .failed = handle_snap_failed,
239};
240
241static void
242handle_global(void *data, struct wl_registry *registry, uint32_t name,
243 const char *interface, uint32_t version)
244{
245 struct app *app = data;
246
247 if (strcmp(interface, wl_shm_interface.name) == 0) {
248 app->shm = wl_registry_bind(registry, name, &wl_shm_interface, 1);
249 return;
250 }
251
252 if (strcmp(interface, "swc_snap") == 0) {
253 uint32_t bind_version = version < 2 ? version : 2;
254
255 app->snap =
256 wl_registry_bind(registry, name, &swc_snap_interface, bind_version);
257 if (app->snap) {
258 swc_snap_add_listener(app->snap, &snap_listener, app);
259 }
260 }
261}
262
263static void
264handle_global_remove(void *data, struct wl_registry *registry, uint32_t name)
265{
266 (void)data;
267 (void)registry;
268 (void)name;
269}
270
271static const struct wl_registry_listener registry_listener = {
272 .global = handle_global,
273 .global_remove = handle_global_remove,
274};
275
276static bool
277app_connect(struct app *app)
278{
279 memset(app, 0, sizeof(*app));
280
281 app->display = wl_display_connect(NULL);
282 if (!app->display) {
283 fprintf(stderr, "swcsnap: failed to connect to Wayland display\n");
284 return false;
285 }
286
287 app->registry = wl_display_get_registry(app->display);
288 wl_registry_add_listener(app->registry, ®istry_listener, app);
289
290 if (wl_display_roundtrip(app->display) < 0 ||
291 wl_display_roundtrip(app->display) < 0) {
292 fprintf(stderr, "swcsnap: roundtrip failed\n");
293 return false;
294 }
295
296 if (!app->shm || !app->snap || !app->have_buffer_info) {
297 fprintf(stderr, "swcsnap: missing wl_shm/swc_snap metadata\n");
298 return false;
299 }
300
301 return true;
302}
303
304static void
305app_disconnect(struct app *app)
306{
307 if (app->snap) {
308 swc_snap_destroy(app->snap);
309 }
310 if (app->shm) {
311 wl_shm_destroy(app->shm);
312 }
313 if (app->registry) {
314 wl_registry_destroy(app->registry);
315 }
316 if (app->display) {
317 wl_display_disconnect(app->display);
318 }
319 memset(app, 0, sizeof(*app));
320}
321
322static bool
323capture_once(struct app *app, struct shm_buffer *buffer)
324{
325 app->capture_ready = false;
326 app->capture_failed = false;
327 app->failure_reason = 0;
328
329 swc_snap_capture(app->snap, buffer->buffer, CAPTURE_FLAGS);
330 wl_display_flush(app->display);
331
332 while (!app->capture_ready && !app->capture_failed) {
333 if (wl_display_dispatch(app->display) < 0) {
334 return false;
335 }
336 }
337
338 return !app->capture_failed;
339}
340
341static void
342timespec_add_ns(struct timespec *ts, long ns)
343{
344 ts->tv_nsec += ns;
345 while (ts->tv_nsec >= 1000000000L) {
346 ts->tv_nsec -= 1000000000L;
347 ts->tv_sec += 1;
348 }
349}
350
351static void
352sleep_until(const struct timespec *deadline)
353{
354 int rc;
355
356 while ((rc = clock_nanosleep(CLOCK_MONOTONIC, TIMER_ABSTIME, deadline,
357 NULL)) != 0) {
358 if (rc != EINTR) {
359 return;
360 }
361 }
362}
363
364static int
365cmd_image(const char *prog, int argc, char **argv)
366{
367 struct app app;
368 struct shm_buffer buffer = {.fd = -1};
369 const char *output = "shot.ppm";
370 int ret = 1;
371
372 if (argc > 2) {
373 fprintf(stderr, "usage: %s image [output.ppm]\n", prog);
374 return 1;
375 }
376 if (argc == 2) {
377 output = argv[1];
378 }
379
380 if (!app_connect(&app)) {
381 goto out0;
382 }
383
384 if (!shm_buffer_create(&buffer, app.shm, (int32_t)app.width,
385 (int32_t)app.height, WL_SHM_FORMAT_ARGB8888)) {
386 fprintf(stderr, "swcsnap: failed to allocate shm buffer\n");
387 goto out1;
388 }
389
390 if (!capture_once(&app, &buffer)) {
391 fprintf(stderr, "swcsnap: capture failed (%u)\n", app.failure_reason);
392 goto out2;
393 }
394
395 if (!write_ppm(output, &buffer)) {
396 fprintf(stderr, "swcsnap: failed to write %s\n", output);
397 goto out2;
398 }
399
400 ret = 0;
401
402out2:
403 shm_buffer_destroy(&buffer);
404out1:
405 app_disconnect(&app);
406out0:
407 return ret;
408}
409
410static int
411cmd_record(const char *prog, int argc, char **argv)
412{
413 struct app app;
414 struct shm_buffer buffer = {.fd = -1};
415 unsigned fps = 60;
416 unsigned frames = 600;
417 const char *output_path = "-";
418 int out_fd = STDOUT_FILENO;
419 struct timespec deadline;
420 long frame_ns;
421 int ret = 1;
422 int opt;
423
424 optind = 1;
425 while ((opt = getopt(argc, argv, "o:n:r:")) != -1) {
426 switch (opt) {
427 case 'o':
428 output_path = optarg;
429 break;
430 case 'n':
431 frames = (unsigned)strtoul(optarg, NULL, 10);
432 break;
433 case 'r':
434 fps = (unsigned)strtoul(optarg, NULL, 10);
435 break;
436 default:
437 fprintf(stderr,
438 "usage: %s record [-o output.raw|-] [-n frames] [-r fps]\n",
439 prog);
440 return 1;
441 }
442 }
443
444 if (optind != argc) {
445 fprintf(stderr,
446 "usage: %s record [-o output.raw|-] [-n frames] [-r fps]\n",
447 prog);
448 return 1;
449 }
450
451 if (fps == 0) {
452 fprintf(stderr, "swcsnap: fps must be > 0\n");
453 return 1;
454 }
455
456 if (strcmp(output_path, "-") != 0) {
457 out_fd = open(output_path, O_CREAT | O_TRUNC | O_WRONLY, 0644);
458 if (out_fd < 0) {
459 perror("open");
460 return 1;
461 }
462 }
463
464 if (!app_connect(&app)) {
465 goto out0;
466 }
467
468 if (!shm_buffer_create(&buffer, app.shm, (int32_t)app.width,
469 (int32_t)app.height, WL_SHM_FORMAT_ARGB8888)) {
470 fprintf(stderr, "swcsnap: failed to allocate shm buffer\n");
471 goto out1;
472 }
473
474 fprintf(stderr, "swcsnap: rawvideo stream %ux%u bgra @ %u fps (%u frames)\n",
475 app.width, app.height, fps, frames);
476
477 frame_ns = 1000000000L / (long)fps;
478 clock_gettime(CLOCK_MONOTONIC, &deadline);
479
480 for (unsigned i = 0; i < frames;) {
481 if ((uint32_t)buffer.width != app.width ||
482 (uint32_t)buffer.height != app.height) {
483 shm_buffer_destroy(&buffer);
484 if (!shm_buffer_create(&buffer, app.shm, (int32_t)app.width,
485 (int32_t)app.height,
486 WL_SHM_FORMAT_ARGB8888)) {
487 fprintf(stderr, "swcsnap: failed to reallocate shm buffer\n");
488 goto out2;
489 }
490 fprintf(stderr, "swcsnap: resized to %ux%u\n", app.width,
491 app.height);
492 }
493
494 if (!capture_once(&app, &buffer)) {
495 if (app.capture_failed &&
496 app.failure_reason == SWC_SNAP_FAILURE_REASON_SIZE_MISMATCH) {
497 continue;
498 }
499 fprintf(stderr, "swcsnap: capture failed (%u)\n",
500 app.failure_reason);
501 goto out2;
502 }
503
504 if (write_all(out_fd, buffer.data, buffer.size) < 0) {
505 perror("write");
506 goto out2;
507 }
508
509 i++;
510 timespec_add_ns(&deadline, frame_ns);
511 sleep_until(&deadline);
512 }
513
514 ret = 0;
515
516out2:
517 shm_buffer_destroy(&buffer);
518out1:
519 app_disconnect(&app);
520out0:
521 if (out_fd != STDOUT_FILENO && out_fd >= 0) {
522 close(out_fd);
523 }
524 return ret;
525}
526
527int
528main(int argc, char **argv)
529{
530 if (argc < 2) {
531 usage(stderr, argv[0]);
532 return 1;
533 }
534
535 if (strcmp(argv[1], "image") == 0) {
536 return cmd_image(argv[0], argc - 1, argv + 1);
537 }
538 if (strcmp(argv[1], "record") == 0) {
539 return cmd_record(argv[0], argc - 1, argv + 1);
540 }
541
542 usage(stderr, argv[0]);
543 return 1;
544}
