
#include "gre.h"
#include "protocol_gre.h"
#include "common/attributes.h"
#include "config.h"
#include "crypto/psk.h"
#include "log-private.h"
#include "logging.h"
#include "rist-private.h"
#include "endian-shim.h"
#include "udp-private.h"
#include "transport-private.h"
#include "eap.h"
#include "peer.h"

#include <errno.h>
#include <stddef.h>
#include <stdint.h>
#include <stdbool.h>
#include <assert.h>
#include <stdlib.h>
#include <string.h>

ssize_t _librist_proto_gre_send_data(struct rist_peer *p, uint8_t payload_type, uint16_t proto, uint8_t *payload, size_t payload_len, uint16_t src_port, uint16_t dst_port, uint8_t gre_version) {
	bool encrypt = (p->key_tx.key_size > 0) && proto != RIST_GRE_PROTOCOL_TYPE_EAPOL;

	/* Our encryption and compression operations directly modify the payload buffer we receive as a pointer
	   so we create a local pointer that points to the payload pointer, if we would either encrypt or compress we instead
	   malloc and mempcy, to ensure our source stays clean. We only do this with RAW data as these buffers are the only
	   assumed to be reused by retransmits */
    bool modifying_payload = encrypt && (payload_type == RIST_PAYLOAD_TYPE_DATA_RAW || payload_type == RIST_PAYLOAD_TYPE_DATA_RAW_RTP_EXT);//Need to make a copy of our data if we'd resend it in the future, otherwise we can safely overwrite the buffer

	uint8_t *payload_wr = payload;
	uint8_t hdr_buf[MAX_GRE_SIZE];
	struct rist_gre_hdr* hdr = (struct rist_gre_hdr *)hdr_buf;
	memset(hdr, 0, sizeof(*hdr));
	size_t hdr_len = sizeof(*hdr);


	struct rist_peer *key_peer = p;
	if (key_peer->parent != NULL && key_peer->parent->multicast_sender)
		key_peer = key_peer->parent;

	uint32_t seq;
    assert(payload != NULL);

	const size_t nonce_offset = sizeof(*hdr);
	hdr->flags2 = (gre_version &0x7) << 3;

	if (encrypt) {
		hdr_len += 4;
		if (gre_version && p->key_tx.key_size == 256)
		{
			hdr->flags2 |= (1 & 1UL) << 6;
		}
	}
	//VSF TR-06-02:_2022-08-11 wants us to send a sequence if we're capable of sending non-RIST traffic
	//For now default to always considering us capable to do that.
	if (true || encrypt) {
		SET_BIT(hdr->flags1, 4); // set seq bit
		seq = key_peer->seq++;
		//Write sequence number
		hdr_buf[hdr_len] = seq >> 24;
		hdr_buf[hdr_len+1] = seq >> 16;
		hdr_buf[hdr_len+2] = seq >> 8;
		hdr_buf[hdr_len+3] = seq;
		hdr_len += 4;
	}

	size_t hdr_payload_offset = hdr_len;

	//If we can use the new VSF ethertype (rist GRE version >= 2), we should, as the other way is considered deprecated
	if ((gre_version >= 2 && (proto == RIST_GRE_PROTOCOL_TYPE_REDUCED || proto == RIST_GRE_PROTOCOL_TYPE_KEEPALIVE || proto == RIST_VSF_PROTOCOL_SUBTYPE_BUFFER_NEGOTIATION))) {
		struct rist_vsf_proto *vsf = (struct rist_vsf_proto *)&hdr_buf[hdr_len];
		hdr_len += sizeof(*vsf);
		vsf->type = RIST_VSF_PROTOCOL_TYPE_RIST;//Network byte order! these are 0 values, so safe to do without htobe16
		if (proto == RIST_GRE_PROTOCOL_TYPE_REDUCED)
			vsf->subtype = RIST_VSF_PROTOCOL_SUBTYPE_REDUCED;
		else if (proto == RIST_VSF_PROTOCOL_SUBTYPE_BUFFER_NEGOTIATION)
			vsf->subtype = htobe16(RIST_VSF_PROTOCOL_SUBTYPE_BUFFER_NEGOTIATION);
		else
			vsf->subtype = htobe16(RIST_VSF_PROTOCOL_SUBTYPE_KEEPALIVE);

		hdr->prot_type = htobe16(RIST_GRE_PROTOCOL_TYPE_VSF);
	} else {
		hdr->prot_type = htobe16(proto);
	}

	if (proto == RIST_GRE_PROTOCOL_TYPE_REDUCED) {
		struct rist_reduced *red = (struct rist_reduced *)(&hdr_buf[hdr_len]);
		hdr_len += sizeof(*red);
		red->dst_port = htobe16(dst_port);
		red->src_port = htobe16(src_port);
	}

	if (encrypt) {
		if (modifying_payload || (hdr_payload_offset != hdr_len && !HAVE_MBEDTLS)) {
			payload_wr = malloc(payload_len + 8);//Let's get rid of this malloc in the hotpath here
			modifying_payload = true;
			assert(payload_wr);
		}
		pthread_mutex_lock(&key_peer->peer_lock);
		struct rist_key *key = &key_peer->key_tx;

#if HAVE_SRP_SUPPORT
        if (librist_peer_should_rollover_passphrase(key_peer)) {
			key_peer->key_tx_odd_active = !key_peer->key_tx_odd_active;
			rist_log_priv(get_cctx(p), RIST_LOG_INFO, "Rolling over to %s passphrase\n", key_peer->key_tx_odd_active? "odd" : "even");
		}
#endif

		if (key_peer->key_tx_odd_active)
			key = &p->key_tx_odd;

		//Data that should be encrypted is part of the header
		if (hdr_payload_offset != hdr_len) {
#if HAVE_MBEDTLS
			//MbedTLS allows us to continue encryption, so we can encrypt hdr & payload in 2 ops
			_librist_crypto_psk_encrypt(key, htobe32(seq), gre_version, &hdr_buf[hdr_payload_offset], &hdr_buf[hdr_payload_offset], hdr_len - hdr_payload_offset);
			_librist_crypto_psk_encrypt_continue(key, payload, payload_wr, payload_len);
#else
			memcpy(payload_wr, &hdr_buf[hdr_payload_offset], hdr_len - hdr_payload_offset);
			memcpy(&payload_wr[hdr_len - hdr_payload_offset], payload, payload_len);
			payload_len += (hdr_len - hdr_payload_offset);
			hdr_len -= (hdr_len - hdr_payload_offset);
			_librist_crypto_psk_encrypt(key, htobe32(seq), gre_version, payload_wr, payload_wr, payload_len);
#endif
		} else {
			//Single encryption pass suffices
			_librist_crypto_psk_encrypt(key, htobe32(seq), gre_version, payload, payload_wr, payload_len);
		}
		pthread_mutex_unlock(&key_peer->peer_lock);
		SET_BIT(hdr->flags1, 5); // set key bit
		//Write key, our nonce is stored in network byte order so just memcpy it
		memcpy(&hdr_buf[nonce_offset], key->gre_nonce, sizeof(p->key_tx.gre_nonce));
	}

	ssize_t ret;
	int errorcode = 0;

	/* Type 8 wrapping: when both sides have negotiated Advanced Profile,
	 * wrap outbound GRE packets in an AP Type 8 envelope so the remote
	 * can accept them on a unified AP port. */
	if (get_cctx(p)->profile == RIST_PROFILE_ADVANCED &&
	    p->is_advanced && p->remote_supports_advanced) {
		uint8_t gre_flat[RIST_MAX_PACKET_SIZE];
		size_t gre_total = hdr_len + payload_len;
		if (gre_total <= sizeof(gre_flat)) {
			memcpy(gre_flat, hdr_buf, hdr_len);
			memcpy(gre_flat + hdr_len, payload_wr, payload_len);
			ret = (rist_adv_send_type8(p, gre_flat, gre_total) == 0)
			      ? (ssize_t)gre_total : -1;
			if (modifying_payload)
				free(payload_wr);
			return ret;
		}
	}

#ifndef _WIN32
	struct msghdr msghdr;
	struct iovec iov[2];
	iov[0].iov_base = hdr_buf;
	iov[0].iov_len = hdr_len;
	iov[1].iov_base = payload_wr;
	iov[1].iov_len = payload_len;
	msghdr.msg_iov = iov;
	msghdr.msg_iovlen = 2;
	msghdr.msg_name = &p->u.address;
	msghdr.msg_namelen = p->address_len;
	msghdr.msg_control = NULL;
	msghdr.msg_controllen = 0;
	msghdr.msg_flags = 0;
	int retries = 0;
	do {
		ret = rist_transport_sendmsg(p, &msghdr, MSG_DONTWAIT);
		if (RIST_UNLIKELY(ret < 0)) {
			errorcode = errno;
			retries++;
		}
		else {
			errorcode = 0;
			break;
		}
	} while (errorcode == EAGAIN && retries < RIST_MAX_SEND_RETRIES);
#else
	WSAMSG msghdr = { 0 };
	WSABUF iov[2];
	DWORD dwBytes = 0;
	iov[0].len = hdr_len;
	iov[0].buf = (char *)hdr_buf;
	iov[1].len = payload_len;
	iov[1].buf = (char *)payload_wr;
	msghdr.name = (struct sockaddr *)&p->u.address;
	msghdr.namelen = p->address_len;
	msghdr.lpBuffers = iov;
	msghdr.dwBufferCount = 2;
	ret = payload_len + hdr_len;
	int retries = 0;
	do {
		if (RIST_UNLIKELY(WSASendMsg(p->sd, &msghdr, MSG_DONTWAIT, &dwBytes, NULL, NULL) != 0)) {
			ret = -1;
			errorcode = WSAGetLastError();
			retries++;
		}
		else {
			errorcode = 0;
			break;
		}
	} while (errorcode == WSAEWOULDBLOCK && retries < RIST_MAX_SEND_RETRIES);
#endif

	if (modifying_payload) {
		free(payload_wr);
	}

	if (RIST_UNLIKELY(errorcode)) {
		_librist_log_send_error(p, errorcode, (size_t)(hdr_len + payload_len),
		                        "main-profile sendmsg");
	} else if (RIST_UNLIKELY(retries > RIST_MAX_SEND_RETRIES / 5)) {
        struct rist_common_ctx *ctx = get_cctx(p);
        rist_log_priv(ctx, RIST_LOG_WARN, "Send Succeded after retries=%d, ret=%d, socket=%d\n", retries, ret, p->sd);
	}

	return ret;
}

