<commit_before>#include "p2ppipe.h"
#include "utils.h"
#include "crypto/sha2.h"

#include <thread>
#include <chrono>
#include <string.h>
#include <unistd.h>
#include <sys/time.h>
#include <stdio.h>

#ifdef WIN32
	#include <winsock2.h>
	#include <ws2tcpip.h>
#else // WIN32
	#include <sys/socket.h>
	#include <netinet/in.h>
	#include <netinet/tcp.h>
	#include <netdb.h>
	#include <fcntl.h>
#endif // !WIN32

void P2PPipe::send_message(const char* command, unsigned char* headerAndData, size_t datalen) {
	prepare_message(command, headerAndData, datalen);
	maybe_do_send_bytes((char*)headerAndData, sizeof(struct bitcoin_msg_header) + datalen);
}

static std::string statefulMessages[] = {"sendheaders", "feefilter" /*, "sendcmpct" - but gets its own processing */};

void P2PPipe::send_hashed_message(const unsigned char* headerAndData, size_t totallen) {
	if (!strncmp(((struct bitcoin_msg_header*)headerAndData)->command, "version", strlen("version"))) {
		std::vector<unsigned char> message(headerAndData, headerAndData + totallen);
		bool insert = !statefulMessagesSent.count("version");
		if (!insert) {
			const unsigned char* oldMessage = statefulMessagesSent["version"].data();

			uint32_t oldMsgLength, newMsgLength;
			memcpy(&oldMsgLength, &((struct bitcoin_msg_header*)oldMessage   )->length, sizeof(oldMsgLength));
			memcpy(&newMsgLength, &((struct bitcoin_msg_header*)headerAndData)->length, sizeof(newMsgLength));
			insert = oldMsgLength != newMsgLength;

			if (!insert && le32toh(newMsgLength) >= sizeof(struct bitcoin_version_start)) {
				struct bitcoin_version_start oldVersion, newVersion;
				memcpy(&oldVersion, oldMessage    + sizeof(struct bitcoin_msg_header), sizeof(oldVersion));
				memcpy(&newVersion, headerAndData + sizeof(struct bitcoin_msg_header), sizeof(newVersion));

				insert = oldVersion.protocol_version != newVersion.protocol_version ||
						oldVersion.services != newVersion.services ||
						memcmp(oldVersion.addr_recv, newVersion.addr_recv, sizeof(oldVersion.addr_recv) - 2) ||
						memcmp(oldVersion.addr_from, newVersion.addr_from, sizeof(oldVersion.addr_from) - 2) ||
						oldVersion.user_agent_length != newVersion.user_agent_length;
			}
		}

		if (insert) {
			statefulMessagesSent["version"] = std::move(message);
			disconnect_from_outside("new outbound version");
		}
		return;
	}

	if (is_connected())
		maybe_do_send_bytes((char*)headerAndData, totallen);
	assert(totallen >= sizeof(struct bitcoin_msg_header));

	for (size_t i = 0; i < sizeof(statefulMessages) / sizeof(std::string); i++)
		if (!strncmp(((struct bitcoin_msg_header*)headerAndData)->command, statefulMessages[i].c_str(), statefulMessages[i].length()))
			statefulMessagesSent[statefulMessages[i]] = std::vector<unsigned char>(headerAndData, headerAndData + totallen);
	if (!strncmp(((struct bitcoin_msg_header*)headerAndData)->command, "sendcmpct", strlen("sendcmpct")) &&
			totallen >= sizeof(struct bitcoin_msg_header) + 9) {
		uint64_t cmpctVersion;
		memcpy(&cmpctVersion, headerAndData + sizeof(struct bitcoin_msg_header) + 1, sizeof(cmpctVersion));
		cmpctVersion = le64toh(cmpctVersion);
		statefulMessagesSent[std::string("sendcmpct") + std::to_string(cmpctVersion)] = std::vector<unsigned char>(headerAndData, headerAndData + totallen);
	}
}

void P2PPipe::on_disconnect() {
	connected = 0;
	read_msg.reset();
	read_pos = 0;
}

void P2PPipe::net_process(const std::function<void(std::string)>& disconnect) {
	connected = 0;

	if (statefulMessagesSent.count("version")) {
		//Rehash the message because send_hashed_message has special processing for version messages
		send_message("version", statefulMessagesSent["version"].data(), statefulMessagesSent["version"].size() - sizeof(struct bitcoin_msg_header));
	} else {
		std::vector<unsigned char> version_msg(generate_version());
		send_message("version", &version_msg[0], version_msg.size() - sizeof(struct bitcoin_msg_header));
	}

	while (true) {
		struct bitcoin_msg_header header;
		if (read_all((char*)&header, sizeof(header)) != sizeof(header))
			return disconnect("failed to read message header");

		if (header.magic != BITCOIN_MAGIC)
			return disconnect("invalid magic bytes");

		header.length = le32toh(header.length);
		if (header.length > 5000000)
			return disconnect("got message too large");

		auto msg = std::make_shared<std::vector<unsigned char> > (sizeof(header) + uint32_t(header.length));
		memcpy(msg->data(), &header, sizeof(header));

		bool check_msg_hash = check_block_msghash || (strncmp(header.command, "block", strlen("block")) &&
			strncmp(header.command, "blocktxn", strlen("blocktxn")) && strncmp(header.command, "cmpctblock", strlen("cmpctblock")));
		if (check_msg_hash) {
			uint32_t hash[8];
			double_sha256_init(hash);

			uint32_t steps = header.length / 64;
			for (uint32_t i = 0; i < steps; i++) {
				unsigned char* writepos = &((*msg)[sizeof(header) + i*64]);
				if (read_all((char*)writepos, 64) != 64)
					return disconnect("failed to read message");
				double_sha256_step(writepos, 64, hash);
			}

			unsigned char* writepos = &((*msg)[sizeof(header) + steps*64]);
			if (read_all((char*)writepos, header.length - steps*64) != ssize_t(header.length - steps*64))
				return disconnect("failed to read message");
			double_sha256_done(writepos, header.length - steps*64, header.length, hash);

			if (memcmp((char*)hash, header.checksum, sizeof(header.checksum)))
				return disconnect("got invalid message checksum");
		} else
			if (read_all((char*)&(*msg)[sizeof(header)], header.length) != ssize_t(header.length))
				return disconnect("failed to read message");

		if (!strncmp(header.command, "version", strlen("version"))) {
			if (connected != 0)
				return disconnect("got invalid version");
			connected = 1;

			if (header.length < sizeof(struct bitcoin_version_start))
				return disconnect("got short version");

			struct bitcoin_version_start their_version;
			memcpy(&their_version, &(*msg)[sizeof(header)], sizeof(their_version));

			struct bitcoin_msg_header new_header;
			send_message("verack", (unsigned char*)&new_header, 0);

			STAMPOUT();
			printf("Connected to bitcoind with version %u\n", le32toh(their_version.protocol_version));

			provide_msg(*msg);

			continue;
		} else if (!strncmp(header.command, "verack", strlen("verack"))) {
			if (connected != 1)
				return disconnect("got invalid verack");
			STAMPOUT();
			printf("Finished connect handshake with bitcoind\n");
			connected = 2;

			for (auto& p : statefulMessagesSent)
				if (p.first != "version") {
					std::string cmpct("sendcmpct");
					if (!p.first.compare(0, cmpct.length(), cmpct))
						send_message("sendcmpct", p.second.data(), p.second.size() - sizeof(struct bitcoin_msg_header));
					else
						send_message(p.first.c_str(), p.second.data(), p.second.size() - sizeof(struct bitcoin_msg_header));
				}

			continue;
		}

		if (connected != 2)
			return disconnect("got non-version, non-verack before version+verack");

		if (!strncmp(header.command, "ping", strlen("ping"))) {
			std::vector<unsigned char> resp(sizeof(struct bitcoin_msg_header) + header.length);
			resp.insert(resp.begin() + sizeof(struct bitcoin_msg_header), msg->begin() + sizeof(header), msg->end());
			send_message("pong", &resp[0], header.length);
		} else if (!strncmp(header.command, "pong", strlen("pong"))) {
			if (msg->size() != 8 + sizeof(header))
				return disconnect("got pong without nonce");
			uint64_t nonce;
			memcpy(&nonce, &(*msg)[sizeof(read_header)], 8);
			pong_received(nonce);
		} else
			provide_msg(*msg);
	}
}

void P2PPipe::send_ping(uint64_t nonce) {
	std::vector<unsigned char> msg(sizeof(struct bitcoin_msg_header) + 8);
	memcpy(&msg[sizeof(struct bitcoin_msg_header)], &nonce, 8);
	send_message("ping", &msg[0], 8);
}

bool P2PPipe::is_connected() const {
	return connected == 2;
}
<commit_msg>Postpend FIBREBentPipe to version string<commit_after>#include "p2ppipe.h"
#include "utils.h"
#include "crypto/sha2.h"

#include <thread>
#include <chrono>
#include <string.h>
#include <unistd.h>
#include <sys/time.h>
#include <stdio.h>
#include <cstddef>

#ifdef WIN32
	#include <winsock2.h>
	#include <ws2tcpip.h>
#else // WIN32
	#include <sys/socket.h>
	#include <netinet/in.h>
	#include <netinet/tcp.h>
	#include <netdb.h>
	#include <fcntl.h>
#endif // !WIN32

void P2PPipe::send_message(const char* command, unsigned char* headerAndData, size_t datalen) {
	prepare_message(command, headerAndData, datalen);
	maybe_do_send_bytes((char*)headerAndData, sizeof(struct bitcoin_msg_header) + datalen);
}

static std::string statefulMessages[] = {"sendheaders", "feefilter" /*, "sendcmpct" - but gets its own processing */};

void P2PPipe::send_hashed_message(const unsigned char* headerAndData, size_t totallen) {
	if (!strncmp(((struct bitcoin_msg_header*)headerAndData)->command, "version", strlen("version"))) {
		std::vector<unsigned char> message(headerAndData, headerAndData + totallen);
		bool insert = !statefulMessagesSent.count("version");
		if (!insert) {
			const unsigned char* oldMessage = statefulMessagesSent["version"].data();

			uint32_t oldMsgLength, newMsgLength;
			memcpy(&oldMsgLength, &((struct bitcoin_msg_header*)oldMessage   )->length, sizeof(oldMsgLength));
			memcpy(&newMsgLength, &((struct bitcoin_msg_header*)headerAndData)->length, sizeof(newMsgLength));
			insert = oldMsgLength != newMsgLength;

			if (!insert && le32toh(newMsgLength) >= sizeof(struct bitcoin_version_start)) {
				struct bitcoin_version_start oldVersion, newVersion;
				memcpy(&oldVersion, oldMessage    + sizeof(struct bitcoin_msg_header), sizeof(oldVersion));
				memcpy(&newVersion, headerAndData + sizeof(struct bitcoin_msg_header), sizeof(newVersion));

				insert = oldVersion.protocol_version != newVersion.protocol_version ||
						oldVersion.services != newVersion.services ||
						memcmp(oldVersion.addr_recv, newVersion.addr_recv, sizeof(oldVersion.addr_recv) - 2) ||
						memcmp(oldVersion.addr_from, newVersion.addr_from, sizeof(oldVersion.addr_from) - 2) ||
						oldVersion.user_agent_length != newVersion.user_agent_length;
			}
		}

		if (insert) {
			statefulMessagesSent["version"] = std::move(message);
			disconnect_from_outside("new outbound version");
		}
		return;
	}

	if (is_connected())
		maybe_do_send_bytes((char*)headerAndData, totallen);
	assert(totallen >= sizeof(struct bitcoin_msg_header));

	for (size_t i = 0; i < sizeof(statefulMessages) / sizeof(std::string); i++)
		if (!strncmp(((struct bitcoin_msg_header*)headerAndData)->command, statefulMessages[i].c_str(), statefulMessages[i].length()))
			statefulMessagesSent[statefulMessages[i]] = std::vector<unsigned char>(headerAndData, headerAndData + totallen);
	if (!strncmp(((struct bitcoin_msg_header*)headerAndData)->command, "sendcmpct", strlen("sendcmpct")) &&
			totallen >= sizeof(struct bitcoin_msg_header) + 9) {
		uint64_t cmpctVersion;
		memcpy(&cmpctVersion, headerAndData + sizeof(struct bitcoin_msg_header) + 1, sizeof(cmpctVersion));
		cmpctVersion = le64toh(cmpctVersion);
		statefulMessagesSent[std::string("sendcmpct") + std::to_string(cmpctVersion)] = std::vector<unsigned char>(headerAndData, headerAndData + totallen);
	}
}

void P2PPipe::on_disconnect() {
	connected = 0;
	read_msg.reset();
	read_pos = 0;
}

