Make the PNG reader code re-useable

kitty/data-types.c

@@ -211,6 +211,7 @@ extern bool init_shaders(PyObject *module);
 extern bool init_mouse(PyObject *module);
 extern bool init_kittens(PyObject *module);
 extern bool init_logging(PyObject *module);
+extern bool init_png_reader(PyObject *module);
 #ifdef __APPLE__
 extern int init_CoreText(PyObject *);
 extern bool init_cocoa(PyObject *module);
@@ -246,6 +247,7 @@ PyInit_fast_data_types(void) {
         if (!init_shaders(m)) return NULL;
         if (!init_mouse(m)) return NULL;
         if (!init_kittens(m)) return NULL;
+        if (!init_png_reader(m)) return NULL;
 #ifdef __APPLE__
         if (!init_CoreText(m)) return NULL;
         if (!init_cocoa(m)) return NULL;

kitty/glfw-wrapper.h

@@ -1679,7 +1679,7 @@ typedef void (*glfwSetCursorPos_func)(GLFWwindow*, double, double);
 glfwSetCursorPos_func glfwSetCursorPos_impl;
 #define glfwSetCursorPos glfwSetCursorPos_impl
 
-typedef GLFWcursor* (*glfwCreateCursor_func)(const GLFWimage*, int, int);
+typedef GLFWcursor* (*glfwCreateCursor_func)(const GLFWimage*, int, int, int);
 glfwCreateCursor_func glfwCreateCursor_impl;
 #define glfwCreateCursor glfwCreateCursor_impl
 

kitty/graphics.c

@@ -15,8 +15,8 @@
 #include <stdlib.h>
 
 #include <zlib.h>
-#include <png.h>
 #include <structmember.h>
+#include "png-reader.h"
 PyTypeObject GraphicsManager_Type;
 
 #define STORAGE_LIMIT (320 * (1024 * 1024))
@@ -227,87 +227,14 @@ err:
     return ok;
 }
 
-struct fake_file { uint8_t *buf; size_t sz, cur; };
-
 static void
