HaloKeymind/src/Mesh.h

#pragma once

#include <Dispatcher.h>

// OTA-over-LoRa transport is understood by every Mesh role even when the OTA manager/installer is not
// compiled in. Repeaters can therefore relay PAYLOAD_TYPE_OTA opaquely while they are on TempRadio.
// OTA traffic always uses the lowest TX priority (selected only after all real traffic).
#ifndef OTA_TX_PRIORITY
#define OTA_TX_PRIORITY 250
#endif
// Keep this many pool slots free while relaying OTA, so best-effort firmware traffic cannot monopolise the
// shared packet pool and starve real traffic.
#ifndef OTA_FWD_MIN_FREE
#define OTA_FWD_MIN_FREE 4
#endif

namespace mesh {

#ifndef MAX_DIRECT_RETRY_SLOTS
  #define MAX_DIRECT_RETRY_SLOTS  6
#endif

#ifndef MAX_FLOOD_RETRY_SLOTS
  #define MAX_FLOOD_RETRY_SLOTS   6
#endif

#ifndef MAX_RECENT_ADVERT_ECHOS
  #define MAX_RECENT_ADVERT_ECHOS  8
#endif
#if MAX_RECENT_ADVERT_ECHOS < 1
  #error "MAX_RECENT_ADVERT_ECHOS must be at least 1"
#endif

#ifndef FLOOD_RETRY_PATH_GATE_DISABLED
  #define FLOOD_RETRY_PATH_GATE_DISABLED  0xFF
#endif

class GroupChannel {
public:
  uint8_t hash[PATH_HASH_SIZE];
  uint8_t secret[PUB_KEY_SIZE];
};

/**
 * An abstraction of the data tables needed to be maintained
*/
class MeshTables {
public:
  virtual bool wasSeen(const Packet* packet) = 0;
  virtual void markSeen(const Packet* packet) = 0;
  virtual void markSent(const Packet* packet) = 0;
  virtual void clear(const Packet* packet) = 0;    // remove this packet hash from table
};

/**
 * \brief  The next layer in the basic Dispatcher task, Mesh recognises the particular Payload TYPES,
 *     and provides virtual methods for sub-classes on handling incoming, and also preparing outbound Packets.
*/
class Mesh : public Dispatcher {
  struct DirectRetryEntry {
    Packet* packet;
    Packet* trigger_packet;
    unsigned long retry_started_at;
    unsigned long echo_wait_started_at;
    unsigned long retry_at;
    uint32_t retry_delay;
    uint32_t message_timestamp;
    uint8_t retry_attempts_sent;
    uint8_t retry_key[MAX_HASH_SIZE];
    uint8_t trace_replacement_key[MAX_HASH_SIZE];
    uint8_t message_replacement_key[MAX_HASH_SIZE];
    uint8_t next_hop_hash[MAX_HASH_SIZE];
    uint8_t next_hop_hash_len;
    uint8_t payload_type;
    uint8_t priority;
    uint8_t progress_marker;
    bool expect_path_growth;
    bool final_hop_retry;
    bool waiting_final_echo;
    bool queued;
    bool has_message_replacement_key;
    bool active;
  };

  struct FloodRetryEntry {
    Packet* packet;
    Packet* trigger_packet;
    unsigned long retry_started_at;
    unsigned long retry_at;
    uint32_t retry_delay;
    uint32_t message_timestamp;
    uint8_t retry_attempts_sent;
    uint8_t retry_key[MAX_HASH_SIZE];
    uint8_t message_replacement_key[MAX_HASH_SIZE];
    uint8_t priority;
    uint8_t progress_marker;
    bool self_advert;
    bool waiting_final_echo;
    bool queued;
    bool has_message_replacement_key;
    bool active;
  };

  struct RecentAdvertEchoEntry {
    uint8_t packet_hash[MAX_HASH_SIZE];
    uint32_t advert_timestamp;
    uint32_t watch_started_at;
    uint8_t progress_marker;
    bool confirmed;
    bool valid;
  };