void P2PPipe::net_process(const std::function<void(std::string)>& disconnect) {
	connected = 0;

	if (statefulMessagesSent.count("version")) {
		// Postpend Bent Pipe info to version
		std::vector<unsigned char> message(statefulMessagesSent["version"]);
		if (message.size() >= sizeof(struct bitcoin_version_start) + sizeof(struct bitcoin_msg_header)) {
			uint8_t ua_length = message[sizeof(struct bitcoin_msg_header) + offsetof(struct bitcoin_version_start, user_agent_length)];
			if (ua_length <= 100 && message.size() > sizeof(struct bitcoin_version_start) + sizeof(struct bitcoin_msg_header) + ua_length) {
				const char* ua_postpend = "FIBREBentPipe:42/";
				message.insert(message.begin() + sizeof(struct bitcoin_msg_header) + sizeof(struct bitcoin_version_start) + ua_length, (unsigned char*)ua_postpend, (unsigned char*) ua_postpend + 17);
				message[sizeof(struct bitcoin_msg_header) + offsetof(struct bitcoin_version_start, user_agent_length)] = ua_length + 17;
			}
		}
		send_message("version", message.data(), message.size() - sizeof(struct bitcoin_msg_header));
	} else {
		std::vector<unsigned char> version_msg(generate_version());
		send_message("version", &version_msg[0], version_msg.size() - sizeof(struct bitcoin_msg_header));
	}

	while (true) {
		struct bitcoin_msg_header header;
		if (read_all((char*)&header, sizeof(header)) != sizeof(header))
			return disconnect("failed to read message header");

		if (header.magic != BITCOIN_MAGIC)
			return disconnect("invalid magic bytes");

		header.length = le32toh(header.length);
		if (header.length > 5000000)
			return disconnect("got message too large");

		auto msg = std::make_shared<std::vector<unsigned char> > (sizeof(header) + uint32_t(header.length));
		memcpy(msg->data(), &header, sizeof(header));

		bool check_msg_hash = check_block_msghash || (strncmp(header.command, "block", strlen("block")) &&
			strncmp(header.command, "blocktxn", strlen("blocktxn")) && strncmp(header.command, "cmpctblock", strlen("cmpctblock")));
		if (check_msg_hash) {
			uint32_t hash[8];
			double_sha256_init(hash);

			uint32_t steps = header.length / 64;
			for (uint32_t i = 0; i < steps; i++) {
				unsigned char* writepos = &((*msg)[sizeof(header) + i*64]);
				if (read_all((char*)writepos, 64) != 64)
					return disconnect("failed to read message");
				double_sha256_step(writepos, 64, hash);
			}

			unsigned char* writepos = &((*msg)[sizeof(header) + steps*64]);
			if (read_all((char*)writepos, header.length - steps*64) != ssize_t(header.length - steps*64))
				return disconnect("failed to read message");
			double_sha256_done(writepos, header.length - steps*64, header.length, hash);

			if (memcmp((char*)hash, header.checksum, sizeof(header.checksum)))
				return disconnect("got invalid message checksum");
		} else
			if (read_all((char*)&(*msg)[sizeof(header)], header.length) != ssize_t(header.length))
				return disconnect("failed to read message");

		if (!strncmp(header.command, "version", strlen("version"))) {
			if (connected != 0)
				return disconnect("got invalid version");
			connected = 1;

			if (header.length < sizeof(struct bitcoin_version_start))
				return disconnect("got short version");

			struct bitcoin_version_start their_version;
			memcpy(&their_version, &(*msg)[sizeof(header)], sizeof(their_version));

			struct bitcoin_msg_header new_header;
			send_message("verack", (unsigned char*)&new_header, 0);

			STAMPOUT();
			printf("Connected to bitcoind with version %u\n", le32toh(their_version.protocol_version));

			provide_msg(*msg);

			continue;
		} else if (!strncmp(header.command, "verack", strlen("verack"))) {
			if (connected != 1)
				return disconnect("got invalid verack");
			STAMPOUT();
			printf("Finished connect handshake with bitcoind\n");
			connected = 2;

			for (auto& p : statefulMessagesSent)
				if (p.first != "version") {
					std::string cmpct("sendcmpct");
					if (!p.first.compare(0, cmpct.length(), cmpct))
						send_message("sendcmpct", p.second.data(), p.second.size() - sizeof(struct bitcoin_msg_header));
					else
						send_message(p.first.c_str(), p.second.data(), p.second.size() - sizeof(struct bitcoin_msg_header));
				}

			continue;
		}

		if (connected != 2)
			return disconnect("got non-version, non-verack before version+verack");

		if (!strncmp(header.command, "ping", strlen("ping"))) {
			std::vector<unsigned char> resp(sizeof(struct bitcoin_msg_header) + header.length);
			resp.insert(resp.begin() + sizeof(struct bitcoin_msg_header), msg->begin() + sizeof(header), msg->end());
			send_message("pong", &resp[0], header.length);
		} else if (!strncmp(header.command, "pong", strlen("pong"))) {
			if (msg->size() != 8 + sizeof(header))
				return disconnect("got pong without nonce");
			uint64_t nonce;
			memcpy(&nonce, &(*msg)[sizeof(read_header)], 8);
			pong_received(nonce);
		} else
			provide_msg(*msg);
	}
}

void P2PPipe::send_ping(uint64_t nonce) {
	std::vector<unsigned char> msg(sizeof(struct bitcoin_msg_header) + 8);
	memcpy(&msg[sizeof(struct bitcoin_msg_header)], &nonce, 8);
	send_message("ping", &msg[0], 8);
}

bool P2PPipe::is_connected() const {
	return connected == 2;
}
<|endoftext|>


<commit_before>/////////////////////////////////////////////////////////////////////////////
//
// BSD 3-Clause License
//
// Copyright (c) 2019, University of California, San Diego.
// All rights reserved.
//
// Redistribution and use in source and binary forms, with or without
// modification, are permitted provided that the following conditions are met:
//
// * Redistributions of source code must retain the above copyright notice, this
//   list of conditions and the following disclaimer.
//
// * Redistributions in binary form must reproduce the above copyright notice,
//   this list of conditions and the following disclaimer in the documentation
//   and/or other materials provided with the distribution.
//
// * Neither the name of the copyright holder nor the names of its
//   contributors may be used to endorse or promote products derived from
//   this software without specific prior written permission.
//
// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
// AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
// ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
// LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
// CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
// SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
// INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
// CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
// ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
// POSSIBILITY OF SUCH DAMAGE.
//
///////////////////////////////////////////////////////////////////////////////

#include "HungarianMatching.h"

#include "utility/Logger.h"

namespace ppl {

HungarianMatching::HungarianMatching(Section& section,
                                     SlotVector& slots,
                                     Logger* logger)
    : netlist_(section.net), slots_(slots)
{
  num_io_pins_ = netlist_.numIOPins();
  num_pin_groups_ = netlist_.numIOGroups();
  begin_slot_ = section.begin_slot;
  end_slot_ = section.end_slot;
  num_slots_ = end_slot_ - begin_slot_;
  non_blocked_slots_ = section.num_slots;
  group_slots_ = 0;
  group_size_ = -1;
  edge_ = section.edge;
  logger_ = logger;
}

void HungarianMatching::findAssignment(std::vector<Constraint>& constraints)
{
  createMatrix(constraints);
  hungarian_solver_.solve(hungarian_matrix_, assignment_);
}

void HungarianMatching::createMatrix(std::vector<Constraint>& constraints)
{
  hungarian_matrix_.resize(non_blocked_slots_);
  int slot_index = 0;
  for (int i = begin_slot_; i < end_slot_; ++i) {
    int pinIndex = 0;
    Point newPos = slots_[i].pos;
    if (slots_[i].blocked) {
      continue;
    }
    hungarian_matrix_[slot_index].resize(num_io_pins_);
    int idx = 0;
    for (IOPin& io_pin : netlist_.getIOPins()) {
      if (!io_pin.isInGroup()) {
        int hpwl = netlist_.computeIONetHPWL(idx, newPos, edge_, constraints);
        hungarian_matrix_[slot_index][pinIndex] = hpwl;
        pinIndex++;
      }
      idx++;
    }
    slot_index++;
  }
}

inline bool samePos(Point& a, Point& b)
{
  return (a.x() == b.x() && a.y() == b.y());
}

void HungarianMatching::getFinalAssignment(std::vector<IOPin>& assigment) const
{
  size_t rows = non_blocked_slots_;
  size_t col = 0;
  int slot_index = 0;
  for (IOPin& io_pin : netlist_.getIOPins()) {
    if (!io_pin.isInGroup()) {
      slot_index = begin_slot_;
      for (size_t row = 0; row < rows; row++) {
        while (slots_[slot_index].blocked && slot_index < slots_.size())
          slot_index++;
        if (assignment_[row] != col) {
          slot_index++;
          continue;
        }
        if (hungarian_matrix_[row][col] == hungarian_fail) {
          logger_->warn(utl::PPL,
                        33,
                        "I/O pin {} cannot be placed in the specified region. "
                        "Not enough space",
                        io_pin.getName().c_str());
        }
        io_pin.setPos(slots_[slot_index].pos);
        io_pin.setLayer(slots_[slot_index].layer);
        assigment.push_back(io_pin);
        slots_[slot_index].used = true;
        break;
      }
      col++;
    }
  }
}

void HungarianMatching::findAssignmentForGroups(std::vector<Constraint>& constraints)
{
  createMatrixForGroups(constraints);

  if (hungarian_matrix_.size() > 0)
    hungarian_solver_.solve(hungarian_matrix_, assignment_);
}

void HungarianMatching::createMatrixForGroups(std::vector<Constraint>& constraints)
{
  for (std::set<int>& io_group : netlist_.getIOGroups()) {
    group_size_ = std::max((int)io_group.size(), group_size_);
  }

  if (group_size_ > 0) {
    for (int i = begin_slot_; i < end_slot_; i+=group_size_) {
      bool blocked = false;
      for (int pin_cnt = 0; pin_cnt < group_size_; pin_cnt++) {
        if (slots_[i + pin_cnt].blocked) {
          blocked = true;
        }
      }
      if (!blocked) {
        group_slots_++;
      }
    }

    hungarian_matrix_.resize(group_slots_);
    int slot_index = 0;
    for (int i = begin_slot_; i < end_slot_; i+=group_size_) {
      int groupIndex = 0;
      Point newPos = slots_[i].pos;

      bool blocked = false;
      for (int pin_cnt = 0; pin_cnt < group_size_; pin_cnt++) {
        if (slots_[i + pin_cnt].blocked) {
          blocked = true;
        }
      }
      if (blocked) {
        continue;
      }

      hungarian_matrix_[slot_index].resize(num_pin_groups_);
      for (std::set<int>& io_group : netlist_.getIOGroups()) {
        int hpwl = 0;
        int pin_count = 0;
        for (int io_idx : io_group) {
          hpwl += netlist_.computeIONetHPWL(io_idx, newPos, edge_, constraints);
          hungarian_matrix_[slot_index][groupIndex] = hpwl;
        }
        groupIndex++;
      }
      slot_index++;
    }
  }
}

void HungarianMatching::getAssignmentForGroups(std::vector<IOPin>& assigment)
{
  if (hungarian_matrix_.size() <= 0)
    return;

  size_t rows = group_slots_;
  size_t col = 0;
  int slot_index = 0;
  for (std::set<int>& io_group : netlist_.getIOGroups()) {
    slot_index = begin_slot_;
    for (size_t row = 0; row < rows; row++) {
      while (slots_[slot_index].blocked && slot_index < slots_.size())
        slot_index++;
      if (assignment_[row] != col) {
        slot_index++;
        continue;
      }
      int pin_cnt = 0;
      for (int pin_idx : io_group) {
        IOPin io_pin = netlist_.getIoPin(pin_idx);
        io_pin.setPos(slots_[slot_index + pin_cnt].pos);
        io_pin.setLayer(slots_[slot_index + pin_cnt].layer);
        assigment.push_back(io_pin);
        slots_[slot_index + pin_cnt].used = true;
        slots_[slot_index + pin_cnt].blocked = true;
        non_blocked_slots_--;
        pin_cnt++;
      }
      break;
    }
    col++;
  }

  hungarian_matrix_.clear();
  assignment_.clear();
}

}  // namespace ppl
<commit_msg>ppl: avoid negative values in hungariam matching matrix<commit_after>/////////////////////////////////////////////////////////////////////////////
//
// BSD 3-Clause License
//
// Copyright (c) 2019, University of California, San Diego.
// All rights reserved.
//
// Redistribution and use in source and binary forms, with or without
// modification, are permitted provided that the following conditions are met:
//
// * Redistributions of source code must retain the above copyright notice, this
//   list of conditions and the following disclaimer.
//
// * Redistributions in binary form must reproduce the above copyright notice,
//   this list of conditions and the following disclaimer in the documentation
//   and/or other materials provided with the distribution.
//
// * Neither the name of the copyright holder nor the names of its
//   contributors may be used to endorse or promote products derived from
//   this software without specific prior written permission.
//
// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
// AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
// ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
// LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
// CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
// SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
// INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
// CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
// ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
// POSSIBILITY OF SUCH DAMAGE.
//
///////////////////////////////////////////////////////////////////////////////

