From e6b2dd18b55d769aa62d2bfb187db946c5d3b48c Mon Sep 17 00:00:00 2001
From: James Piechota <piechota@gmail.com>
Date: Wed, 10 Jul 2024 21:08:40 +0000
Subject: [PATCH] Ensure ar_block_cache is only update from within
 ar_node_worker

---
 apps/arweave/src/ar_block_cache.erl           | 26 +++++++++----------
 apps/arweave/src/ar_events_sup.erl            |  2 +-
 apps/arweave/src/ar_http_iface_middleware.erl | 19 ++++++++++++--
 apps/arweave/src/ar_node_worker.erl           |  4 +++
 4 files changed, 35 insertions(+), 16 deletions(-)

diff --git a/apps/arweave/src/ar_block_cache.erl b/apps/arweave/src/ar_block_cache.erl
index bd25e1d99..fab2fa658 100644
--- a/apps/arweave/src/ar_block_cache.erl
+++ b/apps/arweave/src/ar_block_cache.erl
@@ -215,6 +215,16 @@ get(Tab, H) ->
 			B
 	end.
 
+%% @doc Get the block and its status from cache.
+%% Returns not_found if the block is not in cache.
+get_block_and_status(Tab, H) ->
+	case ets:lookup(Tab, {block, H}) of
+		[] ->
+			not_found;
+		[{_, {B, Status, Timestamp, _Children}}] ->
+			{B, {Status, Timestamp}}
+	end.
+
 %% @doc Get a {block, previous blocks, status} tuple for the earliest block from
 %% the longest chain, which has not been validated yet. The previous blocks are
 %% sorted from newest to oldest. The last one is a block from the current fork.
@@ -297,16 +307,6 @@ tx_id(#tx{ id = ID }) ->
 tx_id(TXID) ->
 	TXID.
 
-%% @doc Get the block and its status from cache.
-%% Returns not_found if the block is not in cache.
-get_block_and_status(Tab, H) ->
-	case ets:lookup(Tab, {block, H}) of
-		[] ->
-			not_found;
-		[{_, {B, Status, Timestamp, _Children}}] ->
-			{B, {Status, Timestamp}}
-	end.
-
 %% @doc Mark the given block as the tip block. Mark the previous blocks as on-chain.
 %% Mark the on-chain blocks from other forks as validated. Raises invalid_tip if
 %% one of the preceeding blocks is not validated. Raises not_found if the block
@@ -458,6 +458,8 @@ get_siblings(Tab, B) ->
 
 update_timestamp(Tab, H, ReceiveTimestamp) ->
 	case ets:lookup(Tab, {block, H}) of
+		[] ->
+			not_found;
 		[{_, {B, Status, Timestamp, Children}}] ->
 			case B#block.receive_timestamp of
 				undefined ->
@@ -472,9 +474,7 @@ update_timestamp(Tab, H, ReceiveTimestamp) ->
 						}, false);
 				_ ->
 					ok
-			end;
-		[] ->
-			not_found
+			end
 	end.
 
 %%%===================================================================
diff --git a/apps/arweave/src/ar_events_sup.erl b/apps/arweave/src/ar_events_sup.erl
index 39fc48ff0..bc734ff8f 100644
--- a/apps/arweave/src/ar_events_sup.erl
+++ b/apps/arweave/src/ar_events_sup.erl
@@ -37,7 +37,7 @@ init([]) ->
 		%% Events: new, ready_for_mining, orphaned, emitting_scheduled,
 		%% preparing_unblacklisting, ready_for_unblacklisting, registered_offset.
 		?CHILD(ar_events, tx, worker),
-		%% Events: discovered, rejected, new, double_signing.
+		%% Events: discovered, rejected, new, double_signing, mined_block_received.
 		?CHILD(ar_events, block, worker),
 		%% Events: unpacked, packed.
 		?CHILD(ar_events, chunk, worker),
diff --git a/apps/arweave/src/ar_http_iface_middleware.erl b/apps/arweave/src/ar_http_iface_middleware.erl
index 1c830e1ab..85d510000 100644
--- a/apps/arweave/src/ar_http_iface_middleware.erl
+++ b/apps/arweave/src/ar_http_iface_middleware.erl
@@ -2321,7 +2321,7 @@ handle_block_announcement(#block_announcement{ indep_hash = H, previous_block =
 		solution_hash = SolutionH }, Req) ->
 	case ar_ignore_registry:member(H) of
 		true ->
-			ar_block_cache:update_timestamp(block_cache, H, erlang:timestamp()),
+			check_block_receive_timestamp(H),
 			{208, #{}, <<>>, Req};
 		false ->
 			case ar_node:get_block_shadow_from_cache(PrevH) of
@@ -2438,7 +2438,7 @@ post_block(check_block_hash_header, Peer, {Req, Pid, Encoding}, ReceiveTimestamp
 				{ok, BH} when byte_size(BH) =< 48 ->
 					case ar_ignore_registry:member(BH) of
 						true ->
-							ar_block_cache:update_timestamp(block_cache, BH, ReceiveTimestamp),
+							check_block_receive_timestamp(BH),
 							{208, #{}, <<"Block already processed.">>, Req};
 						false ->
 							post_block(read_body, Peer, {Req, Pid, Encoding},
@@ -2508,6 +2508,21 @@ post_block(enqueue_block, {B, Peer}, Req, ReceiveTimestamp) ->
 	end,
 	{200, #{}, <<"OK">>, Req}.
 
+check_block_receive_timestamp(H) ->
+	case ar_block_cache:get(block_cache, H) of
+		not_found ->
+			not_found;
+		B ->
+			case B#block.receive_timestamp of
+				undefined ->
+					%% This node mined block H and this is the first time it's been
+					%% gossipped back to it. Update the node's receive_timestamp.
+					ar_events:send(block, {mined_block_received, H, erlang:timestamp()});
+				_ ->
+					ok
+			end
+	end.
+
 handle_post_partial_solution(Req, Pid) ->
 	{ok, Config} = application:get_env(arweave, config),
 	CMExitNode = ar_coordination:is_exit_peer() andalso ar_pool:is_client(),
diff --git a/apps/arweave/src/ar_node_worker.erl b/apps/arweave/src/ar_node_worker.erl
index 17eaf2a8b..9ba604a18 100644
--- a/apps/arweave/src/ar_node_worker.erl
+++ b/apps/arweave/src/ar_node_worker.erl
@@ -534,6 +534,10 @@ handle_info({event, block, {new, B, _Source}}, State) ->
 			{noreply, State}
 	end;
 
+handle_info({event, block, {mined_block_received, H, ReceiveTimestamp}}, State) ->
+	ar_block_cache:update_timestamp(block_cache, H, ReceiveTimestamp),
+	{noreply, State};
+
 handle_info({event, block, _}, State) ->
 	{noreply, State};