fix: Protect native handles with locks

src/c2pa/c2pa.py

@@ -13,6 +13,7 @@
 
 # Version: 0.37.8
 
+import contextlib
 import ctypes
 import enum
 import json
@@ -266,6 +267,8 @@ def __init__(self):
         self._lifecycle_state = LifecycleState.UNINITIALIZED
         self._handle = None
         self._op_lock = threading.RLock()
+        self._inflight = 0
+        self._pending_teardown = None
         record_owner_pid(self)
 
     def _lock(self):
@@ -296,6 +299,39 @@ def _lock(self):
                 pass
         return lock
 
+    @contextlib.contextmanager
+    def _native_call(self):
+        """Hold the handle valid across a native call that goes back
+        and forth to native layers.
+
+        Calls that pass a Stream to the native library run caller-supplied
+        callbacks, so the lock cannot be held across them: the callback may
+        re-enter this API on another thread and deadlock. Instead the call is
+        counted as in flight, and a teardown arriving meanwhile records its
+        intent rather than freeing. The last caller out performs the free.
+
+        The resource is marked closed as soon as the teardown is recorded, so
+        a caller that closed it cannot keep using it while the free is
+        pending.
+        """
+        with self._lock():
+            self._ensure_valid_state()
+            self._inflight = getattr(self, '_inflight', 0) + 1
+        try:
+            yield
+        finally:
+            with self._lock():
+                self._inflight -= 1
+                pending = (self._pending_teardown
+                           if self._inflight == 0 else None)
+                if pending is not None:
+                    self._pending_teardown = None
+            # Released the lock before the free: _teardown takes it again,
+            # and keeping the two acquisitions separate means the counter
+            # update is never held across the release work.
+            if pending is not None:
+                self._teardown(pending)
+
     @staticmethod
     def _free_native_ptr(ptr):
         """Free a native pointer by passing it to c2pa_free.
@@ -356,6 +392,16 @@ def _teardown(self, free_handle: bool):
         thread's state check and its use of the handle in a native call.
         """
         with self._lock():
+            if getattr(self, '_inflight', 0) > 0:
+                # A native call is running that re-enters caller Python and
+                # is still using this handle. Record the intent; whichever
+                # caller leaves _native_call last performs the free. Mark the
+                # resource closed now so it cannot be used while the free is
+                # pending.
+                self._pending_teardown = free_handle
+                self._lifecycle_state = LifecycleState.CLOSED
+                return
+
             if is_foreign_process(self):
                 self._handle = None
                 self._lifecycle_state = LifecycleState.CLOSED
@@ -1735,13 +1781,22 @@ def __init__(
                         check=lambda r: r != 0)
 
                 if signer is not None:
-                    signer._ensure_valid_state()
-                    # A rejected signer is retained, not closed and leaked.
-                    self._signer_callback_cb = signer._callback_cb
-                    signer._consume_no_replacement(
-                        lambda h: _lib.c2pa_context_builder_set_signer(
-                            nb._handle, h),
-                        "Failed to set signer on Context: {}")
+                    # The signer's own in-flight guard: this hands its handle
+                    # to native, so a signer.close() on another thread must
+                    # not free it between the state check and the call. The
+                    # guard also makes the check and the consume atomic.
+                    #
+                    # _consume_no_replacement tears the signer down from
+                    # inside this region. A teardown recorded while the guard
+                    # is held is deferred and performed as the guard unwinds,
+                    # which is still before __init__ returns.
+                    with signer._native_call():
+                        # A rejected signer is retained, not closed and leaked.
+                        self._signer_callback_cb = signer._callback_cb
+                        signer._consume_no_replacement(
+                            lambda h: _lib.c2pa_context_builder_set_signer(
+                                nb._handle, h),
+                            "Failed to set signer on Context: {}")
                     self._has_signer = True
 
                 context_ptr = nb._consume_into(
@@ -2658,12 +2713,19 @@ def _init_from_context(self, context, format_or_path,
             self._own_stream = Stream(stream)
 
         try:
-            # Adopt before the consuming call: _consume_and_swap needs an
-            # active resource, and cleanup then owns the pointer either way.
-            self._create_and_activate(
-                lambda: _lib.c2pa_reader_from_context(
-                    context.execution_context),
-                Reader._ERROR_MESSAGES['reader_error'])
+            # The Context is caller-supplied and may be shared, so its handle
+            # needs its own in-flight guard across the native call: the
+            # execution_context property validates and returns the handle, and
+            # without the guard a context.close() on another thread could free
+            # it before c2pa_reader_from_context reads it.
+            with context._native_call():
+                # Adopt before the consuming call: _consume_and_swap needs an
+                # active resource, and cleanup then owns the pointer either
+                # way.
+                self._create_and_activate(
+                    lambda: _lib.c2pa_reader_from_context(
+                        context.execution_context),
+                    Reader._ERROR_MESSAGES['reader_error'])
 
             if manifest_data is not None:
                 manifest_array = (
@@ -2702,6 +2764,10 @@ def _init_attrs(self):
         # Tracks a file we opened ourselves and must close later.
         self._backing_file = None
 
+        # Fragment streams handed to the native reader by with_fragment,
+        # which it keeps reading from for the rest of its life.
+        self._fragment_streams = []
+
         # Caches for manifest JSON string and parsed data.
         # These are invalidated when with_fragment() is called.
         self._manifest_json_str_cache = None
@@ -2725,6 +2791,12 @@ def _close_streams(self):
                 logger.warning("Failed to close Reader backing file")
             finally:
                 self._backing_file = None
+        for fragment in getattr(self, '_fragment_streams', []):
+            try:
+                fragment.close()
+            except Exception:
+                logger.warning("Failed to close Reader fragment stream")
+        self._fragment_streams = []
 
     def _release(self):
         """Release Reader-specific resources (caches, stream, backing file).
@@ -2787,19 +2859,38 @@ def with_fragment(self, format: Optional[str], stream,
                 cannot be retried: create a new one instead of reusing this
                 instance.
         """
-        self._ensure_valid_state()
-
         format_arg = _format_ffi_arg(_encode_format(format, "Reader"))
 
-        with Stream(stream) as main_obj, Stream(fragment_stream) as frag_obj:
-            self._consume_and_swap(
-                lambda handle: _lib.c2pa_reader_with_fragment(
-                    handle,
-                    format_arg,
-                    main_obj._stream,
-                    frag_obj._stream,
-                ),
-                Reader._ERROR_MESSAGES['fragment_error'])
+        # The native reader keeps reading through both streams after this
+        # returns, so they are owned here and released by _release() rather
+        # than at the end of a with block.
+        main_obj = Stream(stream)
+        frag_obj = Stream(fragment_stream)
+        try:
+            with self._native_call():
+                self._consume_and_swap(
+                    lambda handle: _lib.c2pa_reader_with_fragment(
+                        handle,
+                        format_arg,
+                        main_obj._stream,
+                        frag_obj._stream,
+                    ),
+                    Reader._ERROR_MESSAGES['fragment_error'])
+        except Exception:
+            main_obj.close()
+            frag_obj.close()
+            raise
+
+        # Replace the streams this reader owned, closing the previous ones so
+        # repeated calls do not accumulate them.
+        previous = self._own_stream
+        self._own_stream = main_obj
+        self._fragment_streams.append(frag_obj)
+        if previous is not None and previous is not main_obj:
+            try:
+                previous.close()
+            except Exception:
+                logger.warning("Failed to close previous Reader stream")
 
         # Invalidate caches: processing a new BMFF fragment updates the native
         # reader's state, which can change the manifest data it returns.
@@ -3000,10 +3091,8 @@ def resource_to_stream(self, uri: str, stream: Any) -> int:
         Raises:
             C2paError: If there was an error writing the resource to stream
         """
-        self._ensure_valid_state()
-
         uri_str = uri.encode('utf-8')
-        with Stream(stream) as stream_obj:
+        with self._native_call(), Stream(stream) as stream_obj:
             result = _lib.c2pa_reader_resource_to_stream(
                 self._handle, uri_str, stream_obj._stream)
 
@@ -3451,11 +3540,17 @@ def _init_from_context(self, context, json_str):
         if not context.is_valid:
             raise C2paError("Context is not valid")
 
-        # Adopt before the consuming call: _consume_and_swap needs an
-        # active resource, and cleanup then owns the pointer either way.
-        self._create_and_activate(
-            lambda: _lib.c2pa_builder_from_context(context.execution_context),
-            Builder._ERROR_MESSAGES['builder_error'])
+        # The Context is caller-supplied and may be shared, so its handle
+        # needs its own in-flight guard across the native call: without it a
+        # context.close() on another thread frees the handle between the
+        # is_valid check and c2pa_builder_from_context reading it.
+        with context._native_call():
+            # Adopt before the consuming call: _consume_and_swap needs an
+            # active resource, and cleanup then owns the pointer either way.
+            self._create_and_activate(
+                lambda: _lib.c2pa_builder_from_context(
+                    context.execution_context),
+                Builder._ERROR_MESSAGES['builder_error'])
 
         self._consume_and_swap(
             lambda handle: _lib.c2pa_builder_with_definition(
@@ -3558,10 +3653,8 @@ def add_resource(self, uri: str, stream: Any):
         Raises:
             C2paError: If there was an error adding the resource
         """
-        self._ensure_valid_state()
-
         uri_bytes = _to_utf8_bytes(uri, "resource URI")
-        with Stream(stream) as stream_obj:
+        with self._native_call(), Stream(stream) as stream_obj:
             result = _lib.c2pa_builder_add_resource(
                 self._handle, uri_bytes, stream_obj._stream)
 
@@ -3622,7 +3715,7 @@ def add_ingredient_from_stream(
         ingredient_str = _to_utf8_bytes(ingredient_json, "ingredient JSON")
         format_str = _to_utf8_bytes(format, "ingredient format")
 
-        with Stream(source) as source_stream:
+        with self._native_call(), Stream(source) as source_stream:
             result = (
                 _lib.c2pa_builder_add_ingredient_from_stream(
                     self._handle,
@@ -3671,9 +3764,7 @@ def to_archive(self, stream: Any) -> None:
         Raises:
             C2paError: If there was an error writing the archive
         """
-        self._ensure_valid_state()
-
-        with Stream(stream) as stream_obj:
+        with self._native_call(), Stream(stream) as stream_obj:
             result = _lib.c2pa_builder_to_archive(
                 self._handle, stream_obj._stream)
 
@@ -3698,7 +3789,7 @@ def write_ingredient_archive(self, ingredient_id: str, stream: Any) -> None:
 
         ingredient_id_str = _to_utf8_bytes(ingredient_id, "ingredient_id")
 
-        with Stream(stream) as stream_obj:
+        with self._native_call(), Stream(stream) as stream_obj:
             result = _lib.c2pa_builder_write_ingredient_archive(
                 self._handle, ingredient_id_str, stream_obj._stream)
 
@@ -3718,9 +3809,7 @@ def add_ingredient_from_archive(self, stream: Any) -> None:
         Raises:
             C2paError: If there was an error reading the archive
         """
-        self._ensure_valid_state()
-
-        with Stream(stream) as stream_obj:
+        with self._native_call(), Stream(stream) as stream_obj:
             result = _lib.c2pa_builder_add_ingredient_from_archive(
                 self._handle, stream_obj._stream)
 
@@ -3749,7 +3838,7 @@ def with_archive(self, stream: Any) -> 'Builder':
         """
         self._ensure_valid_state()
 
-        with Stream(stream) as stream_obj:
+        with self._native_call(), Stream(stream) as stream_obj:
             self._consume_and_swap(
                 lambda handle: _lib.c2pa_builder_with_archive(
                     handle, stream_obj._stream),
@@ -3797,23 +3886,36 @@ def _sign_internal(
         manifest_bytes_ptr = ctypes.POINTER(ctypes.c_ubyte)()
 
         try:
-            if signer is not None:
-                result = _lib.c2pa_builder_sign(
-                    self._handle,
-                    format_arg,
-                    source_stream._stream,
-                    dest_stream._stream,
-                    signer._handle,
-                    ctypes.byref(manifest_bytes_ptr)
-                )
-            else:
-                result = _lib.c2pa_builder_sign_context(
-                    self._handle,
-                    format_arg,
-                    source_stream._stream,
-                    dest_stream._stream,
-                    ctypes.byref(manifest_bytes_ptr),
-                )
+            # _native_call covers the signing call only. The close() below is
+            # deliberately outside it, so the deferred teardown it records is
+            # performed on the way out rather than being deferred forever.
+            with self._native_call():
+                if signer is not None:
+                    # c2pa_builder_sign borrows the signer's handle, so the
+                    # signer needs its own in-flight guard: the Builder's
+                    # guard holds only the Builder's handle valid, and a
+                    # signer.close() on another thread would otherwise free
+                    # this handle mid-call. Entered inside self's guard so
+                    # concurrent signs sharing objects acquire in one order.
+                    # The check above is a fast-fail; this re-check inside
+                    # the guard is the one that makes check-then-use atomic.
+                    with signer._native_call():
+                        result = _lib.c2pa_builder_sign(
+                            self._handle,
+                            format_arg,
+                            source_stream._stream,
+                            dest_stream._stream,
+                            signer._handle,
+                            ctypes.byref(manifest_bytes_ptr)
+                        )
+                else:
+                    result = _lib.c2pa_builder_sign_context(
+                        self._handle,
+                        format_arg,
+                        source_stream._stream,
+                        dest_stream._stream,
+                        ctypes.byref(manifest_bytes_ptr),
+                    )
             # Sign borrows the Builder without taking ownership.
             # Closing here ensures resources clean up,
             # and single use/single sign done by a Builder.