// SPDX-License-Identifier: GPL-2.0
#include <linux/export.h>
#include <linux/if_vlan.h>
#include <net/ip.h>
#include <net/tso.h>
#include <linux/dma-mapping.h>
#include <linux/unaligned.h>

void tso_build_hdr(const struct sk_buff *skb, char *hdr, struct tso_t *tso,
		   int size, bool is_last)
{
	int hdr_len = skb_transport_offset(skb) + tso->tlen;
	int mac_hdr_len = skb_network_offset(skb);

	memcpy(hdr, skb->data, hdr_len);
	if (!tso->ipv6) {
		struct iphdr *iph = (void *)(hdr + mac_hdr_len);

		iph->id = htons(tso->ip_id);
		iph->tot_len = htons(size + hdr_len - mac_hdr_len);
		tso->ip_id++;
	} else {
		struct ipv6hdr *iph = (void *)(hdr + mac_hdr_len);

		iph->payload_len = htons(size + tso->tlen);
	}
	hdr += skb_transport_offset(skb);
	if (tso->tlen != sizeof(struct udphdr)) {
		struct tcphdr *tcph = (struct tcphdr *)hdr;

		put_unaligned_be32(tso->tcp_seq, &tcph->seq);

		if (!is_last) {
			/* Clear all special flags for not last packet */
			tcph->psh = 0;
			tcph->fin = 0;
			tcph->rst = 0;
		}
	} else {
		struct udphdr *uh = (struct udphdr *)hdr;

		uh->len = htons(sizeof(*uh) + size);
	}
}
EXPORT_SYMBOL(tso_build_hdr);

void tso_build_data(const struct sk_buff *skb, struct tso_t *tso, int size)
{
	tso->tcp_seq += size; /* not worth avoiding this operation for UDP */
	tso->size -= size;
	tso->data += size;

	if ((tso->size == 0) &&
	    (tso->next_frag_idx < skb_shinfo(skb)->nr_frags)) {
		skb_frag_t *frag = &skb_shinfo(skb)->frags[tso->next_frag_idx];

		/* Move to next segment */
		tso->size = skb_frag_size(frag);
		tso->data = skb_frag_address(frag);
		tso->next_frag_idx++;
	}
}
EXPORT_SYMBOL(tso_build_data);

int tso_start(struct sk_buff *skb, struct tso_t *tso)
{
	int tlen = skb_is_gso_tcp(skb) ? tcp_hdrlen(skb) : sizeof(struct udphdr);
	int hdr_len = skb_transport_offset(skb) + tlen;

	tso->tlen = tlen;
	tso->ip_id = ntohs(ip_hdr(skb)->id);
	tso->tcp_seq = (tlen != sizeof(struct udphdr)) ? ntohl(tcp_hdr(skb)->seq) : 0;
	tso->next_frag_idx = 0;
	tso->ipv6 = vlan_get_protocol(skb) == htons(ETH_P_IPV6);

	/* Build first data */
	tso->size = skb_headlen(skb) - hdr_len;
	tso->data = skb->data + hdr_len;
	if ((tso->size == 0) &&
	    (tso->next_frag_idx < skb_shinfo(skb)->nr_frags)) {
		skb_frag_t *frag = &skb_shinfo(skb)->frags[tso->next_frag_idx];

		/* Move to next segment */
		tso->size = skb_frag_size(frag);
		tso->data = skb_frag_address(frag);
		tso->next_frag_idx++;
	}
	return hdr_len;
}
EXPORT_SYMBOL(tso_start);