#include "HungarianMatching.h"

#include "utility/Logger.h"

namespace ppl {

HungarianMatching::HungarianMatching(Section& section,
                                     SlotVector& slots,
                                     Logger* logger)
    : netlist_(section.net), slots_(slots)
{
  num_io_pins_ = netlist_.numIOPins();
  num_pin_groups_ = netlist_.numIOGroups();
  begin_slot_ = section.begin_slot;
  end_slot_ = section.end_slot;
  num_slots_ = end_slot_ - begin_slot_;
  non_blocked_slots_ = section.num_slots;
  group_slots_ = 0;
  group_size_ = -1;
  edge_ = section.edge;
  logger_ = logger;
}

void HungarianMatching::findAssignment(std::vector<Constraint>& constraints)
{
  createMatrix(constraints);
  hungarian_solver_.solve(hungarian_matrix_, assignment_);
}

void HungarianMatching::createMatrix(std::vector<Constraint>& constraints)
{
  hungarian_matrix_.resize(non_blocked_slots_);
  int slot_index = 0;
  for (int i = begin_slot_; i < end_slot_; ++i) {
    int pinIndex = 0;
    Point newPos = slots_[i].pos;
    if (slots_[i].blocked) {
      continue;
    }
    hungarian_matrix_[slot_index].resize(num_io_pins_);
    int idx = 0;
    for (IOPin& io_pin : netlist_.getIOPins()) {
      if (!io_pin.isInGroup()) {
        int hpwl = netlist_.computeIONetHPWL(idx, newPos, edge_, constraints);
        hungarian_matrix_[slot_index][pinIndex] = hpwl;
        pinIndex++;
      }
      idx++;
    }
    slot_index++;
  }
}

inline bool samePos(Point& a, Point& b)
{
  return (a.x() == b.x() && a.y() == b.y());
}

void HungarianMatching::getFinalAssignment(std::vector<IOPin>& assigment) const
{
  size_t rows = non_blocked_slots_;
  size_t col = 0;
  int slot_index = 0;
  for (IOPin& io_pin : netlist_.getIOPins()) {
    if (!io_pin.isInGroup()) {
      slot_index = begin_slot_;
      for (size_t row = 0; row < rows; row++) {
        while (slots_[slot_index].blocked && slot_index < slots_.size())
          slot_index++;
        if (assignment_[row] != col) {
          slot_index++;
          continue;
        }
        if (hungarian_matrix_[row][col] == hungarian_fail) {
          logger_->warn(utl::PPL,
                        33,
                        "I/O pin {} cannot be placed in the specified region. "
                        "Not enough space",
                        io_pin.getName().c_str());
        }
        io_pin.setPos(slots_[slot_index].pos);
        io_pin.setLayer(slots_[slot_index].layer);
        assigment.push_back(io_pin);
        slots_[slot_index].used = true;
        break;
      }
      col++;
    }
  }
}

void HungarianMatching::findAssignmentForGroups(std::vector<Constraint>& constraints)
{
  createMatrixForGroups(constraints);

  if (hungarian_matrix_.size() > 0)
    hungarian_solver_.solve(hungarian_matrix_, assignment_);
}

void HungarianMatching::createMatrixForGroups(std::vector<Constraint>& constraints)
{
  for (std::set<int>& io_group : netlist_.getIOGroups()) {
    group_size_ = std::max((int)io_group.size(), group_size_);
  }

  if (group_size_ > 0) {
    for (int i = begin_slot_; i < end_slot_; i+=group_size_) {
      bool blocked = false;
      for (int pin_cnt = 0; pin_cnt < group_size_; pin_cnt++) {
        if (slots_[i + pin_cnt].blocked) {
          blocked = true;
        }
      }
      if (!blocked) {
        group_slots_++;
      }
    }

    hungarian_matrix_.resize(group_slots_);
    int slot_index = 0;
    for (int i = begin_slot_; i < end_slot_; i+=group_size_) {
      int groupIndex = 0;
      Point newPos = slots_[i].pos;

      bool blocked = false;
      for (int pin_cnt = 0; pin_cnt < group_size_; pin_cnt++) {
        if (slots_[i + pin_cnt].blocked) {
          blocked = true;
        }
      }
      if (blocked) {
        continue;
      }

      hungarian_matrix_[slot_index].resize(num_pin_groups_);
      for (std::set<int>& io_group : netlist_.getIOGroups()) {
        int group_hpwl = 0;
        int pin_count = 0;
        for (int io_idx : io_group) {
          int pin_hpwl = netlist_.computeIONetHPWL(io_idx, newPos, edge_, constraints);
          if (pin_hpwl == hungarian_fail) {
            group_hpwl = hungarian_fail;
            break;
          } else {
            group_hpwl += netlist_.computeIONetHPWL(io_idx, newPos, edge_, constraints);
          }
        }
        hungarian_matrix_[slot_index][groupIndex] = group_hpwl;
        groupIndex++;
      }
      slot_index++;
    }
  }
}

void HungarianMatching::getAssignmentForGroups(std::vector<IOPin>& assigment)
{
  if (hungarian_matrix_.size() <= 0)
    return;

  size_t rows = group_slots_;
  size_t col = 0;
  int slot_index = 0;
  for (std::set<int>& io_group : netlist_.getIOGroups()) {
    slot_index = begin_slot_;
    for (size_t row = 0; row < rows; row++) {
      while (slots_[slot_index].blocked && slot_index < slots_.size())
        slot_index++;
      if (assignment_[row] != col) {
        slot_index++;
        continue;
      }
      int pin_cnt = 0;
      for (int pin_idx : io_group) {
        IOPin io_pin = netlist_.getIoPin(pin_idx);
        io_pin.setPos(slots_[slot_index + pin_cnt].pos);
        io_pin.setLayer(slots_[slot_index + pin_cnt].layer);
        assigment.push_back(io_pin);
        slots_[slot_index + pin_cnt].used = true;
        slots_[slot_index + pin_cnt].blocked = true;
        non_blocked_slots_--;
        pin_cnt++;
      }
      break;
    }
    col++;
  }

  hungarian_matrix_.clear();
  assignment_.clear();
}

}  // namespace ppl
<|endoftext|>


<commit_before>#include <cmath>
#include <thread>
#include <string>
#include <sstream>
#include <algorithm>
#include <functional>

#include "Tools/general_utils.h"
#include "Tools/system_functions.h"
#include "Tools/Display/bash_tools.h"
#include "Tools/Exception/exception.hpp"
#include "Tools/Display/Statistics/Statistics.hpp"
#include "Tools/Display/Terminal/BFER/Terminal_BFER.hpp"

#ifdef ENABLE_MPI
#include "Module/Monitor/BFER/Monitor_BFER_reduction_mpi.hpp"
#endif

#include "Factory/Module/Monitor/Monitor.hpp"
#include "Factory/Tools/Display/Terminal/BFER/Terminal_BFER.hpp"

#include "BFER.hpp"

using namespace aff3ct;
using namespace aff3ct::simulation;

template <typename B, typename R, typename Q>
BFER<B,R,Q>
::BFER(const factory::BFER::parameters& params_BFER)
: Simulation(params_BFER),
  params_BFER(params_BFER),

  barrier(params_BFER.n_threads),

  bit_rate((float)params_BFER.src->K / (float)params_BFER.cdc->N),

  max_fra(0),

  monitor    (params_BFER.n_threads, nullptr),
  monitor_red(                       nullptr),
  dumper     (params_BFER.n_threads, nullptr),
  dumper_red (                       nullptr),
  terminal   (                       nullptr)
{
	if (params_BFER.n_threads < 1)
	{
		std::stringstream message;
		message << "'n_threads' has to be greater than 0 ('n_threads' = " << params_BFER.n_threads << ").";
		throw tools::invalid_argument(__FILE__, __LINE__, __func__, message.str());
	}

	if (params_BFER.err_track_enable)
	{
		for (auto tid = 0; tid < params_BFER.n_threads; tid++)
			dumper[tid] = new tools::Dumper();

		std::vector<tools::Dumper*> dumpers;
		for (auto tid = 0; tid < params_BFER.n_threads; tid++)
			dumpers.push_back(dumper[tid]);

		dumper_red = new tools::Dumper_reduction(dumpers);
	}

	modules["monitor"] = std::vector<module::Module*>(params_BFER.n_threads, nullptr);
	for (auto tid = 0; tid < params_BFER.n_threads; tid++)
	{
		this->monitor[tid] = this->build_monitor(tid);
		modules["monitor"][tid] = this->monitor[tid];
	}

#ifdef ENABLE_MPI
	// build a monitor to compute BER/FER (reduce the other monitors)
	this->monitor_red = new module::Monitor_BFER_reduction_mpi<B>(this->monitor,
	                                                              std::this_thread::get_id(),
	                                                              params_BFER.mpi_comm_freq);
#else
	// build a monitor to compute BER/FER (reduce the other monitors)
	this->monitor_red = new module::Monitor_BFER_reduction<B>(this->monitor);
#endif
}

template <typename B, typename R, typename Q>
BFER<B,R,Q>
::~BFER()
{
	release_objects();

	if (monitor_red != nullptr) { delete monitor_red; monitor_red = nullptr; }
	if (dumper_red  != nullptr) { delete dumper_red;  dumper_red  = nullptr; }

	for (auto tid = 0; tid < params_BFER.n_threads; tid++)
	{
		if (monitor[tid] != nullptr) { delete monitor[tid]; monitor[tid] = nullptr; }
		if (dumper [tid] != nullptr) { delete dumper [tid]; dumper [tid] = nullptr; }
	}

	if (terminal != nullptr) { delete terminal; terminal = nullptr; }
}

template <typename B, typename R, typename Q>
void BFER<B,R,Q>
::_build_communication_chain()
{
	// build the communication chain in multi-threaded mode
	std::vector<std::thread> threads(params_BFER.n_threads -1);
	for (auto tid = 1; tid < params_BFER.n_threads; tid++)
		threads[tid -1] = std::thread(BFER<B,R,Q>::start_thread_build_comm_chain, this, tid);

	BFER<B,R,Q>::start_thread_build_comm_chain(this, 0);

	// join the slave threads with the master thread
	for (auto tid = 1; tid < params_BFER.n_threads; tid++)
		threads[tid -1].join();
}