  RTCClock* _rtc;
  RNG* _rng;
  MeshTables* _tables;
  DirectRetryEntry _direct_retries[MAX_DIRECT_RETRY_SLOTS];
  FloodRetryEntry _flood_retries[MAX_FLOOD_RETRY_SLOTS];
  RecentAdvertEchoEntry _recent_advert_echoes[MAX_RECENT_ADVERT_ECHOS];
  uint8_t _active_direct_retry_count;
  uint8_t _active_flood_retry_count;
  uint8_t _waiting_direct_retry_count;
  uint8_t _waiting_flood_retry_count;
  uint8_t _next_recent_advert_echo;
  unsigned long _next_direct_retry_timeout;
  unsigned long _next_flood_retry_timeout;

  void removePathPrefix(Packet* packet, uint8_t prefix_count);
  void routeDirectRecvAcks(Packet* packet, uint32_t delay_millis);
  void rebuildNextDirectRetryTimeout();
  void rebuildNextFloodRetryTimeout();
  void clearDirectRetrySlot(int idx);
  void retireDirectRetrySlot(int idx);
  void calculateDirectRetryKey(const Packet* packet, uint8_t* dest_key) const;
  bool calculateTraceReplacementKey(const Packet* packet, uint8_t* dest_key) const;
  void replaceQueuedTraceRetries(const Packet* packet);
  bool cancelDirectRetryOnEcho(const Packet* packet);
  void armDirectRetryOnSendComplete(const Packet* packet);
  void clearPendingDirectRetryOnSendFail(const Packet* packet);
  bool getDirectRetryTarget(const Packet* packet, const uint8_t*& next_hop_hash, uint8_t& next_hop_hash_len,
                            uint8_t& progress_marker, bool& expect_path_growth) const;
  bool canDecodeDirectPayloadForSelf(const Packet* packet);
  void maybeScheduleDirectRetry(const Packet* packet, uint8_t priority, bool final_hop_retry = false);
  void clearFloodRetrySlot(int idx);
  void retireFloodRetrySlot(int idx);
  void replaceQueuedSelfAdvertRetries(const Packet* packet);
  bool cancelFloodRetryOnEcho(const Packet* packet);
  bool getRecentAdvertTimestamp(const Packet* packet, uint32_t& timestamp) const;
  bool isRecentAdvertTimestamp(uint32_t timestamp) const;
  void watchForwardedAdvertEcho(const Packet* packet);
  void observeForwardedAdvertEcho(const Packet* packet);
  bool shouldSuppressEchoedAdvertForward(const Packet* packet) const;
  uint8_t getEffectiveFloodRetryMaxAttempts(const Packet* packet) const;
  uint8_t getEligibleFloodRetryMaxAttempts(const Packet* packet) const;
  void armFloodRetryOnSendComplete(const Packet* packet);
  void clearPendingFloodRetryOnSendFail(const Packet* packet);
  void maybeScheduleFloodRetry(const Packet* packet, uint8_t priority);
  void serviceLoopMaintenance();
  //void routeRecvAcks(Packet* packet, uint32_t delay_millis);
  DispatcherAction forwardMultipartDirect(Packet* pkt);

protected:
  DispatcherAction onRecvPacket(Packet* pkt) override;
  void onTracePacketQueuedForSend(Packet* packet) override;
  void onSendComplete(Packet* packet) override;
  void onSendFail(Packet* packet) override;
  bool allowPacketTransmit(const Packet* packet) const override;
  bool usePassiveChannelCheck(const Packet* packet) const override;
  bool getNextRetryWakeDelay(uint32_t& delay_millis) const;
  bool hasRetryWorkDue() const {
    uint32_t delay_millis;
    return getNextRetryWakeDelay(delay_millis) && delay_millis == 0;
  }

  virtual uint32_t getCADFailRetryDelay() const override;

  /**
   * \brief  Decide what to do with received packet, ie. discard, forward, or hold
   */
  DispatcherAction routeRecvPacket(Packet* packet);

  /**
   * \brief    Called _before_ the packet is dispatched to the on..Recv() methods.
   * \returns  true, if given packet should NOT be processed.
   */
  virtual bool filterRecvFloodPacket(Packet* packet) { return false; }

  /**
   * \brief  Check whether this packet should be forwarded (re-transmitted) or not.
   *     Is sub-classes responsibility to make sure given packet is only transmitted ONCE (by this node)
   */
  virtual bool allowPacketForward(const Packet* packet);

  /**
   * \returns  number of milliseconds delay to apply to retransmitting the given packet.
   */
  virtual uint32_t getRetransmitDelay(const Packet* packet);

