From 405cef8d28d2e46971552efcbecbf0f0c537c654 Mon Sep 17 00:00:00 2001 From: Jack Robison Date: Wed, 28 Dec 2022 14:18:46 -0500 Subject: [PATCH] ResumableSHA256 --- hub/common.py | 38 ++++++++++++++++++++++++++++++++++++++ 1 file changed, 38 insertions(+) diff --git a/hub/common.py b/hub/common.py index 9304689..424dd8f 100644 --- a/hub/common.py +++ b/hub/common.py @@ -7,12 +7,14 @@ import logging import logging.handlers import typing import collections +from ctypes import cast, memmove, POINTER, c_void_p from bisect import insort_right from collections import deque from decimal import Decimal from typing import Iterable, Deque from asyncio import get_event_loop, Event from prometheus_client import Counter +from rehash.structs import EVPobject from hub.schema.tags import clean_tags from hub.schema.url import normalize_name from hub.error import TooManyClaimSearchParametersError @@ -1059,3 +1061,39 @@ async def asyncify_for_loop(gen, ticks_per_sleep: int = 1000): yield item if cnt % ticks_per_sleep == 0: await async_sleep(0) + + +class ResumableSHA256: + __slots__ = ['_hasher'] + + def __init__(self, state: typing.Optional[bytes] = None): + self._hasher = hashlib.sha256() + if state is not None: + ctx = self._get_evp_md_ctx() + ctx_size = ctx.digest.contents.ctx_size + memmove(ctx.md_data, state, ctx_size) + + def _get_evp_md_ctx(self): + c_evp_obj = cast(c_void_p(id(self._hasher)), POINTER(EVPobject)) + if hasattr(c_evp_obj.contents.ctx, "contents"): + return c_evp_obj.contents.ctx.contents + else: + return c_evp_obj.contents.ctx + + def get_state(self) -> bytes: + ctx = self._get_evp_md_ctx() + ctx_size = ctx.digest.contents.ctx_size + hasher_state = ctx.md_data[:ctx_size] + return hasher_state + + def copy(self): + return ResumableSHA256(self.get_state()) + + def update(self, data: bytes): + self._hasher.update(data) + + def digest(self): + return self._hasher.digest() + + def hexdigest(self): + return self.digest().hex()