static int tso_dma_iova_try(struct device *dev, struct tso_dma_map *map,
			    phys_addr_t phys, size_t linear_len,
			    size_t total_len, size_t *offset)
{
	const struct sk_buff *skb;
	unsigned int nr_frags;
	int i;

	if (!dma_iova_try_alloc(dev, &map->iova_state, phys, total_len))
		return 1;

	skb = map->skb;
	nr_frags = skb_shinfo(skb)->nr_frags;

	if (linear_len) {
		if (dma_iova_link(dev, &map->iova_state,
				  phys, *offset, linear_len,
				  DMA_TO_DEVICE, 0))
			goto iova_fail;
		map->linear_len = linear_len;
		*offset += linear_len;
	}

	for (i = 0; i < nr_frags; i++) {
		skb_frag_t *frag = &skb_shinfo(skb)->frags[i];
		unsigned int frag_len = skb_frag_size(frag);

		if (dma_iova_link(dev, &map->iova_state,
				  skb_frag_phys(frag), *offset,
				  frag_len, DMA_TO_DEVICE, 0)) {
			map->nr_frags = i;
			goto iova_fail;
		}
		map->frags[i].len = frag_len;
		*offset += frag_len;
		map->nr_frags = i + 1;
	}

	if (dma_iova_sync(dev, &map->iova_state, 0, total_len))
		goto iova_fail;

	return 0;

iova_fail:
	dma_iova_destroy(dev, &map->iova_state, *offset,
			 DMA_TO_DEVICE, 0);
	memset(&map->iova_state, 0, sizeof(map->iova_state));

	/* reset map state */
	map->frag_idx = -1;
	map->offset = 0;
	map->linear_len = 0;
	map->nr_frags = 0;

	return 1;
}

/**
 * tso_dma_map_init - DMA-map GSO payload regions
 * @map: map struct to initialize
 * @dev: device for DMA mapping
 * @skb: the GSO skb
 * @hdr_len: per-segment header length in bytes
 *
 * DMA-maps the linear payload (after headers) and all frags.
 * Prefers the DMA IOVA API (one contiguous mapping, one IOTLB sync);
 * falls back to per-region dma_map_phys() when IOVA is not available.
 * Positions the iterator at byte 0 of the payload.
 *
 * Return: 0 on success, -ENOMEM on DMA mapping failure (partial mappings
 * are cleaned up internally).
 */
int tso_dma_map_init(struct tso_dma_map *map, struct device *dev,
		     const struct sk_buff *skb, unsigned int hdr_len)
{
	unsigned int linear_len = skb_headlen(skb) - hdr_len;
	unsigned int nr_frags = skb_shinfo(skb)->nr_frags;
	size_t total_len = skb->len - hdr_len;
	size_t offset = 0;
	phys_addr_t phys;
	int i;

	map->dev = dev;
	map->skb = skb;
	map->hdr_len = hdr_len;
	map->frag_idx = -1;
	map->offset = 0;
	map->iova_offset = 0;
	map->total_len = total_len;
	map->linear_len = 0;
	map->nr_frags = 0;
	memset(&map->iova_state, 0, sizeof(map->iova_state));

	if (!total_len)
		return 0;

	if (linear_len)
		phys = virt_to_phys(skb->data + hdr_len);
	else
		phys = skb_frag_phys(&skb_shinfo(skb)->frags[0]);

	if (tso_dma_iova_try(dev, map, phys, linear_len, total_len, &offset)) {
		/* IOVA path failed, map state was reset. Fallback to
		 * per-region dma_map_phys()
		 */
		if (linear_len) {
			map->linear_dma = dma_map_phys(dev, phys, linear_len,
						       DMA_TO_DEVICE, 0);
			if (dma_mapping_error(dev, map->linear_dma))
				return -ENOMEM;
			map->linear_len = linear_len;
		}

		for (i = 0; i < nr_frags; i++) {
			skb_frag_t *frag = &skb_shinfo(skb)->frags[i];
			unsigned int frag_len = skb_frag_size(frag);

			map->frags[i].len = frag_len;
			map->frags[i].dma = dma_map_phys(dev, skb_frag_phys(frag),
							 frag_len, DMA_TO_DEVICE, 0);
			if (dma_mapping_error(dev, map->frags[i].dma)) {
				tso_dma_map_cleanup(map);
				return -ENOMEM;
			}
			map->nr_frags = i + 1;
		}
	}

	if (linear_len == 0 && nr_frags > 0)
		map->frag_idx = 0;

	return 0;
}
EXPORT_SYMBOL(tso_dma_map_init);

/**
 * tso_dma_map_cleanup - unmap all DMA regions in a tso_dma_map
 * @map: the map to clean up
 *
 * Handles both IOVA and fallback paths. For IOVA, calls
 * dma_iova_destroy(). For fallback, unmaps each region individually.
 */