  /**
   * \returns  number of milliseconds delay to apply to retransmitting the given packet, for DIRECT mode.
   */
  virtual uint32_t getDirectRetransmitDelay(const Packet* packet);

  /**
   * \brief  Decide whether a DIRECT packet should retry if the next hop echo is not overheard.
   *         Sub-classes can use recent repeater or other link-quality data to opt in selectively.
   */
  virtual bool allowDirectRetry(const Packet* packet, const uint8_t* next_hop_hash, uint8_t next_hop_hash_len) const;

  /**
   * \brief  Allow subclasses to rewrite a non-TRACE DIRECT packet path when this node can safely skip ahead.
   */
  virtual bool maybeShortCircuitDirect(Packet* packet) { return false; }

  /**
   * \returns  milliseconds to wait for the next-hop echo before queueing a retry of the DIRECT packet.
   */
  virtual uint32_t getDirectRetryEchoDelay(const Packet* packet) const;

  /**
   * \returns  packet-airtime multiplier used by the direct-retry echo window.
   */
  uint8_t getDirectRetryPacketAirtimeFactor(const Packet* packet) const;

  /**
   * \returns  packet-airtime add-on used by the direct-retry echo window.
   */
  uint32_t getDirectRetryPacketAirtimeDelay(const Packet* packet) const;

  /**
   * \returns  maximum number of retry transmissions after the initial direct TX.
   */
  virtual uint8_t getDirectRetryMaxAttempts(const Packet* packet) const;

  /**
   * \returns  delay before a specific retry attempt, where attempt_idx=0 is the first retry.
   */
  virtual uint32_t getDirectRetryAttemptDelay(const Packet* packet, uint8_t attempt_idx);

  /**
   * \brief  Cancel queued or future direct/flood retries for a packet payload hash.
   *         The original transmission and any retry already in progress are allowed to finish.
   * \returns true if at least one active retry sequence was cancelled.
   */
  bool cancelActiveRetries(const uint8_t retry_key[MAX_HASH_SIZE]);

  /**
   * \brief  After a replacement packet is queued, retire the prior retry
   *         sequence and ensure the queued packet owns the retry slot.
   */
  void replaceActiveRetries(const Packet* replacement_packet,
                            const uint8_t retry_key[MAX_HASH_SIZE]);

  /**
   * \brief  Mark a successfully queued text message as the current retry owner.
   *         Older retries with the same caller-defined message key and a
   *         different timestamp are retired across both direct and flood routes.
   */
  void replaceActiveMessageRetries(
      const Packet* replacement_packet,
      const uint8_t message_key[MAX_HASH_SIZE], uint32_t message_timestamp);

  // Cancel every queued/future retry of one route type. A packet already on
  // air is allowed to finish, but no later retry is armed from it.
  void cancelAllDirectRetries();
  void cancelAllFloodRetries();

  /**
   * \returns true while a direct or flood retry sequence owns this packet payload hash.
   */
  bool hasActiveRetries(const uint8_t retry_key[MAX_HASH_SIZE]) const;

  /**
   * \brief  Decide whether a FLOOD packet should retry when no downstream echo is overheard.
   */
  virtual bool allowFloodRetry(const Packet* packet) const;

  /**
   * \brief  Reserve any role-specific state after a free flood retry slot is found.
   * \returns false if role-specific state could not be reserved.
   */
  virtual bool prepareFloodRetry(const Packet* packet) const {
    (void)packet;
    return true;
  }

  /**
   * \returns true only for a zero-hop advert carrying this node's public key.
   */
  bool isSelfOriginAdvert(const Packet* packet) const;

  /**
   * \brief  Return true when this FLOOD packet already carries an application-defined target prefix.
   */
  virtual bool hasFloodRetryTargetPrefix(const Packet* packet) const;

  /**
   * \returns  maximum flood path hash count eligible for retry, or FLOOD_RETRY_PATH_GATE_DISABLED.
   */
  virtual uint8_t getFloodRetryMaxPathLength(const Packet* packet) const;

  /**
   * \returns  the stricter of the general flood retry gate and the group-data-specific gate.
   */
  static uint8_t applyGroupDataFloodRetryPathGate(const Packet* packet,
                                                  uint8_t general_gate,
                                                  uint8_t group_data_gate);

