From caa9c27554320120bc6ab9348b052399eb59598e Mon Sep 17 00:00:00 2001 From: Kovid Goyal Date: Mon, 17 Jul 2023 09:34:52 +0530 Subject: [PATCH] Read serialized signatures --- kittens/transfer/algorithm.c | 198 ++++++++++++++++++++++++++++++++--- 1 file changed, 184 insertions(+), 14 deletions(-) diff --git a/kittens/transfer/algorithm.c b/kittens/transfer/algorithm.c index e8cf9d135..ba352da39 100644 --- a/kittens/transfer/algorithm.c +++ b/kittens/transfer/algorithm.c @@ -8,6 +8,7 @@ #include "data-types.h" #include "binary.h" +#include "kitty-uthash.h" #include #include @@ -141,6 +142,40 @@ rolling_checksum_add_one_byte(rolling_checksum *self, uint8_t first_byte, uint8_ // Python interface {{{ +typedef struct buffer { + uint8_t *data; + size_t len, cap; +} buffer; + +static bool +ensure_space(buffer *b, size_t amt) { + const size_t len = b->len; + if (amt > 0 && b->cap < len + amt) { + size_t newcap = MAX(b->cap * 2, len + (amt * 2)); + b->data = realloc(b->data, newcap); + if (b->data == NULL) { PyErr_NoMemory(); return false; } + b->cap = newcap; + } + return true; +} + +static bool +write_to_buffer(buffer *b, void *data, size_t len) { + if (!ensure_space(b, len)) return false; + memcpy(b->data + b->len, data, len); + b->len += len; + return true; +} + +static void +shift_left(buffer *b, size_t amt) { + if (amt > b->len) amt = b->len; + if (amt > 0) { + b->len -= amt; + memmove(b->data, b->data + amt, b->len); + } +} + // Patcher {{{ typedef struct { PyObject_HEAD @@ -148,7 +183,7 @@ typedef struct { uint64_t signature_idx; size_t block_size; Rsync rsync; - struct { uint8_t *data; size_t len, cap; } buf, block_buf; + buffer buf, block_buf; PyObject *block_buf_view; } Patcher; @@ -177,6 +212,7 @@ Patcher_dealloc(PyObject *self) { if (p->buf.data) free(p->buf.data); Py_CLEAR(p->block_buf_view); if (p->block_buf.data) free(p->block_buf.data); + free_rsync(&p->rsync); Py_TYPE(self)->tp_free(self); } @@ -310,14 +346,7 @@ apply_delta_data(Patcher *self, PyObject *args) { PyObject *read, *write; FREE_BUFFER_AFTER_FUNCTION Py_buffer data = {0}; if (!PyArg_ParseTuple(args, "y*OO", &data, &read, &write)) return NULL; - if (self->buf.cap < (size_t)data.len + self->buf.len) { - size_t newcap = MAX(self->buf.cap * 2, (size_t)data.len); - self->buf.data = realloc(self->buf.data, newcap); - if (self->buf.data == NULL) { return PyErr_NoMemory(); } - self->buf.cap = newcap; - } - memcpy(self->buf.data + self->buf.len, data.buf, data.len); - self->buf.len += (size_t)data.len; + if (!write_to_buffer(&self->buf, data.buf, data.len)) return NULL; size_t pos = 0; Operation op = {0}; while (pos < self->buf.len) { @@ -326,10 +355,7 @@ apply_delta_data(Patcher *self, PyObject *args) { pos += consumed; if (!apply_op(self, op, read, write)) break; } - if (pos > 0) { - self->buf.len -= pos; - memmove(self->buf.data, self->buf.data + pos, self->buf.len); - } + shift_left(&self->buf, pos); if (PyErr_Occurred()) return NULL; Py_RETURN_NONE; } @@ -355,6 +381,150 @@ PyTypeObject Patcher_Type = { }; // }}} Patcher +// Differ {{{ +typedef struct Signature { uint64_t index, strong_hash; } Signature; + +typedef struct SignatureMap { + int weak_hash; + Signature sig, *weak_hash_collisions; + size_t len, cap; + UT_hash_handle hh; +} SignatureMap; + +typedef struct Differ { + PyObject_HEAD + rolling_checksum rc; + uint64_t signature_idx; + Rsync rsync; + bool signature_header_parsed; + buffer buf; + SignatureMap *signature_map; +} Differ; + +static int +Differ_init(PyObject *s, PyObject *args, PyObject *kwds) { + Patcher *self = (Patcher*)s; + static char *kwlist[] = {NULL}; + if (!PyArg_ParseTupleAndKeywords(args, kwds, "", kwlist)) return -1; + const char *err = init_rsync(&self->rsync, default_block_size, 0, 0); + if (err != NULL) { PyErr_SetString(RsyncError, err); return -1; } + return 0; +} + +static void +free_sigmap(SignatureMap *map) { + SignatureMap *current, *tmp; + HASH_ITER(hh, map, current, tmp) { + HASH_DEL(map, current); + free(current->weak_hash_collisions); + free(current); + } +} + +static void +Differ_dealloc(PyObject *self) { + Differ *p = (Differ*)self; + if (p->buf.data) free(p->buf.data); + free_rsync(&p->rsync); + if (p->signature_map) free_sigmap(p->signature_map); + Py_TYPE(self)->tp_free(self); +} + +static void +parse_signature_header(Differ *self) { + if (self->buf.len < 12) return; + uint8_t *p = self->buf.data; + uint32_t x; + if ((x = le16(p)) != 0) { + PyErr_Format(RsyncError, "Invalid version in signature header: %u", x); return; + } p += 2; + if ((x = le16(p)) != 0) { + PyErr_Format(RsyncError, "Invalid checksum type in signature header: %u", x); return; + } p += 2; + if ((x = le16(p)) != 0) { + PyErr_Format(RsyncError, "Invalid strong hash type in signature header: %u", x); return; + } p += 2; + if ((x = le16(p)) != 0) { + PyErr_Format(RsyncError, "Invalid weak hash type in signature header: %u", x); return; + } p += 2; + const char *err = init_rsync(&self->rsync, le32(p), 0, 0); + if (err != NULL) { PyErr_SetString(RsyncError, err); return; } + p += 4; + shift_left(&self->buf, p - self->buf.data); + self->signature_header_parsed = true; +} + +static bool +add_collision(SignatureMap *sm, Signature s) { + if (sm->cap < sm->len + 1) { + size_t new_cap = MAX(sm->cap * 2, 8); + sm->weak_hash_collisions = realloc(sm->weak_hash_collisions, new_cap); + if (!sm->weak_hash_collisions) { PyErr_NoMemory(); return false; } + sm->cap = new_cap; + } + sm->weak_hash_collisions[sm->len++] = s; + return true; +} + +static size_t +parse_signature_block(Differ *self, uint8_t *data, size_t len) { + if (len < 20) return 0; + int weak_hash = le32(data + 8); + SignatureMap *sm; + HASH_FIND_INT(self->signature_map, &weak_hash, sm); + if (sm == NULL) { + sm = calloc(1, sizeof(SignatureMap)); + if (sm == NULL) { PyErr_NoMemory(); return 0; } + sm->weak_hash = weak_hash; + sm->sig.index = le64(data); + sm->sig.strong_hash = le64(data+12); + HASH_ADD_INT(self->signature_map, weak_hash, sm); + } else { + if (!add_collision(sm, (Signature){.index=le64(data), .strong_hash=le64(data+12)})) return 0; + } + return 20; +} + +static PyObject* +add_signature_data(Differ *self, PyObject *args) { + FREE_BUFFER_AFTER_FUNCTION Py_buffer data = {0}; + if (!PyArg_ParseTuple(args, "y*", &data)) return NULL; + if (!write_to_buffer(&self->buf, data.buf, data.len)) return NULL; + if (!self->signature_header_parsed) { + parse_signature_header(self); + if (PyErr_Occurred()) return NULL; + if (!self->signature_header_parsed) { Py_RETURN_NONE; } + } + size_t pos = 0; + while (pos < self->buf.len) { + size_t consumed = parse_signature_block(self, self->buf.data + pos, self->buf.len - pos); + if (!consumed) { break; } + pos += consumed; + } + shift_left(&self->buf, pos); + if (PyErr_Occurred()) return NULL; + Py_RETURN_NONE; +} + +static PyMethodDef Differ_methods[] = { + METHODB(add_signature_data, METH_VARARGS), + {NULL} /* Sentinel */ +}; + + +PyTypeObject Differ_Type = { + PyVarObject_HEAD_INIT(NULL, 0) + .tp_name = "rsync.Differ", + .tp_basicsize = sizeof(Differ), + .tp_dealloc = Differ_dealloc, + .tp_flags = Py_TPFLAGS_DEFAULT, + .tp_doc = "Differ", + .tp_methods = Differ_methods, + .tp_new = PyType_GenericNew, + .tp_init = Differ_init, +}; +// }}} Differ + // Hasher {{{ typedef struct { PyObject_HEAD @@ -538,7 +708,7 @@ exec_module(PyObject *m) { PyModule_AddObject(m, "RsyncError", RsyncError); #define T(which) if (PyType_Ready(& which##_Type) < 0) return -1; Py_INCREF(&which##_Type);\ if (PyModule_AddObject(m, #which, (PyObject *) &which##_Type) < 0) return -1; - T(Hasher); T(Patcher); + T(Hasher); T(Patcher); T(Differ); #undef T return 0; }