void _librist_proto_gre_send_keepalive(struct rist_peer *p, uint8_t gre_version) {
	uint8_t ka_buf[12] = {0};
	struct rist_gre_keepalive *ka = (struct rist_gre_keepalive *)ka_buf;
	memcpy(ka->mac_array, p->mac_addr, sizeof(ka->mac_array));
	SET_BIT(ka->capabilities1, 0); // Null packet deletion
	if (p->sender_ctx && p->sender_ctx->split_mode != LIBRIST_SPLIT_MODE_OFF)
		SET_BIT(ka->capabilities1, 1); // Pair-split active
	SET_BIT(ka->capabilities1, 2); // SMPTE-2022-7
	SET_BIT(ka->capabilities1, 5); // Bonding
	SET_BIT(ka->capabilities2, 5); // Reduced overhead

	struct rist_common_ctx *cctx = get_cctx(p);
	size_t ka_len = sizeof(struct rist_gre_keepalive);
	if (cctx && cctx->profile == RIST_PROFILE_ADVANCED) {
		/* TR-06-3 Section 5.3.6: I bit (bit 31 of 32-bit capabilities) */
		ka_buf[8]  = 0x80; /* I bit = bit 7 of capabilities3 (= bit 31 overall) */
		ka_buf[9]  = 0;
		ka_buf[10] = 0;
		ka_buf[11] = 0;
		ka_len = 12;
	}
	_librist_proto_gre_send_data(p, 0, RIST_GRE_PROTOCOL_TYPE_KEEPALIVE, ka_buf, ka_len, 0, 0, gre_version);
}