  /**
   * \returns  the shared payload/path cap applied after a role chooses its flood retry count.
   */
  static uint8_t applyFloodRetryAttemptPolicy(const Packet* packet,
                                               uint8_t role_max_attempts);

  /**
   * \returns  maximum number of FLOOD retry transmissions after the initial TX.
   */
  virtual uint8_t getFloodRetryMaxAttempts(const Packet* packet) const;

  /**
   * \brief  Return true when a received FLOOD echo is enough to cancel a pending retry.
   */
  virtual bool isFloodRetryEchoTarget(const Packet* packet, uint8_t progress_marker) const;

  /**
   * \returns  delay before a specific flood retry attempt, where attempt_idx=0 is the first retry.
   */
  virtual uint32_t getFloodRetryAttemptDelay(const Packet* packet, uint8_t attempt_idx);

  /**
   * \brief  Optional hook for logging flood-retry lifecycle events.
   *         packet is null when the retained packet has already been released.
   */
  virtual void onFloodRetryEvent(const char* event, const Packet* packet, uint32_t delay_millis, uint8_t retry_attempt) { }

  /**
   * \brief  Called exactly once whenever an active flood-retry slot is released.
   */
  virtual void onFloodRetrySlotReleased(const uint8_t* retry_key) { }

  /**
   * \returns  number of extra (Direct) ACK transmissions wanted.
   */
  virtual uint8_t getExtraAckTransmitCount() const;

  /**
   * \brief  Optional hook for logging direct-retry lifecycle events.
   */
  virtual void onDirectRetryEvent(const char* event, const Packet* packet, uint32_t delay_millis, uint8_t retry_attempt,
                                  const uint8_t* target_hash = NULL, uint8_t target_hash_len = 0,
                                  int16_t payload_type = -1) { }

  /**
   * \brief  Optional hook for link-quality feedback when all direct-retry attempts fail.
   */
  virtual void onDirectRetryFailed(const uint8_t* next_hop_hash, uint8_t next_hop_hash_len) { }

  /**
   * \brief  Optional hook for link-quality feedback when a direct-retry echo is heard.
   */
  virtual void onDirectRetrySucceeded(const uint8_t* next_hop_hash, uint8_t next_hop_hash_len, int8_t snr_x4) { }

  /**
   * \returns  Coding rate to use for a retry attempt, starting from the current/adaptive CR.
   */
  static uint8_t getDirectRetryCodingRateForAttempt(uint8_t start_cr, uint8_t retry_attempt);

  /**
   * \brief  Optional hook to set local-only transmit options on a retry packet before it is queued.
   */
  virtual void configureDirectRetryPacket(Packet* retry, const Packet* original, uint8_t retry_attempt);

  /**
   * \brief  Perform search of local DB of peers/contacts.
   * \returns  Number of peers with matching hash
   */
  virtual int searchPeersByHash(const uint8_t* hash);

  /**
   * \brief  lookup the ECDH shared-secret between this node and peer by idx (calculate if necessary)
   * \param  dest_secret  destination array to copy the secret (must be PUB_KEY_SIZE bytes)
   * \param  peer_idx  index of peer, [0..n) where n is what searchPeersByHash() returned
   */
  virtual void getPeerSharedSecret(uint8_t* dest_secret, int peer_idx) { }

  /**
   * \brief  A (now decrypted) data packet has been received (by a known peer).
   *         NOTE: these can be received multiple times (per sender/msg-id), via different routes
   * \param  type  one of: PAYLOAD_TYPE_TXT_MSG, PAYLOAD_TYPE_REQ, PAYLOAD_TYPE_RESPONSE
   * \param  sender_idx  index of peer, [0..n) where n is what searchPeersByHash() returned
   * \param  secret   the pre-calculated shared-secret (handy for sending response packet)
   * \param  data   decrypted data from payload
  */
  virtual void onPeerDataRecv(Packet* packet, uint8_t type, int sender_idx, const uint8_t* secret, uint8_t* data, size_t len) { }