template <typename B, typename R, typename Q>
void BFER<B,R,Q>
::launch()
{
	this->terminal = this->build_terminal();

	if (!this->params_BFER.err_track_revert)
	{
		this->build_communication_chain();

		if (module::Monitor::is_over())
		{
			this->release_objects();
			return;
		}
	}

	// for each NOISE to be simulated
	for (unsigned noise_idx = 0; noise_idx < params_BFER.noise_range.size(); noise_idx ++)
	{
		auto n = params_BFER.noise_range[noise_idx];

		if (params_BFER.noise_type == "EBN0" || params_BFER.noise_type == "ESN0")
		{
			float esn0, ebn0;
			if (params_BFER.noise_type == "EBN0")
			{
				ebn0 = n;
				esn0 = tools::ebn0_to_esn0(ebn0, bit_rate, params_BFER.mdm->bps);
			}
			else // if (params_BFER.sim->noise_type == "ESN0")
			{
				esn0 = n;
				ebn0 = tools::esn0_to_ebn0(esn0, bit_rate, params_BFER.mdm->bps);
			}

			auto sigma = tools::esn0_to_sigma(esn0, params_BFER.mdm->upf);

			this->noise.set_sigma(sigma, ebn0, esn0);
		}
		else if (params_BFER.noise_type == "ROP")
		{
			this->noise.set_rop(n);
		}
		else // if (params_BFER.noise_type == "EP")
		{
			this->noise.set_ep(n);
		}

		this->terminal->set_noise(this->noise);

		if (this->params_BFER.err_track_revert)
		{
			this->release_objects();
			this->monitor_red->clear_callbacks();

			// dirty hack to override simulation params_BFER
			auto &params_BFER_writable = const_cast<factory::BFER::parameters&>(params_BFER);

			std::stringstream s_noise;
			s_noise << std::setprecision(2) << std::fixed << this->noise.get_noise();

			params_BFER_writable.src     ->path = params_BFER.err_track_path + "_" + s_noise.str() + ".src";
			params_BFER_writable.cdc->enc->path = params_BFER.err_track_path + "_" + s_noise.str() + ".enc";
			params_BFER_writable.chn     ->path = params_BFER.err_track_path + "_" + s_noise.str() + ".chn";

			std::ifstream file(params_BFER.chn->path, std::ios::binary);
			if (file.is_open())
			{
				file.read((char*)&max_fra, sizeof(max_fra));
				file.close();
			}
			else
			{
				std::stringstream message;
				message << "Impossible to read the 'chn' file ('chn' = " << params_BFER.chn->path << ").";
				throw tools::runtime_error(__FILE__, __LINE__, __func__, message.str());
			}

			this->build_communication_chain();

			if (module::Monitor::is_over())
			{
				this->release_objects();
				return;
			}
		}

		if (params_BFER.display_legend)
#ifdef ENABLE_MPI
			if (((!params_BFER.ter->disabled && noise_idx == 0 && !params_BFER.debug) ||
		 	   (params_BFER.statistics && !params_BFER.debug)) && params_BFER.mpi_rank == 0)
#else
			if (((!params_BFER.ter->disabled && noise_idx == 0 && !params_BFER.debug) ||
			    (params_BFER.statistics && !params_BFER.debug)))
#endif
				terminal->legend(std::cout);

		// start the terminal to display BER/FER results
#ifdef ENABLE_MPI
		if (!params_BFER.ter->disabled && params_BFER.ter->frequency != std::chrono::nanoseconds(0) && !params_BFER.debug
		    && params_BFER.mpi_rank == 0)
#else
		if (!params_BFER.ter->disabled && params_BFER.ter->frequency != std::chrono::nanoseconds(0) && !params_BFER.debug)
#endif
			terminal->start_temp_report(params_BFER.ter->frequency);

		try
		{
			this->_launch();
		}
		catch (std::exception const& e)
		{
			module::Monitor::stop();
			terminal->final_report(std::cout); // display final report to not lost last line overwritten by the error messages

			std::cerr << tools::apply_on_each_line(e.what(), &tools::format_error) << std::endl;
			this->simu_error = true;
		}

#ifdef ENABLE_MPI
		if (!params_BFER.ter->disabled && terminal != nullptr && !this->simu_error && params_BFER.mpi_rank == 0)
#else
		if (!params_BFER.ter->disabled && terminal != nullptr && !this->simu_error)
#endif
		{
			if (params_BFER.debug)
				terminal->legend(std::cout);

			terminal->final_report(std::cout);

			if (params_BFER.statistics)
			{
				std::vector<std::vector<const module::Module*>> mod_vec;
				for (auto &vm : modules)
				{
					std::vector<const module::Module*> sub_mod_vec;
					for (auto *m : vm.second)
						sub_mod_vec.push_back(m);
					mod_vec.push_back(sub_mod_vec);
				}

				std::cout << "#" << std::endl;
				tools::Stats::show(mod_vec, true, std::cout);
				std::cout << "#" << std::endl;
			}
		}

		if (this->dumper_red != nullptr && !this->simu_error)
		{
			std::stringstream s_noise;
			s_noise << std::setprecision(2) << std::fixed << this->noise.get_noise();

			this->dumper_red->dump(params_BFER.err_track_path + "_" + s_noise.str());
			this->dumper_red->clear();
		}

		if (!params_BFER.err_track_revert && !module::Monitor::is_interrupt() &&
		    this->monitor_red->get_n_fe() < this->monitor_red->get_fe_limit() &&
		    (max_fra == 0 || this->monitor_red->get_n_fe() < max_fra))
			module::Monitor::stop();

		this->monitor_red->reset();
		for (auto &m : modules)
			for (auto mm : m.second)
				if (mm != nullptr)
					for (auto &t : mm->tasks)
						t->reset_stats();

		if (module::Monitor::is_over())
			break;
	}

	this->release_objects();
}

template <typename B, typename R, typename Q>
void BFER<B,R,Q>
::release_objects()
{
}

template <typename B, typename R, typename Q>
module::Monitor_BFER<B>* BFER<B,R,Q>
::build_monitor(const int tid)
{
	return factory::Monitor_BFER::build<B>(*params_BFER.mnt);
}

template <typename B, typename R, typename Q>
tools::Terminal_BFER<B>* BFER<B,R,Q>
::build_terminal()
{
	return factory::Terminal_BFER::build<B>(*params_BFER.ter, *this->monitor_red);
}

template <typename B, typename R, typename Q>
void BFER<B,R,Q>
::start_thread_build_comm_chain(BFER<B,R,Q> *simu, const int tid)
{
	try
	{
		simu->__build_communication_chain(tid);

		if (simu->params_BFER.err_track_enable)
			simu->monitor[tid]->add_handler_fe(std::bind(&tools::Dumper::add, simu->dumper[tid], std::placeholders::_1, std::placeholders::_2));
	}
	catch (std::exception const& e)
	{
		module::Monitor::stop();
		simu->simu_error = true;

		simu->mutex_exception.lock();

		auto save = tools::exception::no_backtrace;
		tools::exception::no_backtrace = true;
		std::string msg = e.what(); // get only the function signature
		tools::exception::no_backtrace = save;

		if (std::find(simu->prev_err_messages.begin(), simu->prev_err_messages.end(), msg) == simu->prev_err_messages.end())
		{
			std::cerr << tools::apply_on_each_line(e.what(), &tools::format_error) << std::endl; // with backtrace if debug mode
			simu->prev_err_messages.push_back(msg); // save only the function signature
		}
		simu->mutex_exception.unlock();
	}
}

// ==================================================================================== explicit template instantiation
#include "Tools/types.h"
#ifdef MULTI_PREC
template class aff3ct::simulation::BFER<B_8,R_8,Q_8>;
template class aff3ct::simulation::BFER<B_16,R_16,Q_16>;
template class aff3ct::simulation::BFER<B_32,R_32,Q_32>;
template class aff3ct::simulation::BFER<B_64,R_64,Q_64>;
#else
template class aff3ct::simulation::BFER<B,R,Q>;
#endif
// ==================================================================================== explicit template instantiation

<commit_msg>Simulate EP in reverse order; Add a noise type check else throw<commit_after>#include <cmath>
#include <thread>
#include <string>
#include <sstream>
#include <algorithm>
#include <functional>

#include "Tools/general_utils.h"
#include "Tools/system_functions.h"
#include "Tools/Display/bash_tools.h"
#include "Tools/Exception/exception.hpp"
#include "Tools/Display/Statistics/Statistics.hpp"
#include "Tools/Display/Terminal/BFER/Terminal_BFER.hpp"

#ifdef ENABLE_MPI
#include "Module/Monitor/BFER/Monitor_BFER_reduction_mpi.hpp"
#endif

#include "Factory/Module/Monitor/Monitor.hpp"
#include "Factory/Tools/Display/Terminal/BFER/Terminal_BFER.hpp"

#include "BFER.hpp"

using namespace aff3ct;
using namespace aff3ct::simulation;

template <typename B, typename R, typename Q>
BFER<B,R,Q>
::BFER(const factory::BFER::parameters& params_BFER)
: Simulation(params_BFER),
  params_BFER(params_BFER),

  barrier(params_BFER.n_threads),

  bit_rate((float)params_BFER.src->K / (float)params_BFER.cdc->N),

  max_fra(0),

  monitor    (params_BFER.n_threads, nullptr),
  monitor_red(                       nullptr),
  dumper     (params_BFER.n_threads, nullptr),
  dumper_red (                       nullptr),
  terminal   (                       nullptr)
{
	if (params_BFER.n_threads < 1)
	{
		std::stringstream message;
		message << "'n_threads' has to be greater than 0 ('n_threads' = " << params_BFER.n_threads << ").";
		throw tools::invalid_argument(__FILE__, __LINE__, __func__, message.str());
	}

	if (params_BFER.err_track_enable)
	{
		for (auto tid = 0; tid < params_BFER.n_threads; tid++)
			dumper[tid] = new tools::Dumper();

		std::vector<tools::Dumper*> dumpers;
		for (auto tid = 0; tid < params_BFER.n_threads; tid++)
			dumpers.push_back(dumper[tid]);

		dumper_red = new tools::Dumper_reduction(dumpers);
	}

	modules["monitor"] = std::vector<module::Module*>(params_BFER.n_threads, nullptr);
	for (auto tid = 0; tid < params_BFER.n_threads; tid++)
	{
		this->monitor[tid] = this->build_monitor(tid);
		modules["monitor"][tid] = this->monitor[tid];
	}

#ifdef ENABLE_MPI
	// build a monitor to compute BER/FER (reduce the other monitors)
	this->monitor_red = new module::Monitor_BFER_reduction_mpi<B>(this->monitor,
	                                                              std::this_thread::get_id(),
	                                                              params_BFER.mpi_comm_freq);
#else
	// build a monitor to compute BER/FER (reduce the other monitors)
	this->monitor_red = new module::Monitor_BFER_reduction<B>(this->monitor);
#endif
}

template <typename B, typename R, typename Q>
BFER<B,R,Q>
::~BFER()
{
	release_objects();

	if (monitor_red != nullptr) { delete monitor_red; monitor_red = nullptr; }
	if (dumper_red  != nullptr) { delete dumper_red;  dumper_red  = nullptr; }

	for (auto tid = 0; tid < params_BFER.n_threads; tid++)
	{
		if (monitor[tid] != nullptr) { delete monitor[tid]; monitor[tid] = nullptr; }
		if (dumper [tid] != nullptr) { delete dumper [tid]; dumper [tid] = nullptr; }
	}

	if (terminal != nullptr) { delete terminal; terminal = nullptr; }
}

template <typename B, typename R, typename Q>
void BFER<B,R,Q>
::_build_communication_chain()
{
	// build the communication chain in multi-threaded mode
	std::vector<std::thread> threads(params_BFER.n_threads -1);
	for (auto tid = 1; tid < params_BFER.n_threads; tid++)
		threads[tid -1] = std::thread(BFER<B,R,Q>::start_thread_build_comm_chain, this, tid);

	BFER<B,R,Q>::start_thread_build_comm_chain(this, 0);

	// join the slave threads with the master thread
	for (auto tid = 1; tid < params_BFER.n_threads; tid++)
		threads[tid -1].join();
}

template <typename B, typename R, typename Q>
void BFER<B,R,Q>
::launch()
{
	this->terminal = this->build_terminal();

	if (!this->params_BFER.err_track_revert)
	{
		this->build_communication_chain();

		if (module::Monitor::is_over())
		{
			this->release_objects();
			return;
		}
	}

	// for each NOISE to be simulated
	for (unsigned noise_idx = 0; noise_idx < params_BFER.noise_range.size(); noise_idx ++)
	{
		auto n = params_BFER.noise_range[noise_idx];

		if (params_BFER.noise_type == "EBN0" || params_BFER.noise_type == "ESN0")
		{
			float esn0, ebn0;
			if (params_BFER.noise_type == "EBN0")
			{
				ebn0 = n;
				esn0 = tools::ebn0_to_esn0(ebn0, bit_rate, params_BFER.mdm->bps);
			}
			else // if (params_BFER.sim->noise_type == "ESN0")
			{
				esn0 = n;
				ebn0 = tools::esn0_to_ebn0(esn0, bit_rate, params_BFER.mdm->bps);
			}

			auto sigma = tools::esn0_to_sigma(esn0, params_BFER.mdm->upf);

			this->noise.set_sigma(sigma, ebn0, esn0);
		}
		else if (params_BFER.noise_type == "ROP")
		{
			this->noise.set_rop(n);
		}
		else if (params_BFER.noise_type == "EP")
		{
            n = params_BFER.noise_range[params_BFER.noise_range.size() - noise_idx -1];

            this->noise.set_ep(n);
		}
		else
        {
            std::stringstream message;
            message << "Unknown noise type ('noise_type' = " << params_BFER.noise_type << ").";
            throw tools::runtime_error(__FILE__, __LINE__, __func__, message.str());
        }

		this->terminal->set_noise(this->noise);

		if (this->params_BFER.err_track_revert)
		{
			this->release_objects();
			this->monitor_red->clear_callbacks();

			// dirty hack to override simulation params_BFER
			auto &params_BFER_writable = const_cast<factory::BFER::parameters&>(params_BFER);

			std::stringstream s_noise;
			s_noise << std::setprecision(2) << std::fixed << this->noise.get_noise();

			params_BFER_writable.src     ->path = params_BFER.err_track_path + "_" + s_noise.str() + ".src";
			params_BFER_writable.cdc->enc->path = params_BFER.err_track_path + "_" + s_noise.str() + ".enc";
			params_BFER_writable.chn     ->path = params_BFER.err_track_path + "_" + s_noise.str() + ".chn";

			std::ifstream file(params_BFER.chn->path, std::ios::binary);
			if (file.is_open())
			{
				file.read((char*)&max_fra, sizeof(max_fra));
				file.close();
			}
			else
			{
				std::stringstream message;
				message << "Impossible to read the 'chn' file ('chn' = " << params_BFER.chn->path << ").";
				throw tools::runtime_error(__FILE__, __LINE__, __func__, message.str());
			}

			this->build_communication_chain();

			if (module::Monitor::is_over())
			{
				this->release_objects();
				return;
			}
		}

		if (params_BFER.display_legend)
#ifdef ENABLE_MPI
			if (((!params_BFER.ter->disabled && noise_idx == 0 && !params_BFER.debug) ||
		 	   (params_BFER.statistics && !params_BFER.debug)) && params_BFER.mpi_rank == 0)
#else
			if (((!params_BFER.ter->disabled && noise_idx == 0 && !params_BFER.debug) ||
			    (params_BFER.statistics && !params_BFER.debug)))
#endif
				terminal->legend(std::cout);

		// start the terminal to display BER/FER results
#ifdef ENABLE_MPI
		if (!params_BFER.ter->disabled && params_BFER.ter->frequency != std::chrono::nanoseconds(0) && !params_BFER.debug
		    && params_BFER.mpi_rank == 0)
#else
		if (!params_BFER.ter->disabled && params_BFER.ter->frequency != std::chrono::nanoseconds(0) && !params_BFER.debug)
#endif
			terminal->start_temp_report(params_BFER.ter->frequency);

		try
		{
			this->_launch();
		}
		catch (std::exception const& e)
		{
			module::Monitor::stop();
			terminal->final_report(std::cout); // display final report to not lost last line overwritten by the error messages

			std::cerr << tools::apply_on_each_line(e.what(), &tools::format_error) << std::endl;
			this->simu_error = true;
		}

#ifdef ENABLE_MPI
		if (!params_BFER.ter->disabled && terminal != nullptr && !this->simu_error && params_BFER.mpi_rank == 0)
#else
		if (!params_BFER.ter->disabled && terminal != nullptr && !this->simu_error)
#endif
		{
			if (params_BFER.debug)
				terminal->legend(std::cout);

			terminal->final_report(std::cout);

			if (params_BFER.statistics)
			{
				std::vector<std::vector<const module::Module*>> mod_vec;
				for (auto &vm : modules)
				{
					std::vector<const module::Module*> sub_mod_vec;
					for (auto *m : vm.second)
						sub_mod_vec.push_back(m);
					mod_vec.push_back(sub_mod_vec);
				}

				std::cout << "#" << std::endl;
				tools::Stats::show(mod_vec, true, std::cout);
				std::cout << "#" << std::endl;
			}
		}

		if (this->dumper_red != nullptr && !this->simu_error)
		{
			std::stringstream s_noise;
			s_noise << std::setprecision(2) << std::fixed << this->noise.get_noise();

			this->dumper_red->dump(params_BFER.err_track_path + "_" + s_noise.str());
			this->dumper_red->clear();
		}

		if (!params_BFER.err_track_revert && !module::Monitor::is_interrupt() &&
		    this->monitor_red->get_n_fe() < this->monitor_red->get_fe_limit() &&
		    (max_fra == 0 || this->monitor_red->get_n_fe() < max_fra))
			module::Monitor::stop();

		this->monitor_red->reset();
		for (auto &m : modules)
			for (auto mm : m.second)
				if (mm != nullptr)
					for (auto &t : mm->tasks)
						t->reset_stats();

		if (module::Monitor::is_over())
			break;
	}

	this->release_objects();
}

