fix: Avoid ub when casting function pointer to regular pointer

Tradias · Oct 6, 2023 · cd384a4 · cd384a4
1 parent 05ce0ee
commit cd384a4
Show file tree

Hide file tree

Showing 5 changed files with 28 additions and 17 deletions.
diff --git a/src/agrpc/detail/basic_sender.hpp b/src/agrpc/detail/basic_sender.hpp
@@ -159,7 +159,7 @@ class BasicSenderRunningOperation : public detail::BaseForSenderImplementationTy
     template <AllocationType, int Id>
     void set_on_complete() noexcept
     {
-        detail::OperationBaseAccess::get_on_complete(*this) = &do_complete<Id>;
+        detail::OperationBaseAccess::set_on_complete(*this, &do_complete<Id>);
     }
 
     template <AllocationType, class... Args>
@@ -169,17 +169,14 @@ class BasicSenderRunningOperation : public detail::BaseForSenderImplementationTy
         detail::satisfy_receiver(static_cast<Receiver&&>(receiver()), static_cast<Args&&>(args)...);
     }
 
-    void put_into_scratch_space(void* ptr) noexcept
-    {
-        detail::OperationBaseAccess::get_on_complete(*this) = reinterpret_cast<detail::OperationOnComplete>(ptr);
-    }
+    void put_into_scratch_space(void* ptr) noexcept { detail::OperationBaseAccess::set_scratch_space(*this, ptr); }
 
     [[nodiscard]] void* get_scratch_space() const noexcept
     {
-        return reinterpret_cast<void*>(detail::OperationBaseAccess::get_on_complete(*this));
+        return detail::OperationBaseAccess::get_scratch_space(*this);
     }
 
-    void restore_scratch_space() noexcept { detail::OperationBaseAccess::get_on_complete(*this) = &do_complete; }
+    void restore_scratch_space() noexcept { detail::OperationBaseAccess::set_on_complete(*this, &do_complete); }
 
   private:
     detail::CompressedPair<detail::CompressedPair<Receiver, StopCallback>, Implementation> impl_;

diff --git a/src/agrpc/detail/operation_base.hpp b/src/agrpc/detail/operation_base.hpp
@@ -53,7 +53,11 @@ class OperationBase
   private:
     friend detail::OperationBaseAccess;
 
-    OperationOnComplete on_complete_;
+    union
+    {
+        OperationOnComplete on_complete_;
+        void* scratch_space_;
+    };
 };
 
 class QueueableOperationBase : public detail::OperationBase
@@ -71,9 +75,19 @@ class QueueableOperationBase : public detail::OperationBase
 
 struct OperationBaseAccess
 {
-    static auto& get_on_complete(detail::OperationBase& operation) noexcept { return operation.on_complete_; }
+    static void set_on_complete(detail::OperationBase& operation, OperationOnComplete on_complete) noexcept
+    {
+        operation.on_complete_ = on_complete;
+    }
 
     static auto get_on_complete(const detail::OperationBase& operation) noexcept { return operation.on_complete_; }
+
+    static void set_scratch_space(detail::OperationBase& operation, void* ptr) noexcept
+    {
+        operation.scratch_space_ = ptr;
+    }
+
+    static void* get_scratch_space(const detail::OperationBase& operation) noexcept { return operation.scratch_space_; }
 };
 
 [[nodiscard]] constexpr bool is_ok(OperationResult result) noexcept { return result == OperationResult::OK; }

diff --git a/src/agrpc/detail/repeatedly_request_context.hpp b/src/agrpc/detail/repeatedly_request_context.hpp
@@ -107,8 +107,8 @@ class RepeatedlyRequestOperation : public detail::QueueableOperationBase,
 
     void add_completing_operation(agrpc::GrpcContext& grpc_context)
     {
-        detail::OperationBaseAccess::get_on_complete(*this) =
-            detail::DO_COMPLETE_NO_ARG_HANDLER<RepeatedlyRequestOperation>;
+        detail::OperationBaseAccess::set_on_complete(*this,
+                                                     detail::DO_COMPLETE_NO_ARG_HANDLER<RepeatedlyRequestOperation>);
         detail::GrpcContextImplementation::add_local_operation(grpc_context, this);
     }
 

diff --git a/src/agrpc/detail/sender_implementation_operation.hpp b/src/agrpc/detail/sender_implementation_operation.hpp
@@ -105,7 +105,7 @@ struct SenderImplementationOperation : public detail::BaseForSenderImplementatio
     template <AllocationType AllocType, int Id>
     void set_on_complete() noexcept
     {
-        detail::OperationBaseAccess::get_on_complete(*this) = &do_complete<AllocType, Id>;
+        detail::OperationBaseAccess::set_on_complete(*this, &do_complete<AllocType, Id>);
     }
 
     template <AllocationType AllocType, class... Args>

diff --git a/src/agrpc/detail/server_write_reactor.hpp b/src/agrpc/detail/server_write_reactor.hpp
@@ -63,15 +63,15 @@ class ServerWriteReactor : public detail::ServerWriteReactorStepBase, public det
 
     void write(const Response& response)
     {
-        get_on_complete() = &ServerWriteReactor::do_write_done;
+        set_on_complete(&ServerWriteReactor::do_write_done);
         writer_.Write(response, static_cast<StepBase*>(this));
     }
 
     [[nodiscard]] bool is_finished() const noexcept { return &ServerWriteReactor::do_finish_done == get_on_complete(); }
 
     void finish(const grpc::Status& status)
     {
-        get_on_complete() = &ServerWriteReactor::do_finish_done;
+        set_on_complete(&ServerWriteReactor::do_finish_done);
         writer_.Finish(status, static_cast<StepBase*>(this));
     }
 
@@ -80,17 +80,17 @@ class ServerWriteReactor : public detail::ServerWriteReactorStepBase, public det
   private:
     auto get_allocator() const noexcept { return grpc_context_->get_allocator(); }
 
-    auto& get_on_complete() noexcept
+    void set_on_complete(detail::OperationOnComplete on_complete) noexcept
     {
-        return detail::OperationBaseAccess::get_on_complete(*static_cast<StepBase*>(this));
+        detail::OperationBaseAccess::set_on_complete(*static_cast<StepBase*>(this), on_complete);
     }
 
     auto get_on_complete() const noexcept
     {
         return detail::OperationBaseAccess::get_on_complete(*static_cast<const StepBase*>(this));
     }
 
-    void set_step_done() noexcept { get_on_complete() = nullptr; }
+    void set_step_done() noexcept { set_on_complete(nullptr); }
 
     bool is_finishing_or_writing() const noexcept { return nullptr != get_on_complete(); }