package main

import (
	"encoding/binary"
	"errors"
	"flag"
	"fmt"
	"log"
	"math"
	"net"
	"net/http"
	"net/netip"
	"sync"
	"syscall"

	_ "net/http/pprof"

	"golang.org/x/net/ipv4"
	"golang.org/x/net/ipv6"
)

const (
	NUM_BUFFERS = 1
)

var (
	OOB_SIZE = syscall.CmsgSpace(syscall.SizeofInet6Pktinfo) // ipv6 address + iface index, struct in6_pktinfo
)

func (vx *vx46) transform46(msgs4 []ipv4.Message, msgs6 []ipv6.Message) error {
	// embed the "client" ipv4 into the src address for the packet to the upstream vxlan ipv6 endpoint
	// the destination is the upstream vxlan endpoint
	for i := range msgs4 {
		if len(msgs6[i].OOB) < OOB_SIZE {
			msgs6[i].OOB = make([]byte, OOB_SIZE)
			msgs6[i].NN = OOB_SIZE
		}

		var hdr syscall.Cmsghdr
		hdr.SetLen(syscall.CmsgLen(syscall.SizeofInet6Pktinfo))
		hdr.Level = syscall.IPPROTO_IPV6
		hdr.Type = syscall.IPV6_PKTINFO

		oobdata := msgs6[i].OOB[:0]

		// set cmsg header
		oobdata = binary.NativeEndian.AppendUint64(oobdata, hdr.Len)
		oobdata = binary.NativeEndian.AppendUint32(oobdata, uint32(hdr.Level))
		oobdata = binary.NativeEndian.AppendUint32(oobdata, uint32(hdr.Type))

		// set the address
		// prefix + embeded v4 + port
		oobdata = append(oobdata, vx.natprefix[:10]...)
		inUDPAddr4 := msgs4[i].Addr.(*net.UDPAddr)
		oobdata = append(oobdata, inUDPAddr4.IP.To4()[:4]...)
		oobdata = binary.BigEndian.AppendUint16(oobdata, uint16(inUDPAddr4.Port))

		// let the kernel decide which interface to use
		oobdata = binary.NativeEndian.AppendUint32(oobdata, 0) // iface

		msgs6[i].Buffers[0] = msgs4[i].Buffers[0][:msgs4[i].N]
		msgs6[i].Addr = vx.upstream
		msgs6[i].N = msgs4[i].N
	}

	return nil
}

func (vx *vx46) forward46() error {
	msgs4 := make([]ipv4.Message, vx.buffers)
	msgs6 := make([]ipv6.Message, vx.buffers)

	for i := range msgs6 {
		msgs4[i].Buffers = [][]byte{make([]byte, vx.mtu)}
		msgs6[i].Buffers = [][]byte{nil}
	}

	for {
		n, err := vx.pc4.ReadBatch(msgs4[:], 0)
		if n > vx.buffers*3/4 {
			//log.Printf("forward46: %d in", n)
		}
		if err != nil {
			return err
		}
		err = vx.transform46(msgs4[:n], msgs6[:n])
		if err != nil {
			return err
		}
		queue := msgs6[:n]
		for len(queue) > 0 {
			outn, err := vx.pc6.WriteBatch(queue, 0)
			if err != nil {
				return err
			}
			queue = queue[outn:]
		}
	}
}

