// mqtt_concatenator.cpp (vulnerable)
// Concatenates one payload from EACH of 3 MQTT topics, then publishes the combined buffer.

#include <mosquitto.h>
#include <iostream>
#include <cstring>
#include <array>

const char* BROKER = "192.168.1.89";
const char* TOPICS[3] = {"lidar/scan", "lidar/car2/scan", "lidar/car3/scan"};

// Big buffer for demo
//static uint8_t buffer[262144];  // 256 KB
static uint8_t buffer[600];  // 256 KB

static size_t  offset = 0;
static bool    overflow_imminent = false;
static std::array<bool,3> got = {false, false, false};

static inline int topic_index(const char* topic) {
  if (std::strcmp(topic, "lidar/scan") == 0)       return 0;
  if (std::strcmp(topic, "lidar/car2/scan") == 0)  return 1;
  if (std::strcmp(topic, "lidar/car3/scan") == 0)  return 2;
  return -1; // unexpected topic
}

static inline bool all_got() {
  return got[0] && got[1] && got[2];
}

static void reset_cycle() {
  offset = 0;
  overflow_imminent = false;
  got = {false, false, false};
}

static void attempt_publish(struct mosquitto* mosq) {
  if (!all_got()) return;

  if (!overflow_imminent)
    std::cout << "Safe: combined " << offset << " bytes\n";
  else
    std::cout << "Overflow: tried " << offset << " > buffer\n";

  mosquitto_publish(mosq, /*mid*/NULL, "lidar/cpp",
                    static_cast<int>(offset), buffer, /*qos*/0, /*retain*/false);
  std::cout << "Published " << offset << " bytes\n";

  reset_cycle();
}

static void on_connect(struct mosquitto* mosq, void*, int rc) {
  if (rc == 0) {
    std::cout << "Connected to broker\n";
    for (int i = 0; i < 3; ++i) {
      mosquitto_subscribe(mosq, NULL, TOPICS[i], 0);
      std::cout << "Subscribed to: " << TOPICS[i] << "\n";
    }
  } else {
    std::cerr << "Connect failed, rc=" << rc << "\n";
  }
}

static void on_message(struct mosquitto* mosq, void*, const struct mosquitto_message* msg) {
  // Always log every message in the requested style
  std::cout << "CPP VNF Received payload from " << msg->topic
            << " (" << msg->payloadlen << " bytes)\n";

  int idx = topic_index(msg->topic);
  if (idx < 0) return; // ignore unexpected topics

  // Only the first message per topic in a cycle contributes to the concat
  if (!got[idx]) {
    if (offset + static_cast<size_t>(msg->payloadlen) > sizeof(buffer))
      overflow_imminent = true;

    std::memcpy(buffer + offset, msg->payload, static_cast<size_t>(msg->payloadlen));
    offset += static_cast<size_t>(msg->payloadlen);
    got[idx] = true;

    if (all_got())
      attempt_publish(mosq);
  }
}

int main() {
  mosquitto_lib_init();

  struct mosquitto* mosq = mosquitto_new(/*id*/NULL, /*clean_session*/true, /*userdata*/NULL);
  if (!mosq) {
    std::cerr << "mosquitto_new failed\n";
    mosquitto_lib_cleanup();
    return 1;
  }

  mosquitto_connect_callback_set(mosq, on_connect);
  mosquitto_message_callback_set(mosq, on_message);

  int rc = mosquitto_connect(mosq, BROKER, 1883, 60);
  if (rc != MOSQ_ERR_SUCCESS) {
    std::cerr << "mosquitto_connect failed: " << mosquitto_strerror(rc) << "\n";
    mosquitto_destroy(mosq);
    mosquitto_lib_cleanup();
    return 1;
  }

  std::cout << "MQTT Concatenator running …\n";
  mosquitto_loop_forever(mosq, -1, 1);

  mosquitto_destroy(mosq);
  mosquitto_lib_cleanup();
  return 0;
}