template <typename B, typename R, typename Q>
void BFER<B,R,Q>
::release_objects()
{
}

template <typename B, typename R, typename Q>
module::Monitor_BFER<B>* BFER<B,R,Q>
::build_monitor(const int tid)
{
	return factory::Monitor_BFER::build<B>(*params_BFER.mnt);
}

template <typename B, typename R, typename Q>
tools::Terminal_BFER<B>* BFER<B,R,Q>
::build_terminal()
{
	return factory::Terminal_BFER::build<B>(*params_BFER.ter, *this->monitor_red);
}

template <typename B, typename R, typename Q>
void BFER<B,R,Q>
::start_thread_build_comm_chain(BFER<B,R,Q> *simu, const int tid)
{
	try
	{
		simu->__build_communication_chain(tid);

		if (simu->params_BFER.err_track_enable)
			simu->monitor[tid]->add_handler_fe(std::bind(&tools::Dumper::add, simu->dumper[tid], std::placeholders::_1, std::placeholders::_2));
	}
	catch (std::exception const& e)
	{
		module::Monitor::stop();
		simu->simu_error = true;

		simu->mutex_exception.lock();

		auto save = tools::exception::no_backtrace;
		tools::exception::no_backtrace = true;
		std::string msg = e.what(); // get only the function signature
		tools::exception::no_backtrace = save;

		if (std::find(simu->prev_err_messages.begin(), simu->prev_err_messages.end(), msg) == simu->prev_err_messages.end())
		{
			std::cerr << tools::apply_on_each_line(e.what(), &tools::format_error) << std::endl; // with backtrace if debug mode
			simu->prev_err_messages.push_back(msg); // save only the function signature
		}
		simu->mutex_exception.unlock();
	}
}

// ==================================================================================== explicit template instantiation
#include "Tools/types.h"
#ifdef MULTI_PREC
template class aff3ct::simulation::BFER<B_8,R_8,Q_8>;
template class aff3ct::simulation::BFER<B_16,R_16,Q_16>;
template class aff3ct::simulation::BFER<B_32,R_32,Q_32>;
template class aff3ct::simulation::BFER<B_64,R_64,Q_64>;
#else
template class aff3ct::simulation::BFER<B,R,Q>;
#endif
// ==================================================================================== explicit template instantiation

<|endoftext|>