  /**
   * \brief  A TRACE packet has been received. (and has reached the end of its given path)
   *         NOTE: this may have been initiated by another node.
   * \param  tag         a random (unique-ish) tag set by initiator
   * \param  auth_code   a code to authenticate the packet
   * \param  flags       zero for now
   * \param  path_snrs   single byte SNR*4 for each hop in the path
   * \param  path_hashes hashes of each repeater in the path
   * \param  path_len    length of the path_snrs[] and path_hashes[] arrays
  */
  virtual void onTraceRecv(Packet* packet, uint32_t tag, uint32_t auth_code, uint8_t flags, const uint8_t* path_snrs, const uint8_t* path_hashes, uint8_t path_len) { }

  /**
   * \brief  A path TO peer (sender_idx) has been received. (also with optional 'extra' data encoded)
   *         NOTE: these can be received multiple times (per sender), via different routes
   * \param  sender_idx  index of peer, [0..n) where n is what searchPeersByHash() returned
   * \param  secret   the pre-calculated shared-secret (handy for sending response packet)
   * \returns   true, if path was accepted and that reciprocal path should be sent
  */
  virtual bool onPeerPathRecv(Packet* packet, int sender_idx, const uint8_t* secret, uint8_t* path, uint8_t path_len, uint8_t extra_type, uint8_t* extra, uint8_t extra_len) { return false; }

  /**
   * \brief  A new incoming Advertisement has been received.
   *         NOTE: these can be received multiple times (per id/timestamp), via different routes
  */
  virtual void onAdvertRecv(Packet* packet, const Identity& id, uint32_t timestamp, const uint8_t* app_data, size_t app_data_len) { }

  /**
   * \brief  A (now decrypted) data packet has been received.
   *         NOTE: these can be received multiple times (per sender/contents), via different routes
   * \param  secret  ECDH shared secret
   * \param  sender  public key provided by sender
  */
  virtual void onAnonDataRecv(Packet* packet, const uint8_t* secret, const Identity& sender, uint8_t* data, size_t len) { }

  /**
   * \brief  A path TO 'sender' has been received. (also with optional 'extra' data encoded)
   *         NOTE: these can be received multiple times (per sender), via different routes
  */
  virtual void onPathRecv(Packet* packet, Identity& sender, uint8_t* path, uint8_t path_len, uint8_t extra_type, uint8_t* extra, uint8_t extra_len) { }

  /**
   * \brief  A control packet has been received.
  */
  virtual void onControlDataRecv(Packet* packet) { }

  /**
   * \brief  A packet with PAYLOAD_TYPE_RAW_CUSTOM has been received.
  */
  virtual void onRawDataRecv(Packet* packet) { }

  /** True only while the role has actually switched into a temporary-radio window. */
  virtual bool isTempRadioActive() const { return false; }

#if defined(ENABLE_OTA)
  /**
   * \brief  An OTA-over-LoRa packet (PAYLOAD_TYPE_OTA) has been received. Subclasses forward the
   *         payload bytes to their OtaManager. See docs/ota_protocol.md.
   */
  virtual void onOtaRecv(Packet* packet) { }

  /** \returns  max OTA flood reach in hops - accept up to N, relay while < N (0=direct). `ota config hops`. */
  virtual uint8_t getOtaHopLimit() const;

  // OTA mesh-integration is centralized in Mesh::begin()/loop()/dispatch, so every role (repeater,
  // companion, room, sensor, ...) gets fetch/serve/apply without per-example wiring.
  static void otaSendAdapter(void* ctx, const uint8_t* msg, uint16_t len, bool flood);
  unsigned long _next_ota_tick = 0;
  unsigned long _next_ota_announce = 0;   // advertisements are scheduled only while temp radio is active
  uint8_t       _ota_announce_count = 0;  // adverts sent so far (boot burst before settling to daily)
  bool          _ota_resumed = false;     // one-shot: resumed an interrupted fetch staged in flash on boot
  bool          _ota_autoinstall_tried = false;  // attempted auto-install for the current COMPLETE fetch
  bool          _ota_temp_was_active = false;    // detects entry into a temporary-radio window
#endif

  /**
   * \brief  Perform search of local DB of matching GroupChannels.
   * \param  channels  OUT - store matching channels in this array, up to max_matches
   * \returns  Number of channels with matching hash
   */
  virtual int searchChannelsByHash(const uint8_t* hash, GroupChannel channels[], int max_matches);

  /**
   * \brief  A structurally valid encrypted group packet has been observed.
   *         Called once for each unseen packet before local channel matching and
   *         independently of whether the packet will be forwarded.
   */
  virtual void onGroupPacketRecv(Packet* packet) { }