void tso_dma_map_cleanup(struct tso_dma_map *map)
{
	int i;

	if (dma_use_iova(&map->iova_state)) {
		dma_iova_destroy(map->dev, &map->iova_state, map->total_len,
				 DMA_TO_DEVICE, 0);
		memset(&map->iova_state, 0, sizeof(map->iova_state));
	} else {
		if (map->linear_len)
			dma_unmap_phys(map->dev, map->linear_dma,
				       map->linear_len, DMA_TO_DEVICE, 0);

		for (i = 0; i < map->nr_frags; i++)
			dma_unmap_phys(map->dev, map->frags[i].dma,
				       map->frags[i].len, DMA_TO_DEVICE, 0);
	}

	map->linear_len = 0;
	map->nr_frags = 0;
}
EXPORT_SYMBOL(tso_dma_map_cleanup);

/**
 * tso_dma_map_count - count descriptors for a payload range
 * @map: the payload map
 * @len: number of payload bytes in this segment
 *
 * Counts how many contiguous DMA region chunks the next @len bytes
 * will span, without advancing the iterator. On the IOVA path this
 * is always 1 (contiguous). On the fallback path, uses region sizes
 * from the current position.
 *
 * Return: the number of descriptors needed for @len bytes of payload.
 */
unsigned int tso_dma_map_count(struct tso_dma_map *map, unsigned int len)
{
	unsigned int offset = map->offset;
	int idx = map->frag_idx;
	unsigned int count = 0;

	if (!len)
		return 0;

	if (dma_use_iova(&map->iova_state))
		return 1;

	while (len > 0) {
		unsigned int region_len, chunk;

		if (idx == -1)
			region_len = map->linear_len;
		else
			region_len = map->frags[idx].len;

		chunk = min(len, region_len - offset);
		len -= chunk;
		count++;
		offset = 0;
		idx++;
	}

	return count;
}
EXPORT_SYMBOL(tso_dma_map_count);

/**
 * tso_dma_map_next - yield the next DMA address range
 * @map: the payload map
 * @addr: output DMA address
 * @chunk_len: output chunk length
 * @mapping_len: full DMA mapping length when this chunk starts a new
 *               mapping region, or 0 when continuing a previous one.
 *               On the IOVA path this is always 0 (driver must not
 *               do per-region unmaps; use tso_dma_map_cleanup instead).
 * @seg_remaining: bytes left in current segment
 *
 * Yields the next (dma_addr, chunk_len) pair and advances the iterator.
 * On the IOVA path, the entire payload is contiguous so each segment
 * is always a single chunk.
 *
 * Return: true if a chunk was yielded, false when @seg_remaining is 0.
 */
bool tso_dma_map_next(struct tso_dma_map *map, dma_addr_t *addr,
		      unsigned int *chunk_len, unsigned int *mapping_len,
		      unsigned int seg_remaining)
{
	unsigned int region_len, chunk;

	if (!seg_remaining)
		return false;

	/* IOVA path: contiguous DMA range, no region boundaries */
	if (dma_use_iova(&map->iova_state)) {
		*addr = map->iova_state.addr + map->iova_offset;
		*chunk_len = seg_remaining;
		*mapping_len = 0;
		map->iova_offset += seg_remaining;
		return true;
	}

	/* Fallback path: per-region iteration */

	if (map->frag_idx == -1) {
		region_len = map->linear_len;
		chunk = min(seg_remaining, region_len - map->offset);
		*addr = map->linear_dma + map->offset;
	} else {
		region_len = map->frags[map->frag_idx].len;
		chunk = min(seg_remaining, region_len - map->offset);
		*addr = map->frags[map->frag_idx].dma + map->offset;
	}

	*mapping_len = (map->offset == 0) ? region_len : 0;
	*chunk_len = chunk;
	map->offset += chunk;

	if (map->offset >= region_len) {
		map->frag_idx++;
		map->offset = 0;
	}

	return true;
}
EXPORT_SYMBOL(tso_dma_map_next);