-read_png_from_buffer(png_structp png, png_bytep out, png_size_t length) {
+png_error_handler(const char *code, const char *msg) {
-    struct fake_file *f = png_get_io_ptr(png);
+    set_add_response(code, "%s", msg);
-    if (f) {
-        size_t amt = MIN(length, f->sz - f->cur);
-        memcpy(out, f->buf + f->cur, amt);
-        f->cur += amt;
-    }
-}
-
-static void
-read_png_error_handler(png_structp png_ptr, png_const_charp msg) {
-    jmp_buf *jb;
-    set_add_response("EBADPNG", msg);
-    jb = png_get_error_ptr(png_ptr);
-    if (jb == NULL) fatal("read_png_error_handler: could not retrieve jmp_buf");
-    longjmp(*jb, 1);
-}
-
-static void
-read_png_warn_handler(png_structp UNUSED png_ptr, png_const_charp UNUSED msg) {
-    // ignore warnings
-}
-
-struct png_jmp_data { uint8_t *decompressed; bool ok; png_bytep *row_pointers; int width, height; size_t sz; };
-
-static void
-inflate_png_inner(struct png_jmp_data *d, uint8_t *buf, size_t bufsz) {
-    struct fake_file f = {.buf = buf, .sz = bufsz};
-    png_structp png = NULL;
-    png_infop info = NULL;
-    jmp_buf jb;
-    png = png_create_read_struct(PNG_LIBPNG_VER_STRING, &jb, read_png_error_handler, read_png_warn_handler);
-    if (!png) ABRT(ENOMEM, "Failed to create PNG read structure");
-    info = png_create_info_struct(png);
-    if (!info) ABRT(ENOMEM, "Failed to create PNG info structure");
-
-    if (setjmp(jb)) goto err;
-
-    png_set_read_fn(png, &f, read_png_from_buffer);
-    png_read_info(png, info);
-    png_byte color_type, bit_depth;
-    d->width      = png_get_image_width(png, info);
-    d->height     = png_get_image_height(png, info);
-    color_type = png_get_color_type(png, info);
-    bit_depth  = png_get_bit_depth(png, info);
-
-    // Ensure we get RGBA data out of libpng
-    if (bit_depth == 16) png_set_strip_16(png);
-    if (color_type == PNG_COLOR_TYPE_PALETTE) png_set_palette_to_rgb(png);
-    // PNG_COLOR_TYPE_GRAY_ALPHA is always 8 or 16bit depth.
-    if (color_type == PNG_COLOR_TYPE_GRAY && bit_depth < 8) png_set_expand_gray_1_2_4_to_8(png);
-
-    if (png_get_valid(png, info, PNG_INFO_tRNS)) png_set_tRNS_to_alpha(png);
-
-    // These color_type don't have an alpha channel then fill it with 0xff.
-    if (color_type == PNG_COLOR_TYPE_RGB || color_type == PNG_COLOR_TYPE_GRAY || color_type == PNG_COLOR_TYPE_PALETTE) png_set_filler(png, 0xFF, PNG_FILLER_AFTER);
-
-    if (color_type == PNG_COLOR_TYPE_GRAY || color_type == PNG_COLOR_TYPE_GRAY_ALPHA) png_set_gray_to_rgb(png);
-    png_read_update_info(png, info);
-
-    int rowbytes = png_get_rowbytes(png, info);
-    d->sz = rowbytes * d->height * sizeof(png_byte);
-    d->decompressed = malloc(d->sz + 16);
-    if (d->decompressed == NULL) ABRT(ENOMEM, "Out of memory allocating decompression buffer for PNG");
-    d->row_pointers = malloc(d->height * sizeof(png_bytep));
-    if (d->row_pointers == NULL) ABRT(ENOMEM, "Out of memory allocating row_pointers buffer for PNG");
-    for (int i = 0; i < d->height; i++) d->row_pointers[i] = d->decompressed + i * rowbytes;
-    png_read_image(png, d->row_pointers);
-
-    d->ok = true;
-err:
-    if (png) png_destroy_read_struct(&png, info ? &info : NULL, NULL);
-    return;
 }
 
 static inline bool
 inflate_png(GraphicsManager UNUSED *self, Image *img, uint8_t *buf, size_t bufsz) {
-    struct png_jmp_data d = {0};
+    png_read_data d = {.err_handler=png_error_handler};
     inflate_png_inner(&d, buf, bufsz);
     if (d.ok) {
         free_load_data(&img->load_data);

kitty/png-reader.c

@@ -0,0 +1,126 @@
+/*
+ * png-reader.c
+ * Copyright (C) 2018 Kovid Goyal <kovid at kovidgoyal.net>
+ *
+ * Distributed under terms of the GPL3 license.
+ */
+
+#include "png-reader.h"
+
+
+struct fake_file { const uint8_t *buf; size_t sz, cur; };
+
+static void
+read_png_from_buffer(png_structp png, png_bytep out, png_size_t length) {
+    struct fake_file *f = png_get_io_ptr(png);
+    if (f) {
+        size_t amt = MIN(length, f->sz - f->cur);
+        memcpy(out, f->buf + f->cur, amt);
+        f->cur += amt;
+    }
+}
+
+struct custom_error_handler {
+    jmp_buf jb;
+    png_error_handler_func err_handler;
+};
+
+static void
+read_png_error_handler(png_structp png_ptr, png_const_charp msg) {
+    struct custom_error_handler *eh;
+    eh = png_get_error_ptr(png_ptr);
+    if (eh == NULL) fatal("read_png_error_handler: could not retrieve error handler");
+    if(eh->err_handler) eh->err_handler("EBADPNG", msg);
+    longjmp(eh->jb, 1);
+}
+
+static void
+read_png_warn_handler(png_structp UNUSED png_ptr, png_const_charp UNUSED msg) {
+    // ignore warnings
+}
+
+#define ABRT(code, msg) { if(d->err_handler) d->err_handler(#code, msg); goto err; }
+
+void
+inflate_png_inner(png_read_data *d, const uint8_t *buf, size_t bufsz) {
+    struct fake_file f = {.buf = buf, .sz = bufsz};
+    png_structp png = NULL;
+    png_infop info = NULL;
+    struct custom_error_handler eh = {.err_handler = d->err_handler};
+    png = png_create_read_struct(PNG_LIBPNG_VER_STRING, &eh, read_png_error_handler, read_png_warn_handler);
+    if (!png) ABRT(ENOMEM, "Failed to create PNG read structure");
+    info = png_create_info_struct(png);
+    if (!info) ABRT(ENOMEM, "Failed to create PNG info structure");
+
+    if (setjmp(eh.jb)) goto err;
+
+    png_set_read_fn(png, &f, read_png_from_buffer);
+    png_read_info(png, info);
+    png_byte color_type, bit_depth;
+    d->width      = png_get_image_width(png, info);
+    d->height     = png_get_image_height(png, info);
+    color_type = png_get_color_type(png, info);
+    bit_depth  = png_get_bit_depth(png, info);
+
+    // Ensure we get RGBA data out of libpng
+    if (bit_depth == 16) png_set_strip_16(png);
+    if (color_type == PNG_COLOR_TYPE_PALETTE) png_set_palette_to_rgb(png);
+    // PNG_COLOR_TYPE_GRAY_ALPHA is always 8 or 16bit depth.
+    if (color_type == PNG_COLOR_TYPE_GRAY && bit_depth < 8) png_set_expand_gray_1_2_4_to_8(png);
+
+    if (png_get_valid(png, info, PNG_INFO_tRNS)) png_set_tRNS_to_alpha(png);
+
+    // These color_type don't have an alpha channel then fill it with 0xff.
+    if (color_type == PNG_COLOR_TYPE_RGB || color_type == PNG_COLOR_TYPE_GRAY || color_type == PNG_COLOR_TYPE_PALETTE) png_set_filler(png, 0xFF, PNG_FILLER_AFTER);
+
+    if (color_type == PNG_COLOR_TYPE_GRAY || color_type == PNG_COLOR_TYPE_GRAY_ALPHA) png_set_gray_to_rgb(png);
+    png_read_update_info(png, info);
+
+    int rowbytes = png_get_rowbytes(png, info);
+    d->sz = rowbytes * d->height * sizeof(png_byte);
+    d->decompressed = malloc(d->sz + 16);
+    if (d->decompressed == NULL) ABRT(ENOMEM, "Out of memory allocating decompression buffer for PNG");
+    d->row_pointers = malloc(d->height * sizeof(png_bytep));
+    if (d->row_pointers == NULL) ABRT(ENOMEM, "Out of memory allocating row_pointers buffer for PNG");
+    for (int i = 0; i < d->height; i++) d->row_pointers[i] = d->decompressed + i * rowbytes;
+    png_read_image(png, d->row_pointers);
+
+    d->ok = true;
+err:
+    if (png) png_destroy_read_struct(&png, info ? &info : NULL, NULL);
+    return;
+}
+
+static void
+png_error_handler(const char *code, const char *msg) {
+    PyErr_Format(PyExc_ValueError, "[%s] %s", code, msg);
+}
+
+static PyObject*
+load_png_data(PyObject *self UNUSED, PyObject *args) {
+    Py_ssize_t sz;
+    const char *data;
+    if (!PyArg_ParseTuple(args, "s#", &data, &sz)) return NULL;
+    png_read_data d = {.err_handler=png_error_handler};
+    inflate_png_inner(&d, (const uint8_t*)data, sz);
+    PyObject *ans = NULL;
+    if (d.ok && !PyErr_Occurred()) {
+        ans = PyBytes_FromStringAndSize((const char*)d.decompressed, d.sz);
+    } else {
+        if (!PyErr_Occurred()) PyErr_SetString(PyExc_ValueError, "Unknown error while reading PNG data");
+    }
+    free(d.decompressed);
+    free(d.row_pointers);
+    return ans;
+}
+
+static PyMethodDef module_methods[] = {
+    METHODB(load_png_data, METH_VARARGS),
+    {NULL, NULL, 0, NULL}        /* Sentinel */
+};
+
+bool
+init_png_reader(PyObject *module) {
+    if (PyModule_AddFunctions(module, module_methods) != 0) return false;
+    return true;
+}

kitty/png-reader.h

@@ -0,0 +1,21 @@
+/*
+ * Copyright (C) 2018 Kovid Goyal <kovid at kovidgoyal.net>
+ *
+ * Distributed under terms of the GPL3 license.
+ */
+
+#pragma once
+
+#include "data-types.h"
+#include <png.h>
+typedef void(*png_error_handler_func)(const char*, const char*);
+typedef struct {
+    uint8_t *decompressed;
+    bool ok;
+    png_bytep *row_pointers;
+    int width, height;
+    size_t sz;
+    png_error_handler_func err_handler;
+} png_read_data;
+
+void inflate_png_inner(png_read_data *d, const uint8_t *buf, size_t bufsz);

kitty_tests/graphics.py

@@ -10,7 +10,7 @@ from base64 import standard_b64decode, standard_b64encode
 from io import BytesIO
 
 from kitty.fast_data_types import (
-    parse_bytes, set_send_to_gpu, shm_unlink, shm_write
+    load_png_data, parse_bytes, set_send_to_gpu, shm_unlink, shm_write
 )
 
 from . import BaseTest
@@ -209,8 +209,12 @@ class TestGraphics(BaseTest):
     def test_load_png_simple(self):
         # 1x1 transparent PNG
         png_data = standard_b64decode('iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+P+/HgAFhAJ/wlseKgAAAABJRU5ErkJggg==')
+        expected = b'\x00\xff\xff\x7f'
+        self.ae(load_png_data(png_data), expected)
         s, g, l, sl = load_helpers(self)
-        sl(png_data, f=100, expecting_data=b'\x00\xff\xff\x7f')
+        sl(png_data, f=100, expecting_data=expected)
+        # test error handling for loading bad png data
+        self.assertRaisesRegex(ValueError, '[EBADPNG]', load_png_data, b'dsfsdfsfsfd')
 
     def test_image_put(self):
         cw, ch = 10, 20