  /**
   * \brief  An encrypted group data packet has been received.
   *         NOTE: the same payload can be received multiple times, via different routes
   * \param  type  one of: PAYLOAD_TYPE_GRP_TXT, PAYLOAD_TYPE_GRP_DATA
   * \param  channel  the matching GroupChannel
  */
  virtual void onGroupDataRecv(Packet* packet, uint8_t type, const GroupChannel& channel, uint8_t* data, size_t len) { }

  /**
   * \brief  A simple ACK packet has been received.
   *         NOTE: same ACK can be received multiple times, via different routes
  */
  virtual void onAckRecv(Packet* packet, uint32_t ack_crc) { }

  Mesh(Radio& radio, MillisecondClock& ms, RNG& rng, RTCClock& rtc, PacketManager& mgr, MeshTables& tables)
    : Dispatcher(radio, ms, mgr), _rng(&rng), _rtc(&rtc), _tables(&tables)
  {
  }

  MeshTables* getTables() const { return _tables; }

public:
  void begin();
  void loop();

  LocalIdentity self_id;

  RNG* getRNG() const { return _rng; }
  RTCClock* getRTCClock() const { return _rtc; }

  Packet* createAdvert(const LocalIdentity& id, const uint8_t* app_data=NULL, size_t app_data_len=0);
  Packet* createDatagram(uint8_t type, const Identity& dest, const uint8_t* secret, const uint8_t* data, size_t len);
  Packet* createAnonDatagram(uint8_t type, const LocalIdentity& sender, const Identity& dest, const uint8_t* secret, const uint8_t* data, size_t data_len);
  Packet* createGroupDatagram(uint8_t type, const GroupChannel& channel, const uint8_t* data, size_t data_len);
  Packet* createAck(const uint8_t* ack_hash, uint8_t ack_len);
  Packet* createAck(uint32_t ack_crc);
  Packet* createMultiAck(const uint8_t* ack_hash, uint8_t ack_len, uint8_t remaining);
  Packet* createMultiAck(uint32_t ack_crc, uint8_t remaining);
  Packet* createPathReturn(const uint8_t* dest_hash, const uint8_t* secret, const uint8_t* path, uint8_t path_len, uint8_t extra_type, const uint8_t*extra, size_t extra_len);
  Packet* createPathReturn(const Identity& dest, const uint8_t* secret, const uint8_t* path, uint8_t path_len, uint8_t extra_type, const uint8_t*extra, size_t extra_len);
  Packet* createRawData(const uint8_t* data, size_t len);

#if defined(ENABLE_OTA)
  // Build a PAYLOAD_TYPE_OTA packet from raw OTA message bytes (route set by sendOtaFlood).
  Packet* createOtaPacket(const uint8_t* data, size_t len);
  // Flood-send at the lowest priority (so OTA never competes with mesh traffic).
  void sendOtaFlood(Packet* packet, uint32_t delay_millis = 0);
#endif
  Packet* createTrace(uint32_t tag, uint32_t auth_code, uint8_t flags = 0);
  Packet* createControlData(const uint8_t* data, size_t len);

  /**
   * \brief  send a locally-generated Packet with flood routing
  */
  bool sendFlood(Packet* packet, uint32_t delay_millis=0, uint8_t path_hash_size=1);

  /**
   * \brief  send a locally-generated Packet with flood routing
   * \param transport_codes   array of 2 codes to attach to packet
  */
  bool sendFlood(Packet* packet, uint16_t* transport_codes, uint32_t delay_millis=0, uint8_t path_hash_size=1);

  /**
   * \brief  send a locally-generated Packet with Direct routing
  */
  bool sendDirect(Packet* packet, const uint8_t* path, uint8_t path_len, uint32_t delay_millis=0);

  /**
   * \brief  send a locally-generated Packet to just neighbor nodes (zero hops)
  */
  void sendZeroHop(Packet* packet, uint32_t delay_millis=0);

  /**
   * \brief  send a locally-generated Packet to just neighbor nodes (zero hops), with specific transport codes
   * \param transport_codes   array of 2 codes to attach to packet
  */
  void sendZeroHop(Packet* packet, uint16_t* transport_codes, uint32_t delay_millis=0);

};

}