int _librist_proto_gre_parse_keepalive(const uint8_t buf[], size_t buflen, struct rist_keepalive_info  *info) {
	memset(info, 0, sizeof(*info));
	if (buflen < sizeof(struct rist_gre_keepalive)) {
		return -1;
	}
	struct rist_gre_keepalive *ka = (struct rist_gre_keepalive *)buf;
	memcpy(&info->ka, ka, SIZEOF_GRE_KEEPALIVE);
	memcpy(info->mac, ka->mac_array, sizeof(ka->mac_array));
	info->x = CHECK_BIT(ka->capabilities1, 7);
	info->r = CHECK_BIT(ka->capabilities1, 6);
	info->b = CHECK_BIT(ka->capabilities1, 5);
	info->a = CHECK_BIT(ka->capabilities1, 4);
	info->p = CHECK_BIT(ka->capabilities1, 3);
	info->e = CHECK_BIT(ka->capabilities1, 2);
	info->l = CHECK_BIT(ka->capabilities1, 1);
	info->n = CHECK_BIT(ka->capabilities1, 0);
	info->d = CHECK_BIT(ka->capabilities2, 7);
	info->t = CHECK_BIT(ka->capabilities2, 6);
	info->v = CHECK_BIT(ka->capabilities2, 5);
	info->j = CHECK_BIT(ka->capabilities2, 4);
	info->f = CHECK_BIT(ka->capabilities2, 3);
	/* TR-06-3 extended capabilities (4 bytes at offset 8, if present) */
	if (buflen >= sizeof(*ka) + 4) {
		info->adv_i = !!(buf[sizeof(*ka)]     & 0x80); /* bit 31 → I */
		info->adv_g = !!(buf[sizeof(*ka)]     & 0x40); /* bit 30 → G */
		info->adv_c = !!(buf[sizeof(*ka)]     & 0x20); /* bit 29 → C */
		info->json_len = buflen - sizeof(*ka) - 4;
		if (info->json_len > 0)
			info->json = (const char *)&buf[sizeof(*ka) + 4];
	} else {
		info->json_len = buflen - sizeof(*ka);
		if (info->json_len > 0)
			info->json = (const char *)&buf[sizeof(*ka)];
	}
	return 0;
}

