package main

import (
	"errors"
	"fmt"
	"syscall"
	"unsafe"
)

const (
	getInfo    = 64
	setReplace = 64
)

func readTable(fd int) (tableInfo, error) {
	var info tableInfo
	copy(info.Name[:], "nat")
	size := uint32(unsafe.Sizeof(info))
	// #nosec G103 -- Direct Syscall6 pointer arguments remain fixed until the call returns.
	_, _, err := syscall.Syscall6(syscall.SYS_GETSOCKOPT, uintptr(fd), syscall.IPPROTO_IP, getInfo, uintptr(unsafe.Pointer(&info)), uintptr(unsafe.Pointer(&size)), 0)
	if err != 0 {
		return tableInfo{}, err
	}
	return info, nil
}

func newReplace(value config, count uint32) *replacement {
	result := &replacement{counters: make([]counter, count)}
	// ipt_replace embeds this address, outside the syscall's directly pinned arguments.
	result.pins.Pin(&result.counters[0])
	result.data = table(value, count, result.address())
	return result
}

func (value *replacement) address() uint64 {
	// #nosec G103 -- newReplace pins this buffer until the replacement syscall completes.
	return uint64(uintptr(unsafe.Pointer(&value.counters[0])))
}

func writeTable(fd int, value *replacement) error {
	// #nosec G103 -- The payload is a direct syscall argument; its embedded counter address is pinned.
	_, _, err := syscall.Syscall6(syscall.SYS_SETSOCKOPT, uintptr(fd), syscall.IPPROTO_IP, setReplace, uintptr(unsafe.Pointer(&value.data[0])), uintptr(len(value.data)), 0)
	if err != 0 {
		return err
	}
	return nil
}

func install(value config) (err error) {
	if unsafe.Sizeof(uintptr(0)) != 8 {
		return fmt.Errorf("unsupported iptables ABI")
	}
	fd, err := syscall.Socket(syscall.AF_INET, syscall.SOCK_RAW|syscall.SOCK_CLOEXEC, syscall.IPPROTO_RAW)
	if err != nil {
		return fmt.Errorf("open NAT socket: %w", err)
	}
	defer func() { err = errors.Join(err, syscall.Close(fd)) }()
	info, err := readTable(fd)
	if err != nil {
		return fmt.Errorf("read NAT table: %w", err)
	}
	if info.Count == 0 {
		return fmt.Errorf("NAT table has no counter entries")
	}
	update := newReplace(value, info.Count)
	defer update.pins.Unpin()
	if err := writeTable(fd, update); err != nil {
		return fmt.Errorf("install NAT table: %w", err)
	}
	return nil
}