<commit_before>/*
 * Copyright 2018 The Cartographer Authors
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#include "cartographer/mapping/2d/tsdf_range_data_inserter_2d.h"

#include "cartographer/mapping/internal/2d/normal_estimation_2d.h"
#include "cartographer/mapping/internal/2d/ray_to_pixel_mask.h"

namespace cartographer {
namespace mapping {
namespace {

// Factor for subpixel accuracy of start and end point for ray casts.
constexpr int kSubpixelScale = 1000;
// Minimum distance between range observation and origin. Otherwise, range
// observations are discarded.
constexpr float kMinRangeMeters = 1e-6f;
constexpr float kSqrtTwoPi = std::sqrt(2.0 * M_PI);

void GrowAsNeeded(const sensor::RangeData& range_data,
                  const float truncation_distance, TSDF2D* const tsdf) {
  Eigen::AlignedBox2f bounding_box(range_data.origin.head<2>());
  for (const Eigen::Vector3f& hit : range_data.returns) {
    const Eigen::Vector3f direction = (hit - range_data.origin).normalized();
    const Eigen::Vector3f end_position = hit + truncation_distance * direction;
    bounding_box.extend(end_position.head<2>());
  }
  // Padding around bounding box to avoid numerical issues at cell boundaries.
  constexpr float kPadding = 1e-6f;
  tsdf->GrowLimits(bounding_box.min() - kPadding * Eigen::Vector2f::Ones());
  tsdf->GrowLimits(bounding_box.max() + kPadding * Eigen::Vector2f::Ones());
}

float GaussianKernel(const float x, const float sigma) {
  return 1.0 / (kSqrtTwoPi * sigma) * std::exp(-0.5 * x * x / (sigma * sigma));
}

std::pair<Eigen::Array2i, Eigen::Array2i> SuperscaleRay(
    const Eigen::Vector2f& begin, const Eigen::Vector2f& end,
    TSDF2D* const tsdf) {
  const MapLimits& limits = tsdf->limits();
  const double superscaled_resolution = limits.resolution() / kSubpixelScale;
  const MapLimits superscaled_limits(
      superscaled_resolution, limits.max(),
      CellLimits(limits.cell_limits().num_x_cells * kSubpixelScale,
                 limits.cell_limits().num_y_cells * kSubpixelScale));

  const Eigen::Array2i superscaled_begin =
      superscaled_limits.GetCellIndex(begin);
  const Eigen::Array2i superscaled_end = superscaled_limits.GetCellIndex(end);
  return std::make_pair(superscaled_begin, superscaled_end);
}

struct RangeDataSorter {
  RangeDataSorter(Eigen::Vector3f origin) { origin_ = origin.head<2>(); }
  bool operator()(Eigen::Vector3f lhs, Eigen::Vector3f rhs) {
    const Eigen::Vector2f delta_lhs = (lhs.head<2>() - origin_).normalized();
    const Eigen::Vector2f delta_rhs = (lhs.head<2>() - origin_).normalized();
    if ((delta_lhs[1] < 0.f) != (delta_rhs[1] < 0.f)) {
      return delta_lhs[1] < 0.f;
    } else if (delta_lhs[1] < 0.f) {
      return delta_lhs[0] < delta_rhs[0];
    } else {
      return delta_lhs[0] > delta_rhs[0];
    }
  }

 private:
  Eigen::Vector2f origin_;
};

float ComputeRangeWeightFactor(float range, int exponent) {
  float weight = 0.f;
  if (std::abs(range) > kMinRangeMeters) {
    weight = 1.f / (std::pow(range, exponent));
  }
  return weight;
}
}  // namespace

proto::TSDFRangeDataInserterOptions2D CreateTSDFRangeDataInserterOptions2D(
    common::LuaParameterDictionary* parameter_dictionary) {
  proto::TSDFRangeDataInserterOptions2D options;
  options.set_truncation_distance(
      parameter_dictionary->GetDouble("truncation_distance"));
  options.set_maximum_weight(parameter_dictionary->GetDouble("maximum_weight"));
  options.set_update_free_space(
      parameter_dictionary->GetBool("update_free_space"));
  *options
       .mutable_normal_estimation_options() = CreateNormalEstimationOptions2D(
      parameter_dictionary->GetDictionary("normal_estimation_options").get());
  options.set_project_sdf_distance_to_scan_normal(
      parameter_dictionary->GetBool("project_sdf_distance_to_scan_normal"));
  options.set_update_weight_range_exponent(
      parameter_dictionary->GetInt("update_weight_range_exponent"));
  options.set_update_weight_angle_scan_normal_to_ray_kernel_bandwith(
      parameter_dictionary->GetDouble(
          "update_weight_angle_scan_normal_to_ray_kernel_bandwith"));
  options.set_update_weight_distance_cell_to_hit_kernel_bandwith(
      parameter_dictionary->GetDouble(
          "update_weight_distance_cell_to_hit_kernel_bandwith"));
  return options;
}

TSDFRangeDataInserter2D::TSDFRangeDataInserter2D(
    const proto::TSDFRangeDataInserterOptions2D& options)
    : options_(options) {}

// Casts a ray from origin towards hit for each hit in range data.
// If 'options.update_free_space' is 'true', all cells along the ray
// until 'truncation_distance' behind hit are updated. Otherwise, only the cells
// within 'truncation_distance' around hit are updated.
void TSDFRangeDataInserter2D::Insert(const sensor::RangeData& range_data,
                                     GridInterface* grid) const {
  const float truncation_distance =
      static_cast<float>(options_.truncation_distance());
  TSDF2D* tsdf = static_cast<TSDF2D*>(grid);
  GrowAsNeeded(range_data, truncation_distance, tsdf);

  // Compute normals if needed.
  bool scale_update_weight_angle_scan_normal_to_ray =
      options_.update_weight_angle_scan_normal_to_ray_kernel_bandwith() != 0.f;
  sensor::RangeData sorted_range_data = range_data;
  std::vector<float> normals;
  if (options_.project_sdf_distance_to_scan_normal() ||
      scale_update_weight_angle_scan_normal_to_ray) {
    std::sort(sorted_range_data.returns.begin(),
              sorted_range_data.returns.end(),
              RangeDataSorter(sorted_range_data.origin));
    normals = EstimateNormals(sorted_range_data,
                              options_.normal_estimation_options());
  }

  const Eigen::Vector2f origin = sorted_range_data.origin.head<2>();
  for (size_t hit_index = 0; hit_index < sorted_range_data.returns.size();
       ++hit_index) {
    const Eigen::Vector2f hit = sorted_range_data.returns[hit_index].head<2>();
    const float normal = normals.empty()
                             ? std::numeric_limits<float>::quiet_NaN()
                             : normals[hit_index];
    InsertHit(options_, hit, origin, normal, tsdf);
  }
  tsdf->FinishUpdate();
}

void TSDFRangeDataInserter2D::InsertHit(
    const proto::TSDFRangeDataInserterOptions2D& options,
    const Eigen::Vector2f& hit, const Eigen::Vector2f& origin, float normal,
    TSDF2D* tsdf) const {
  const Eigen::Vector2f ray = hit - origin;
  const float range = ray.norm();
  const float truncation_distance =
      static_cast<float>(options_.truncation_distance());
  if (range < truncation_distance) return;
  const float truncation_ratio = truncation_distance / range;
  const Eigen::Vector2f ray_begin =
      options_.update_free_space() ? origin
                                   : origin + (1.0f - truncation_ratio) * ray;
  const Eigen::Vector2f ray_end = origin + (1.0f + truncation_ratio) * ray;
  std::pair<Eigen::Array2i, Eigen::Array2i> superscaled_ray =
      SuperscaleRay(ray_begin, ray_end, tsdf);
  std::vector<Eigen::Array2i> ray_mask = RayToPixelMask(
      superscaled_ray.first, superscaled_ray.second, kSubpixelScale);

  // Precompute weight factors.
  float weight_factor_angle_ray_normal = 1.f;
  if (options_.update_weight_angle_scan_normal_to_ray_kernel_bandwith() !=
      0.f) {
    const Eigen::Vector2f negative_ray = -ray;
    float angle_ray_normal =
        common::NormalizeAngleDifference(normal - common::atan2(negative_ray));
    weight_factor_angle_ray_normal = GaussianKernel(
        angle_ray_normal,
        options_.update_weight_angle_scan_normal_to_ray_kernel_bandwith());
  }
  float weight_factor_range = 1.f;
  if (options_.update_weight_range_exponent() != 0) {
    weight_factor_range = ComputeRangeWeightFactor(
        range, options_.update_weight_range_exponent());
  }

  // Update Cells.
  for (const Eigen::Array2i& cell_index : ray_mask) {
    if (tsdf->CellIsUpdated(cell_index)) continue;
    Eigen::Vector2f cell_center = tsdf->limits().GetCellCenter(cell_index);
    float distance_cell_to_origin = (cell_center - origin).norm();
    float update_tsd = range - distance_cell_to_origin;
    if (options_.project_sdf_distance_to_scan_normal()) {
      float normal_orientation = normal;
      update_tsd = (cell_center - hit)
                       .dot(Eigen::Vector2f{std::cos(normal_orientation),
                                            std::sin(normal_orientation)});
    }
    update_tsd =
        common::Clamp(update_tsd, -truncation_distance, truncation_distance);
    float update_weight = weight_factor_range * weight_factor_angle_ray_normal;
    if (options_.update_weight_distance_cell_to_hit_kernel_bandwith() != 0.f) {
      update_weight *= GaussianKernel(
          update_tsd,
          options_.update_weight_distance_cell_to_hit_kernel_bandwith());
    }
    UpdateCell(cell_index, update_tsd, update_weight, tsdf);
  }
}

void TSDFRangeDataInserter2D::UpdateCell(const Eigen::Array2i& cell,
                                         float update_sdf, float update_weight,
                                         TSDF2D* tsdf) const {
  if (update_weight == 0.f) return;
  const std::pair<float, float> tsd_and_weight = tsdf->GetTSDAndWeight(cell);
  float updated_weight = tsd_and_weight.second + update_weight;
  float updated_sdf = (tsd_and_weight.first * tsd_and_weight.second +
                       update_sdf * update_weight) /
                      updated_weight;
  updated_weight =
      std::min(updated_weight, static_cast<float>(options_.maximum_weight()));
  tsdf->SetCell(cell, updated_sdf, updated_weight);
}

}  // namespace mapping
}  // namespace cartographer
<commit_msg>Make kSqrtTwoPi const. (#1259)<commit_after>/*
 * Copyright 2018 The Cartographer Authors
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#include "cartographer/mapping/2d/tsdf_range_data_inserter_2d.h"

#include "cartographer/mapping/internal/2d/normal_estimation_2d.h"
#include "cartographer/mapping/internal/2d/ray_to_pixel_mask.h"

namespace cartographer {
namespace mapping {
namespace {

// Factor for subpixel accuracy of start and end point for ray casts.
constexpr int kSubpixelScale = 1000;
// Minimum distance between range observation and origin. Otherwise, range
// observations are discarded.
constexpr float kMinRangeMeters = 1e-6f;
const float kSqrtTwoPi = std::sqrt(2.0 * M_PI);

void GrowAsNeeded(const sensor::RangeData& range_data,
                  const float truncation_distance, TSDF2D* const tsdf) {
  Eigen::AlignedBox2f bounding_box(range_data.origin.head<2>());
  for (const Eigen::Vector3f& hit : range_data.returns) {
    const Eigen::Vector3f direction = (hit - range_data.origin).normalized();
    const Eigen::Vector3f end_position = hit + truncation_distance * direction;
    bounding_box.extend(end_position.head<2>());
  }
  // Padding around bounding box to avoid numerical issues at cell boundaries.
  constexpr float kPadding = 1e-6f;
  tsdf->GrowLimits(bounding_box.min() - kPadding * Eigen::Vector2f::Ones());
  tsdf->GrowLimits(bounding_box.max() + kPadding * Eigen::Vector2f::Ones());
}

float GaussianKernel(const float x, const float sigma) {
  return 1.0 / (kSqrtTwoPi * sigma) * std::exp(-0.5 * x * x / (sigma * sigma));
}

std::pair<Eigen::Array2i, Eigen::Array2i> SuperscaleRay(
    const Eigen::Vector2f& begin, const Eigen::Vector2f& end,
    TSDF2D* const tsdf) {
  const MapLimits& limits = tsdf->limits();
  const double superscaled_resolution = limits.resolution() / kSubpixelScale;
  const MapLimits superscaled_limits(
      superscaled_resolution, limits.max(),
      CellLimits(limits.cell_limits().num_x_cells * kSubpixelScale,
                 limits.cell_limits().num_y_cells * kSubpixelScale));

  const Eigen::Array2i superscaled_begin =
      superscaled_limits.GetCellIndex(begin);
  const Eigen::Array2i superscaled_end = superscaled_limits.GetCellIndex(end);
  return std::make_pair(superscaled_begin, superscaled_end);
}

struct RangeDataSorter {
  RangeDataSorter(Eigen::Vector3f origin) { origin_ = origin.head<2>(); }
  bool operator()(Eigen::Vector3f lhs, Eigen::Vector3f rhs) {
    const Eigen::Vector2f delta_lhs = (lhs.head<2>() - origin_).normalized();
    const Eigen::Vector2f delta_rhs = (lhs.head<2>() - origin_).normalized();
    if ((delta_lhs[1] < 0.f) != (delta_rhs[1] < 0.f)) {
      return delta_lhs[1] < 0.f;
    } else if (delta_lhs[1] < 0.f) {
      return delta_lhs[0] < delta_rhs[0];
    } else {
      return delta_lhs[0] > delta_rhs[0];
    }
  }

 private:
  Eigen::Vector2f origin_;
};

float ComputeRangeWeightFactor(float range, int exponent) {
  float weight = 0.f;
  if (std::abs(range) > kMinRangeMeters) {
    weight = 1.f / (std::pow(range, exponent));
  }
  return weight;
}
}  // namespace

proto::TSDFRangeDataInserterOptions2D CreateTSDFRangeDataInserterOptions2D(
    common::LuaParameterDictionary* parameter_dictionary) {
  proto::TSDFRangeDataInserterOptions2D options;
  options.set_truncation_distance(
      parameter_dictionary->GetDouble("truncation_distance"));
  options.set_maximum_weight(parameter_dictionary->GetDouble("maximum_weight"));
  options.set_update_free_space(
      parameter_dictionary->GetBool("update_free_space"));
  *options
       .mutable_normal_estimation_options() = CreateNormalEstimationOptions2D(
      parameter_dictionary->GetDictionary("normal_estimation_options").get());
  options.set_project_sdf_distance_to_scan_normal(
      parameter_dictionary->GetBool("project_sdf_distance_to_scan_normal"));
  options.set_update_weight_range_exponent(
      parameter_dictionary->GetInt("update_weight_range_exponent"));
  options.set_update_weight_angle_scan_normal_to_ray_kernel_bandwith(
      parameter_dictionary->GetDouble(
          "update_weight_angle_scan_normal_to_ray_kernel_bandwith"));
  options.set_update_weight_distance_cell_to_hit_kernel_bandwith(
      parameter_dictionary->GetDouble(
          "update_weight_distance_cell_to_hit_kernel_bandwith"));
  return options;
}

TSDFRangeDataInserter2D::TSDFRangeDataInserter2D(
    const proto::TSDFRangeDataInserterOptions2D& options)
    : options_(options) {}

// Casts a ray from origin towards hit for each hit in range data.
// If 'options.update_free_space' is 'true', all cells along the ray
// until 'truncation_distance' behind hit are updated. Otherwise, only the cells
// within 'truncation_distance' around hit are updated.
void TSDFRangeDataInserter2D::Insert(const sensor::RangeData& range_data,
                                     GridInterface* grid) const {
  const float truncation_distance =
      static_cast<float>(options_.truncation_distance());
  TSDF2D* tsdf = static_cast<TSDF2D*>(grid);
  GrowAsNeeded(range_data, truncation_distance, tsdf);

  // Compute normals if needed.
  bool scale_update_weight_angle_scan_normal_to_ray =
      options_.update_weight_angle_scan_normal_to_ray_kernel_bandwith() != 0.f;
  sensor::RangeData sorted_range_data = range_data;
  std::vector<float> normals;
  if (options_.project_sdf_distance_to_scan_normal() ||
      scale_update_weight_angle_scan_normal_to_ray) {
    std::sort(sorted_range_data.returns.begin(),
              sorted_range_data.returns.end(),
              RangeDataSorter(sorted_range_data.origin));
    normals = EstimateNormals(sorted_range_data,
                              options_.normal_estimation_options());
  }

  const Eigen::Vector2f origin = sorted_range_data.origin.head<2>();
  for (size_t hit_index = 0; hit_index < sorted_range_data.returns.size();
       ++hit_index) {
    const Eigen::Vector2f hit = sorted_range_data.returns[hit_index].head<2>();
    const float normal = normals.empty()
                             ? std::numeric_limits<float>::quiet_NaN()
                             : normals[hit_index];
    InsertHit(options_, hit, origin, normal, tsdf);
  }
  tsdf->FinishUpdate();
}

void TSDFRangeDataInserter2D::InsertHit(
    const proto::TSDFRangeDataInserterOptions2D& options,
    const Eigen::Vector2f& hit, const Eigen::Vector2f& origin, float normal,
    TSDF2D* tsdf) const {
  const Eigen::Vector2f ray = hit - origin;
  const float range = ray.norm();
  const float truncation_distance =
      static_cast<float>(options_.truncation_distance());
  if (range < truncation_distance) return;
  const float truncation_ratio = truncation_distance / range;
  const Eigen::Vector2f ray_begin =
      options_.update_free_space() ? origin
                                   : origin + (1.0f - truncation_ratio) * ray;
  const Eigen::Vector2f ray_end = origin + (1.0f + truncation_ratio) * ray;
  std::pair<Eigen::Array2i, Eigen::Array2i> superscaled_ray =
      SuperscaleRay(ray_begin, ray_end, tsdf);
  std::vector<Eigen::Array2i> ray_mask = RayToPixelMask(
      superscaled_ray.first, superscaled_ray.second, kSubpixelScale);

  // Precompute weight factors.
  float weight_factor_angle_ray_normal = 1.f;
  if (options_.update_weight_angle_scan_normal_to_ray_kernel_bandwith() !=
      0.f) {
    const Eigen::Vector2f negative_ray = -ray;
    float angle_ray_normal =
        common::NormalizeAngleDifference(normal - common::atan2(negative_ray));
    weight_factor_angle_ray_normal = GaussianKernel(
        angle_ray_normal,
        options_.update_weight_angle_scan_normal_to_ray_kernel_bandwith());
  }
  float weight_factor_range = 1.f;
  if (options_.update_weight_range_exponent() != 0) {
    weight_factor_range = ComputeRangeWeightFactor(
        range, options_.update_weight_range_exponent());
  }

  // Update Cells.
  for (const Eigen::Array2i& cell_index : ray_mask) {
    if (tsdf->CellIsUpdated(cell_index)) continue;
    Eigen::Vector2f cell_center = tsdf->limits().GetCellCenter(cell_index);
    float distance_cell_to_origin = (cell_center - origin).norm();
    float update_tsd = range - distance_cell_to_origin;
    if (options_.project_sdf_distance_to_scan_normal()) {
      float normal_orientation = normal;
      update_tsd = (cell_center - hit)
                       .dot(Eigen::Vector2f{std::cos(normal_orientation),
                                            std::sin(normal_orientation)});
    }
    update_tsd =
        common::Clamp(update_tsd, -truncation_distance, truncation_distance);
    float update_weight = weight_factor_range * weight_factor_angle_ray_normal;
    if (options_.update_weight_distance_cell_to_hit_kernel_bandwith() != 0.f) {
      update_weight *= GaussianKernel(
          update_tsd,
          options_.update_weight_distance_cell_to_hit_kernel_bandwith());
    }
    UpdateCell(cell_index, update_tsd, update_weight, tsdf);
  }
}

void TSDFRangeDataInserter2D::UpdateCell(const Eigen::Array2i& cell,
                                         float update_sdf, float update_weight,
                                         TSDF2D* tsdf) const {
  if (update_weight == 0.f) return;
  const std::pair<float, float> tsd_and_weight = tsdf->GetTSDAndWeight(cell);
  float updated_weight = tsd_and_weight.second + update_weight;
  float updated_sdf = (tsd_and_weight.first * tsd_and_weight.second +
                       update_sdf * update_weight) /
                      updated_weight;
  updated_weight =
      std::min(updated_weight, static_cast<float>(options_.maximum_weight()));
  tsdf->SetCell(cell, updated_sdf, updated_weight);
}

}  // namespace mapping
}  // namespace cartographer
<|endoftext|>


<commit_before>#define LOG_MODULE PacketLogModuleDnsLayer

#include "DnsResource.h"
#include "Logger.h"
#include <sstream>
#include <string.h>
#include "EndianPortable.h"

namespace pcpp
{

IDnsResource::IDnsResource(DnsLayer* dnsLayer, size_t offsetInLayer)
	: m_DnsLayer(dnsLayer), m_OffsetInLayer(offsetInLayer), m_NextResource(NULL)
{
	char decodedName[256];
	m_NameLength = decodeName((const char*)getRawData(), decodedName);
	m_DecodedName = decodedName;
}

IDnsResource::IDnsResource(uint8_t* emptyRawData)
	: m_DnsLayer(NULL), m_OffsetInLayer(0), m_NextResource(NULL), m_DecodedName(""), m_NameLength(0), m_ExternalRawData(emptyRawData)
{
}

uint8_t* IDnsResource::getRawData() const
{
	if (m_DnsLayer == NULL)
		return m_ExternalRawData;

	return m_DnsLayer->m_Data + m_OffsetInLayer;
}

size_t IDnsResource::decodeName(const char* encodedName, char* result, int iteration)
{
	size_t encodedNameLength = 0;
	char* resultPtr = result;
	resultPtr[0] = 0;

	size_t curOffsetInLayer = (uint8_t*)encodedName - m_DnsLayer->m_Data;
	if (curOffsetInLayer + 1 > m_DnsLayer->m_DataLen)
		return encodedNameLength;

	if (iteration > 20)
		return encodedNameLength;

	uint8_t wordLength = encodedName[0];

	// A string to parse
	while (wordLength != 0)
	{
		// A pointer to another place in the packet
		if ((wordLength & 0xc0) == 0xc0)
		{
			if (curOffsetInLayer + 2 > m_DnsLayer->m_DataLen)
				return encodedNameLength;

			uint16_t offsetInLayer = (wordLength & 0x3f)*256 + (0xFF & encodedName[1]);
			if (offsetInLayer < sizeof(dnshdr) || offsetInLayer >= m_DnsLayer->m_DataLen)
			{
				LOG_ERROR("DNS parsing error: name pointer is illegal");
				return 0;
			}

			char tempResult[256];
			int i = 0;
			decodeName((const char*)(m_DnsLayer->m_Data + offsetInLayer), tempResult, iteration+1);
			while (tempResult[i] != 0)
			{
				resultPtr[0] = tempResult[i++];
				resultPtr++;
			}

			resultPtr[0] = 0;

			// in this case the length of the pointer is: 1B for 0xc0 + 1B for the offset itself
			return encodedNameLength + sizeof(uint16_t);
		}
		else
		{
			// return if next word would be outside of the DNS layer or overflow the buffer behind resultPtr
			if (curOffsetInLayer + wordLength + 1 > m_DnsLayer->m_DataLen || encodedNameLength + wordLength > 255)
			{
				// add the last '\0' to the decoded string
				if (encodedNameLength == 256)
					resultPtr--;
				else
					encodedNameLength++;

				resultPtr[0] = 0;
				return encodedNameLength;
			}

			memcpy(resultPtr, encodedName+1, wordLength);
			resultPtr += wordLength;
			resultPtr[0] = '.';
			resultPtr++;
			encodedName += wordLength + 1;
			encodedNameLength += wordLength + 1;

			curOffsetInLayer = (uint8_t*)encodedName - m_DnsLayer->m_Data;
			if (curOffsetInLayer + 1 > m_DnsLayer->m_DataLen)
			{
				// add the last '\0' to the decoded string
				if (encodedNameLength == 256)
					resultPtr--;
				else
					encodedNameLength++;

				resultPtr[0] = 0;
				return encodedNameLength;
			}

			wordLength = encodedName[0];
		}
	}

	// remove the last "."
	if (resultPtr > result)
	{
		result[resultPtr - result - 1] = 0;
	}

	// add the last '\0' to encodedNameLength
	resultPtr[0] = 0;
	encodedNameLength++;

	return encodedNameLength;
}

void IDnsResource::encodeName(const std::string& decodedName, char* result, size_t& resultLen)
{
	resultLen = 0;
	std::stringstream strstream(decodedName);
    std::string word;
    while (getline(strstream, word, '.'))
    {
    	// pointer to a different hostname in the packet
    	if (word[0] == '#')
    	{
    		// convert the number from string to int
    		std::stringstream stream(word.substr(1));
    		int pointerInPacket = 0;
    		stream >> pointerInPacket;

    		// verify it's indeed a number and that is in the range of [0-255]
    		if (stream.fail() || pointerInPacket < 0 || pointerInPacket > 0xff)
    		{
    			LOG_ERROR("Error encoding the string '%s'", decodedName.c_str());
    			return;
    		}

    		// set the pointer to the encoded string result
    		result[0] = (uint8_t)0xc0;
    		result[1] = (uint8_t)pointerInPacket;
    		result += 2;
    		resultLen += 2;
    		return; // pointer always comes last
    	}

    	result[0] = word.length();
    	result++;
    	memcpy(result, word.c_str(), word.length());
    	result += word.length();
    	resultLen += word.length() + 1;
    }

    result[0] = 0;
    resultLen++;
}

DnsType IDnsResource::getDnsType() const
{
	uint16_t dnsType = *(uint16_t*)(getRawData() + m_NameLength);
	return (DnsType)be16toh(dnsType);
}

void IDnsResource::setDnsType(DnsType newType)
{
	uint16_t newTypeAsInt = htobe16((uint16_t)newType);
	memcpy(getRawData() + m_NameLength, &newTypeAsInt, sizeof(uint16_t));
}

DnsClass IDnsResource::getDnsClass() const
{
	uint16_t dnsClass = *(uint16_t*)(getRawData() + m_NameLength + sizeof(uint16_t));
	return (DnsClass)be16toh(dnsClass);
}

void IDnsResource::setDnsClass(DnsClass newClass)
{
	uint16_t newClassAsInt = htobe16((uint16_t)newClass);
	memcpy(getRawData() + m_NameLength + sizeof(uint16_t), &newClassAsInt, sizeof(uint16_t));
}

bool IDnsResource::setName(const std::string& newName)
{
	char encodedName[256];
	size_t encodedNameLen = 0;
	encodeName(newName, encodedName, encodedNameLen);
	if (m_DnsLayer != NULL)
	{
		if (encodedNameLen > m_NameLength)
		{
			if (!m_DnsLayer->extendLayer(m_OffsetInLayer, encodedNameLen-m_NameLength, this))
			{
				LOG_ERROR("Couldn't set name for DNS query, unable to extend layer");
				return false;
			}
		}
		else if (encodedNameLen < m_NameLength)
		{
			if (!m_DnsLayer->shortenLayer(m_OffsetInLayer, m_NameLength-encodedNameLen, this))
			{
				LOG_ERROR("Couldn't set name for DNS query, unable to shorten layer");
				return false;
			}
		}
	}
	else
	{
		size_t size = getSize();
		char* tempData = new char[size];
		memcpy(tempData, m_ExternalRawData, size);
		memcpy(m_ExternalRawData + encodedNameLen, tempData, size);
		delete[] tempData;
	}

	memcpy(getRawData(), encodedName, encodedNameLen);
	m_NameLength = encodedNameLen;
	m_DecodedName = newName;

	return true;
}

void IDnsResource::setDnsLayer(DnsLayer* dnsLayer, size_t offsetInLayer)
{
	memcpy(dnsLayer->m_Data + offsetInLayer, m_ExternalRawData, getSize());
	m_DnsLayer = dnsLayer;
	m_OffsetInLayer = offsetInLayer;
	m_ExternalRawData = NULL;
}

uint32_t DnsResource::getTTL() const
{
	uint32_t ttl = *(uint32_t*)(getRawData() + m_NameLength + 2*sizeof(uint16_t));
	return be32toh(ttl);
}

void DnsResource::setTTL(uint32_t newTTL)
{
	newTTL = htobe32(newTTL);
	memcpy(getRawData() + m_NameLength + 2*sizeof(uint16_t), &newTTL, sizeof(uint32_t));
}

size_t DnsResource::getDataLength() const
{
	uint16_t dataLength = *(uint16_t*)(getRawData() + m_NameLength + 2*sizeof(uint16_t) + sizeof(uint32_t));
	return be16toh(dataLength);
}

DnsResourceDataPtr DnsResource::getData() const
{
	uint8_t* resourceRawData = getRawData() + m_NameLength + 3*sizeof(uint16_t) + sizeof(uint32_t);
	size_t dataLength = getDataLength();

	switch (getDnsType())
	{
	case DNS_TYPE_A:
	{
		return DnsResourceDataPtr(new IPv4DnsResourceData(resourceRawData, dataLength));
	}

	case DNS_TYPE_AAAA:
	{
		return DnsResourceDataPtr(new IPv6DnsResourceData(resourceRawData, dataLength));
	}

	case DNS_TYPE_NS:
	case DNS_TYPE_CNAME:
	case DNS_TYPE_DNAM:
	case DNS_TYPE_PTR:
	{
		return DnsResourceDataPtr(new StringDnsResourceData(resourceRawData, dataLength, const_cast<IDnsResource*>(static_cast<const IDnsResource*>(this))));
	}

	case DNS_TYPE_MX:
	{
		return DnsResourceDataPtr(new MxDnsResourceData(resourceRawData, dataLength, const_cast<IDnsResource*>(static_cast<const IDnsResource*>(this))));
	}

	default:
	{
		return DnsResourceDataPtr(new GenericDnsResourceData(resourceRawData, dataLength));
	}

	}
}

size_t DnsResource::getDataOffset() const
{
	return (size_t)(m_OffsetInLayer + m_NameLength + 3*sizeof(uint16_t) + sizeof(uint32_t));
}

bool DnsResource::setData(IDnsResourceData* data)
{
	// convert data to byte array according to the DNS type
	size_t dataLength = 0;
	uint8_t dataAsByteArr[256];

	if (data == NULL)
	{
		LOG_ERROR("Given data is NULL");
		return false;
	}

	switch (getDnsType())
	{
	case DNS_TYPE_A:
	{
		if (!data->isTypeOf<IPv4DnsResourceData>())
		{
			LOG_ERROR("DNS record is of type A but given data isn't of type IPv4DnsResourceData");
			return false;
		}
		break;
	}

	case DNS_TYPE_AAAA:
	{
		if (!data->isTypeOf<IPv6DnsResourceData>())
		{
			LOG_ERROR("DNS record is of type AAAA but given data isn't of type IPv6DnsResourceData");
			return false;
		}
		break;
	}

	case DNS_TYPE_NS:
	case DNS_TYPE_CNAME:
	case DNS_TYPE_DNAM:
	case DNS_TYPE_PTR:
	{
		if (!data->isTypeOf<StringDnsResourceData>())
		{
			LOG_ERROR("DNS record is of type NS, CNAME, DNAM or PTR but given data isn't of type StringDnsResourceData");
			return false;
		}
		break;
	}

	case DNS_TYPE_MX:
	{
		if (!data->isTypeOf<MxDnsResourceData>())
		{
			LOG_ERROR("DNS record is of type MX but given data isn't of type MxDnsResourceData");
			return false;
		}
		break;
	}

	default:
	{
		// do nothing
	}

	}

	// convert the IDnsResourceData to byte array
	if (!data->toByteArr(dataAsByteArr, dataLength, this))
	{
		LOG_ERROR("Cannot convert DNS resource data to byte array, data is probably invalid");
		return false;
	}

	size_t dataLengthOffset = m_NameLength + (2*sizeof(uint16_t)) + sizeof(uint32_t);
	size_t dataOffset = dataLengthOffset + sizeof(uint16_t);

	if (m_DnsLayer != NULL)
	{
		size_t curLength = getDataLength();
		if (dataLength > curLength)
		{
			if (!m_DnsLayer->extendLayer(m_OffsetInLayer + dataOffset, dataLength-curLength, this))
			{
				LOG_ERROR("Couldn't set data for DNS query, unable to extend layer");
				return false;
			}
		}
		else if (dataLength < curLength)
		{
			if (!m_DnsLayer->shortenLayer(m_OffsetInLayer + dataOffset, curLength-dataLength, this))
			{
				LOG_ERROR("Couldn't set data for DNS query, unable to shorten layer");
				return false;
			}
		}
	}

	// write data to resource
	memcpy(getRawData() + dataOffset, dataAsByteArr, dataLength);
	//update data length in resource
	dataLength = htobe16((uint16_t)dataLength);
	memcpy(getRawData() + dataLengthOffset, &dataLength, sizeof(uint16_t));

	return true;
}

uint16_t DnsResource::getCustomDnsClass() const
{
	uint16_t value = *(uint16_t*)(getRawData() + m_NameLength + sizeof(uint16_t));
	return be16toh(value);
}

void DnsResource::setCustomDnsClass(uint16_t customValue)
{
	memcpy(getRawData() + m_NameLength + sizeof(uint16_t), &customValue, sizeof(uint16_t));
}

}
<commit_msg>Bugfix #424<commit_after>#define LOG_MODULE PacketLogModuleDnsLayer

#include "DnsResource.h"
#include "Logger.h"
#include <sstream>
#include <string.h>
#include "EndianPortable.h"

namespace pcpp
{

IDnsResource::IDnsResource(DnsLayer* dnsLayer, size_t offsetInLayer)
	: m_DnsLayer(dnsLayer), m_OffsetInLayer(offsetInLayer), m_NextResource(NULL)
{
	char decodedName[256];
	m_NameLength = decodeName((const char*)getRawData(), decodedName);
	m_DecodedName = decodedName;
}

IDnsResource::IDnsResource(uint8_t* emptyRawData)
	: m_DnsLayer(NULL), m_OffsetInLayer(0), m_NextResource(NULL), m_DecodedName(""), m_NameLength(0), m_ExternalRawData(emptyRawData)
{
}

uint8_t* IDnsResource::getRawData() const
{
	if (m_DnsLayer == NULL)
		return m_ExternalRawData;

	return m_DnsLayer->m_Data + m_OffsetInLayer;
}

size_t IDnsResource::decodeName(const char* encodedName, char* result, int iteration)
{
	size_t encodedNameLength = 0;
	size_t decodedNameLength = 0;
	char* resultPtr = result;
	resultPtr[0] = 0;

	size_t curOffsetInLayer = (uint8_t*)encodedName - m_DnsLayer->m_Data;
	if (curOffsetInLayer + 1 > m_DnsLayer->m_DataLen)
		return encodedNameLength;

	if (iteration > 20)
		return encodedNameLength;

	uint8_t wordLength = encodedName[0];

	// A string to parse
	while (wordLength != 0)
	{
		// A pointer to another place in the packet
		if ((wordLength & 0xc0) == 0xc0)
		{
			if (curOffsetInLayer + 2 > m_DnsLayer->m_DataLen)
				return encodedNameLength;

			uint16_t offsetInLayer = (wordLength & 0x3f)*256 + (0xFF & encodedName[1]);
			if (offsetInLayer < sizeof(dnshdr) || offsetInLayer >= m_DnsLayer->m_DataLen)
			{
				LOG_ERROR("DNS parsing error: name pointer is illegal");
				return 0;
			}

			char tempResult[256];
			int i = 0;
			decodeName((const char*)(m_DnsLayer->m_Data + offsetInLayer), tempResult, iteration+1);
			while (tempResult[i] != 0 && decodedNameLength < 255)
			{
				resultPtr[0] = tempResult[i++];
				resultPtr++;
				decodedNameLength++;
			}

			resultPtr[0] = 0;

			// in this case the length of the pointer is: 1 byte for 0xc0 + 1 byte for the offset itself
			return encodedNameLength + sizeof(uint16_t);
		}
		else
		{
			// return if next word would be outside of the DNS layer or overflow the buffer behind resultPtr
			if (curOffsetInLayer + wordLength + 1 > m_DnsLayer->m_DataLen || encodedNameLength + wordLength > 255)
			{
				// add the last '\0' to the decoded string
				if (encodedNameLength == 256)
				{
					resultPtr--;
					decodedNameLength--;
				}
				else
				{
					encodedNameLength++;
				}

				resultPtr[0] = 0;
				return encodedNameLength;
			}

			memcpy(resultPtr, encodedName+1, wordLength);
			resultPtr += wordLength;
			resultPtr[0] = '.';
			resultPtr++;
			decodedNameLength += wordLength + 1;
			encodedName += wordLength + 1;
			encodedNameLength += wordLength + 1;

			curOffsetInLayer = (uint8_t*)encodedName - m_DnsLayer->m_Data;
			if (curOffsetInLayer + 1 > m_DnsLayer->m_DataLen)
			{
				// add the last '\0' to the decoded string
				if (encodedNameLength == 256)
				{
					decodedNameLength--;
					resultPtr--;
				}
				else
				{
					encodedNameLength++;
				}

				resultPtr[0] = 0;
				return encodedNameLength;
			}

			wordLength = encodedName[0];
		}
	}

	// remove the last "."
	if (resultPtr > result)
	{
		result[resultPtr - result - 1] = 0;
	}

	// add the last '\0' to encodedNameLength
	resultPtr[0] = 0;
	encodedNameLength++;

	return encodedNameLength;
}

void IDnsResource::encodeName(const std::string& decodedName, char* result, size_t& resultLen)
{
	resultLen = 0;
	std::stringstream strstream(decodedName);
    std::string word;
    while (getline(strstream, word, '.'))
    {
    	// pointer to a different hostname in the packet
    	if (word[0] == '#')
    	{
    		// convert the number from string to int
    		std::stringstream stream(word.substr(1));
    		int pointerInPacket = 0;
    		stream >> pointerInPacket;

    		// verify it's indeed a number and that is in the range of [0-255]
    		if (stream.fail() || pointerInPacket < 0 || pointerInPacket > 0xff)
    		{
    			LOG_ERROR("Error encoding the string '%s'", decodedName.c_str());
    			return;
    		}

    		// set the pointer to the encoded string result
    		result[0] = (uint8_t)0xc0;
    		result[1] = (uint8_t)pointerInPacket;
    		result += 2;
    		resultLen += 2;
    		return; // pointer always comes last
    	}

    	result[0] = word.length();
    	result++;
    	memcpy(result, word.c_str(), word.length());
    	result += word.length();
    	resultLen += word.length() + 1;
    }

    result[0] = 0;
    resultLen++;
}

DnsType IDnsResource::getDnsType() const
{
	uint16_t dnsType = *(uint16_t*)(getRawData() + m_NameLength);
	return (DnsType)be16toh(dnsType);
}

void IDnsResource::setDnsType(DnsType newType)
{
	uint16_t newTypeAsInt = htobe16((uint16_t)newType);
	memcpy(getRawData() + m_NameLength, &newTypeAsInt, sizeof(uint16_t));
}

DnsClass IDnsResource::getDnsClass() const
{
	uint16_t dnsClass = *(uint16_t*)(getRawData() + m_NameLength + sizeof(uint16_t));
	return (DnsClass)be16toh(dnsClass);
}

void IDnsResource::setDnsClass(DnsClass newClass)
{
	uint16_t newClassAsInt = htobe16((uint16_t)newClass);
	memcpy(getRawData() + m_NameLength + sizeof(uint16_t), &newClassAsInt, sizeof(uint16_t));
}

bool IDnsResource::setName(const std::string& newName)
{
	char encodedName[256];
	size_t encodedNameLen = 0;
	encodeName(newName, encodedName, encodedNameLen);
	if (m_DnsLayer != NULL)
	{
		if (encodedNameLen > m_NameLength)
		{
			if (!m_DnsLayer->extendLayer(m_OffsetInLayer, encodedNameLen-m_NameLength, this))
			{
				LOG_ERROR("Couldn't set name for DNS query, unable to extend layer");
				return false;
			}
		}
		else if (encodedNameLen < m_NameLength)
		{
			if (!m_DnsLayer->shortenLayer(m_OffsetInLayer, m_NameLength-encodedNameLen, this))
			{
				LOG_ERROR("Couldn't set name for DNS query, unable to shorten layer");
				return false;
			}
		}
	}
	else
	{
		size_t size = getSize();
		char* tempData = new char[size];
		memcpy(tempData, m_ExternalRawData, size);
		memcpy(m_ExternalRawData + encodedNameLen, tempData, size);
		delete[] tempData;
	}

	memcpy(getRawData(), encodedName, encodedNameLen);
	m_NameLength = encodedNameLen;
	m_DecodedName = newName;

	return true;
}

void IDnsResource::setDnsLayer(DnsLayer* dnsLayer, size_t offsetInLayer)
{
	memcpy(dnsLayer->m_Data + offsetInLayer, m_ExternalRawData, getSize());
	m_DnsLayer = dnsLayer;
	m_OffsetInLayer = offsetInLayer;
	m_ExternalRawData = NULL;
}

uint32_t DnsResource::getTTL() const
{
	uint32_t ttl = *(uint32_t*)(getRawData() + m_NameLength + 2*sizeof(uint16_t));
	return be32toh(ttl);
}

void DnsResource::setTTL(uint32_t newTTL)
{
	newTTL = htobe32(newTTL);
	memcpy(getRawData() + m_NameLength + 2*sizeof(uint16_t), &newTTL, sizeof(uint32_t));
}

size_t DnsResource::getDataLength() const
{
	uint16_t dataLength = *(uint16_t*)(getRawData() + m_NameLength + 2*sizeof(uint16_t) + sizeof(uint32_t));
	return be16toh(dataLength);
}

DnsResourceDataPtr DnsResource::getData() const
{
	uint8_t* resourceRawData = getRawData() + m_NameLength + 3*sizeof(uint16_t) + sizeof(uint32_t);
	size_t dataLength = getDataLength();

	switch (getDnsType())
	{
	case DNS_TYPE_A:
	{
		return DnsResourceDataPtr(new IPv4DnsResourceData(resourceRawData, dataLength));
	}

	case DNS_TYPE_AAAA:
	{
		return DnsResourceDataPtr(new IPv6DnsResourceData(resourceRawData, dataLength));
	}

	case DNS_TYPE_NS:
	case DNS_TYPE_CNAME:
	case DNS_TYPE_DNAM:
	case DNS_TYPE_PTR:
	{
		return DnsResourceDataPtr(new StringDnsResourceData(resourceRawData, dataLength, const_cast<IDnsResource*>(static_cast<const IDnsResource*>(this))));
	}

	case DNS_TYPE_MX:
	{
		return DnsResourceDataPtr(new MxDnsResourceData(resourceRawData, dataLength, const_cast<IDnsResource*>(static_cast<const IDnsResource*>(this))));
	}

	default:
	{
		return DnsResourceDataPtr(new GenericDnsResourceData(resourceRawData, dataLength));
	}

	}
}

size_t DnsResource::getDataOffset() const
{
	return (size_t)(m_OffsetInLayer + m_NameLength + 3*sizeof(uint16_t) + sizeof(uint32_t));
}

bool DnsResource::setData(IDnsResourceData* data)
{
	// convert data to byte array according to the DNS type
	size_t dataLength = 0;
	uint8_t dataAsByteArr[256];

	if (data == NULL)
	{
		LOG_ERROR("Given data is NULL");
		return false;
	}

	switch (getDnsType())
	{
	case DNS_TYPE_A:
	{
		if (!data->isTypeOf<IPv4DnsResourceData>())
		{
			LOG_ERROR("DNS record is of type A but given data isn't of type IPv4DnsResourceData");
			return false;
		}
		break;
	}

	case DNS_TYPE_AAAA:
	{
		if (!data->isTypeOf<IPv6DnsResourceData>())
		{
			LOG_ERROR("DNS record is of type AAAA but given data isn't of type IPv6DnsResourceData");
			return false;
		}
		break;
	}

	case DNS_TYPE_NS:
	case DNS_TYPE_CNAME:
	case DNS_TYPE_DNAM:
	case DNS_TYPE_PTR:
	{
		if (!data->isTypeOf<StringDnsResourceData>())
		{
			LOG_ERROR("DNS record is of type NS, CNAME, DNAM or PTR but given data isn't of type StringDnsResourceData");
			return false;
		}
		break;
	}

	case DNS_TYPE_MX:
	{
		if (!data->isTypeOf<MxDnsResourceData>())
		{
			LOG_ERROR("DNS record is of type MX but given data isn't of type MxDnsResourceData");
			return false;
		}
		break;
	}

	default:
	{
		// do nothing
	}

	}

	// convert the IDnsResourceData to byte array
	if (!data->toByteArr(dataAsByteArr, dataLength, this))
	{
		LOG_ERROR("Cannot convert DNS resource data to byte array, data is probably invalid");
		return false;
	}

	size_t dataLengthOffset = m_NameLength + (2*sizeof(uint16_t)) + sizeof(uint32_t);
	size_t dataOffset = dataLengthOffset + sizeof(uint16_t);

	if (m_DnsLayer != NULL)
	{
		size_t curLength = getDataLength();
		if (dataLength > curLength)
		{
			if (!m_DnsLayer->extendLayer(m_OffsetInLayer + dataOffset, dataLength-curLength, this))
			{
				LOG_ERROR("Couldn't set data for DNS query, unable to extend layer");
				return false;
			}
		}
		else if (dataLength < curLength)
		{
			if (!m_DnsLayer->shortenLayer(m_OffsetInLayer + dataOffset, curLength-dataLength, this))
			{
				LOG_ERROR("Couldn't set data for DNS query, unable to shorten layer");
				return false;
			}
		}
	}

	// write data to resource
	memcpy(getRawData() + dataOffset, dataAsByteArr, dataLength);
	//update data length in resource
	dataLength = htobe16((uint16_t)dataLength);
	memcpy(getRawData() + dataLengthOffset, &dataLength, sizeof(uint16_t));

	return true;
}

uint16_t DnsResource::getCustomDnsClass() const
{
	uint16_t value = *(uint16_t*)(getRawData() + m_NameLength + sizeof(uint16_t));
	return be16toh(value);
}

void DnsResource::setCustomDnsClass(uint16_t customValue)
{
	memcpy(getRawData() + m_NameLength + sizeof(uint16_t), &customValue, sizeof(uint16_t));
}

}
<|endoftext|>