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 struct timespec now;
355 struct timespec remaining;
356
357 for (;;) {
358 if (clock_gettime(CLOCK_MONOTONIC, &now) != 0) {
359 return;
360 }
361
362 if (now.tv_sec > deadline->tv_sec ||
363 (now.tv_sec == deadline->tv_sec && now.tv_nsec >= deadline->tv_nsec)) {
364 return;
365 }
366
367 remaining.tv_sec = deadline->tv_sec - now.tv_sec;
368 remaining.tv_nsec = deadline->tv_nsec - now.tv_nsec;
369 if (remaining.tv_nsec < 0) {
370 remaining.tv_nsec += 1000000000L;
371 remaining.tv_sec--;
372 }
373
374 if (nanosleep(&remaining, NULL) == 0 || errno != EINTR) {
375 return;
376 }
377 }
378}
379
380static int
381cmd_image(const char *prog, int argc, char **argv)
382{
383 struct app app;
384 struct shm_buffer buffer = {.fd = -1};
385 const char *output = "shot.ppm";
386 int ret = 1;
387
388 if (argc > 2) {
389 fprintf(stderr, "usage: %s image [output.ppm]\n", prog);
390 return 1;
391 }
392 if (argc == 2) {
393 output = argv[1];
394 }
395
396 if (!app_connect(&app)) {
397 goto out0;
398 }
399
400 if (!shm_buffer_create(&buffer, app.shm, (int32_t)app.width,
401 (int32_t)app.height, WL_SHM_FORMAT_ARGB8888)) {
402 fprintf(stderr, "swcsnap: failed to allocate shm buffer\n");
403 goto out1;
404 }
405
406 if (!capture_once(&app, &buffer)) {
407 fprintf(stderr, "swcsnap: capture failed (%u)\n", app.failure_reason);
408 goto out2;
409 }
410
411 if (!write_ppm(output, &buffer)) {
412 fprintf(stderr, "swcsnap: failed to write %s\n", output);
413 goto out2;
414 }
415
416 ret = 0;
417
418out2:
419 shm_buffer_destroy(&buffer);
420out1:
421 app_disconnect(&app);
422out0:
423 return ret;
424}
425
426static int
427cmd_record(const char *prog, int argc, char **argv)
428{
429 struct app app;
430 struct shm_buffer buffer = {.fd = -1};
431 unsigned fps = 60;
432 unsigned frames = 600;
433 const char *output_path = "-";
434 int out_fd = STDOUT_FILENO;
435 struct timespec deadline;
436 long frame_ns;
437 int ret = 1;
438 int opt;
439
440 optind = 1;
441 while ((opt = getopt(argc, argv, "o:n:r:")) != -1) {
442 switch (opt) {
443 case 'o':
444 output_path = optarg;
445 break;
446 case 'n':
447 frames = (unsigned)strtoul(optarg, NULL, 10);
448 break;
449 case 'r':
450 fps = (unsigned)strtoul(optarg, NULL, 10);
451 break;
452 default:
453 fprintf(stderr,
454 "usage: %s record [-o output.raw|-] [-n frames] [-r fps]\n",
455 prog);
456 return 1;
457 }
458 }
459
460 if (optind != argc) {
461 fprintf(stderr,
462 "usage: %s record [-o output.raw|-] [-n frames] [-r fps]\n",
463 prog);
464 return 1;
465 }
466
467 if (fps == 0) {
468 fprintf(stderr, "swcsnap: fps must be > 0\n");
469 return 1;
470 }
471
472 if (strcmp(output_path, "-") != 0) {
473 out_fd = open(output_path, O_CREAT | O_TRUNC | O_WRONLY, 0644);
474 if (out_fd < 0) {
475 perror("open");
476 return 1;
477 }
478 }
479
480 if (!app_connect(&app)) {
481 goto out0;
482 }
483
484 if (!shm_buffer_create(&buffer, app.shm, (int32_t)app.width,
485 (int32_t)app.height, WL_SHM_FORMAT_ARGB8888)) {
486 fprintf(stderr, "swcsnap: failed to allocate shm buffer\n");
487 goto out1;
488 }
489
490 fprintf(stderr, "swcsnap: rawvideo stream %ux%u bgra @ %u fps (%u frames)\n",
491 app.width, app.height, fps, frames);
492
493 frame_ns = 1000000000L / (long)fps;
494 clock_gettime(CLOCK_MONOTONIC, &deadline);
495
496 for (unsigned i = 0; i < frames;) {
497 if ((uint32_t)buffer.width != app.width ||
498 (uint32_t)buffer.height != app.height) {
499 shm_buffer_destroy(&buffer);
500 if (!shm_buffer_create(&buffer, app.shm, (int32_t)app.width,
501 (int32_t)app.height,
502 WL_SHM_FORMAT_ARGB8888)) {
503 fprintf(stderr, "swcsnap: failed to reallocate shm buffer\n");
504 goto out2;
505 }
506 fprintf(stderr, "swcsnap: resized to %ux%u\n", app.width,
507 app.height);
508 }
509
510 if (!capture_once(&app, &buffer)) {
511 if (app.capture_failed &&
512 app.failure_reason == SWC_SNAP_FAILURE_REASON_SIZE_MISMATCH) {
513 continue;
514 }
515 fprintf(stderr, "swcsnap: capture failed (%u)\n",
516 app.failure_reason);
517 goto out2;
518 }
519
520 if (write_all(out_fd, buffer.data, buffer.size) < 0) {
521 perror("write");
522 goto out2;
523 }
524
525 i++;
526 timespec_add_ns(&deadline, frame_ns);
527 sleep_until(&deadline);
528 }
529
530 ret = 0;
531
532out2:
533 shm_buffer_destroy(&buffer);
534out1:
535 app_disconnect(&app);
536out0:
537 if (out_fd != STDOUT_FILENO && out_fd >= 0) {
538 close(out_fd);
539 }
540 return ret;
541}
542
543int
544main(int argc, char **argv)
545{
546 if (argc < 2) {
547 usage(stderr, argv[0]);
548 return 1;
549 }
550
551 if (strcmp(argv[1], "image") == 0) {
552 return cmd_image(argv[0], argc - 1, argv + 1);
553 }
554 if (strcmp(argv[1], "record") == 0) {
555 return cmd_record(argv[0], argc - 1, argv + 1);
556 }
557
558 usage(stderr, argv[0]);
559 return 1;
560}