func (vx *vx46) transform64(msgs6 []ipv6.Message, msgs4 []ipv4.Message) error {
	// embed the "client" ipv4 into the src address for the packet to the upstream vxlan ipv6 endpoint
	// the destination is the upstream vxlan endpoint
	for i := range msgs6 {
		var hdr syscall.Cmsghdr
		hdr.Len = binary.NativeEndian.Uint64(msgs6[i].OOB[0:])
		hdr.Level = int32(binary.NativeEndian.Uint32(msgs6[i].OOB[8:]))
		hdr.Type = int32(binary.NativeEndian.Uint32(msgs6[i].OOB[12:]))

		if hdr.Len != uint64(syscall.CmsgLen(syscall.SizeofInet6Pktinfo)) {
			log.Fatal("OOB: unexpected Len: ", hdr.Len, syscall.CmsgLen(syscall.SizeofInet6Pktinfo))
		}
		if hdr.Level != syscall.IPPROTO_IPV6 {
			log.Fatal("OOB: unexpected Level", hdr.Level)
		}
		if hdr.Type != syscall.IPV6_PKTINFO {
			log.Fatal("OOB: unexpected Type", hdr.Type)
		}

		dstaddr6 := msgs6[i].OOB[16:32]

		msgs4[i].Buffers[0] = msgs6[i].Buffers[0][:msgs6[i].N]
		if msgs4[i].Addr == nil {
			msgs4[i].Addr = &net.UDPAddr{}
		}
		addr := msgs4[i].Addr.(*net.UDPAddr)
		addr.IP = dstaddr6[10:14]
		addr.Port = int(binary.BigEndian.Uint16(dstaddr6[14:16]))
		msgs4[i].N = msgs6[i].N
		msgs4[i].OOB = nil
		msgs4[i].NN = 0
	}

	return nil
}

func (vx *vx46) forward64() error {
	msgs4 := make([]ipv4.Message, vx.buffers)
	msgs6 := make([]ipv6.Message, vx.buffers)

	for i := range msgs6 {
		msgs6[i].Buffers = [][]byte{make([]byte, vx.mtu)}
		msgs6[i].OOB = make([]byte, OOB_SIZE)
		msgs4[i].Buffers = [][]byte{nil}
	}

	for {
		n, err := vx.pc6.ReadBatch(msgs6[:], 0)
		if n > vx.buffers*3/4 {
			//log.Printf("forward64: %d in", n)
		}
		if err != nil {
			return err
		}
		err = vx.transform64(msgs6[:n], msgs4[:n])
		if err != nil {
			return err
		}
		queue := msgs4[:n]
		for len(queue) > 0 {
			outn, err := vx.pc4.WriteBatch(queue, 0)
			if err != nil {
				return err
			}
			queue = queue[outn:]
		}
	}
}

func (vx *vx46) forward() error {
	l4, err := net.ListenUDP("udp4", &net.UDPAddr{Port: int(vx.port)})
	if err != nil {
		return err
	}

	l6, err := net.ListenUDP("udp6", &net.UDPAddr{Port: int(vx.port)})
	if err != nil {
		return err
	}

	vx.pc4 = ipv4.NewPacketConn(l4)
	vx.pc6 = ipv6.NewPacketConn(l6)

	vx.pc6.SetControlMessage(ipv6.FlagDst, true)

	var wg sync.WaitGroup
	var err4 error
	var err6 error

	wg.Add(1)
	go func() {
		err4 = vx.forward46()
		l6.Close()
		wg.Done()
	}()
	wg.Add(1)
	go func() {
		err6 = vx.forward64()
		l4.Close()
		wg.Done()
	}()

	wg.Wait()

	return errors.Join(err4, err6)
}

type vx46 struct {
	pc4          *ipv4.PacketConn
	pc6          *ipv6.PacketConn
	natprefix    [16]byte
	upstreamAddr netip.Addr
	upstream     *net.UDPAddr
	port         uint16
	mtu          uint16
	buffers      int
}