void _librist_proto_gre_send_buffer_negotiation(struct rist_peer *p, uint16_t sender_max_buffer, uint16_t receiver_current_buffer) {
	struct rist_buffer_negotiation bn;
	bn.sender_max_buffer_size_ms = htobe16(sender_max_buffer);
	bn.receiver_current_buffer_size_ms = htobe16(receiver_current_buffer);
	bn.protocol_type = 0;
	_librist_proto_gre_send_data(p, 0, RIST_VSF_PROTOCOL_SUBTYPE_BUFFER_NEGOTIATION, (uint8_t *)&bn, sizeof(bn), 0, 0, 2);
}

int _librist_proto_gre_parse_buffer_negotiation(struct rist_peer *p, uint8_t buf[], size_t buflen, uint16_t *sender_max_buffer, uint16_t *receiver_current_buffer) {
	if (buflen < sizeof(struct rist_buffer_negotiation)) {
		return -1;
	}
	struct rist_buffer_negotiation *bn = (struct rist_buffer_negotiation *)buf;
	if (bn->protocol_type != 0) {
		bn->sender_max_buffer_size_ms = 0;
		bn->receiver_current_buffer_size_ms = 0;
		_librist_proto_gre_send_data(p, 0, RIST_VSF_PROTOCOL_SUBTYPE_BUFFER_NEGOTIATION, buf, buflen, 0, 0, 2);
		return -2;
	}
	*sender_max_buffer = be16toh(bn->sender_max_buffer_size_ms);
	*receiver_current_buffer = be16toh(bn->receiver_current_buffer_size_ms);
	return 0;
}