func vx46forward(natprefix netip.Addr, upstreamAddr netip.Addr, port uint16, mtu uint16) error {
	upstream := netip.AddrPortFrom(upstreamAddr, port)
	p, err := net.ListenUDP("udp", &net.UDPAddr{Port: int(port)})
	if err != nil {
		return err
	}

	// we need to know to which address incoming ipv6 packets were sent to
	// extract the client ipv4 and port
	ipv6.NewPacketConn(p).SetControlMessage(ipv6.FlagDst, true)

	defer p.Close()

	b := make([]byte, mtu) // inner ethernet header + vxlan header
	oob := make([]byte, OOB_SIZE)
	for {
		n, oobn, _, ingressSrcAddrPort, err := p.ReadMsgUDPAddrPort(b[:], oob[:])

		cm := ipv6.ControlMessage{}
		if err := cm.Parse(oob[:oobn]); err != nil {
			return err
		}

		ingressDstAddr, ok := netip.AddrFromSlice(cm.Dst)
		if !ok {
			continue
		}

		var ingressDstAddrPort netip.AddrPort
		ingressDstAddrPort = netip.AddrPortFrom(ingressDstAddr.Unmap(), port)
		ingressSrcAddrPort = netip.AddrPortFrom(ingressSrcAddrPort.Addr().Unmap(), ingressSrcAddrPort.Port())

		var egressDstAddrPort netip.AddrPort
		var egressOOB []byte

		if ingressSrcAddrPort.Addr().Is4() {
			// embed the "client" ipv4 into the src address for the packet to the upstream vxlan ipv6 endpoint
			// the destination is the upstream vxlan endpoint
			inaddr4 := ingressSrcAddrPort.Addr().As4()
			egressSrcAddr := natprefix.As16()
			copy(egressSrcAddr[10:14], inaddr4[:])
			binary.BigEndian.PutUint16(egressSrcAddr[14:16], ingressSrcAddrPort.Port())

			cm := ipv6.ControlMessage{Src: net.IP(egressSrcAddr[:])}
			egressOOB = cm.Marshal()
			egressDstAddrPort = upstream
		} else {
			inaddr6 := ingressDstAddrPort.Addr().As16()
			egressDstAddrPort = netip.AddrPortFrom(
				netip.AddrFrom4([4]byte(inaddr6[10:14])), // extract client ipv4
				binary.BigEndian.Uint16(inaddr6[14:16]),  // extract client port
			)
		}

		outn, outoobn, err := p.WriteMsgUDPAddrPort(b[:n], egressOOB, egressDstAddrPort)
		if err != nil {
			return err
		}
		if outn != n {
			return fmt.Errorf("dropped bytes: %d sent of %d", outn, n)
		}
		if outoobn != len(egressOOB) {
			return fmt.Errorf("dropped oob bytes: %d sent of %d", outoobn, len(egressOOB))
		}
	}
}
func main() {
	natprefixStr := flag.String("prefix", "", "local IPv6 base address for a /80 to use for communication with upstream")
	upstreamStr := flag.String("upstream", "", "IPv6 address of the upstream VXLAN endpoint")
	portInt := flag.Int("port", 8472, "port for vxlan communication")
	mtuInt := flag.Uint("mtu", 1422, "buffer size")
	buffersInt := flag.Int("buffers", 64, "number of buffers for I/O batching")
	pprofBool := flag.Bool("pprof", false, "enable pprof http endpoint on http://localhost:6060")

	flag.Parse()

	natprefix, err := netip.ParseAddr(*natprefixStr)
	if err != nil {
		log.Fatalf("Invalid prefix: %s", err)
	}
	upstreamAddr, err := netip.ParseAddr(*upstreamStr)
	if err != nil {
		log.Fatalf("Invalid upstream: %s", err)
	}
	if *portInt > math.MaxUint16 {
		log.Fatalf("port out of range: %d", *portInt)
	}
	port := uint16(*portInt)

	if *mtuInt > math.MaxUint16 {
		log.Fatalf("mtu out of range: %d", *mtuInt)
	}
	mtu := uint16(*mtuInt + 14 + 8) // inner ethernet header + vxlan header

	if *buffersInt < 1 {
		log.Fatalf("buffers < 1: %d", *mtuInt)
	}

	if *pprofBool {
		go func() {
			log.Println(http.ListenAndServe("localhost:6060", nil))
		}()
	}

	if *buffersInt == 1 {
		log.Println(vx46forward(natprefix, upstreamAddr, port, mtu))
	} else {
		vx := vx46{
			natprefix: natprefix.As16(),
			upstream: &net.UDPAddr{
				IP:   net.IP(upstreamAddr.AsSlice()),
				Port: *portInt,
			},
			port:    port,
			mtu:     mtu,
			buffers: *buffersInt,
		}
		log.Println(vx.forward())
	}
}