path: root/tunnel/tools/libwg-go/nat-tun.go

/* SPDX-Identifier-License: MIT
 *
 * Copyright (C) 2019-2021 WireGuard LLC. All Rights Reserved.
 */

package main

// TODO debug IPv6 NAT
// TODO implement TCP session tracking according to RFC 7857,

import (
	"encoding/binary"
	"fmt"
	"os"

	"golang.zx2c4.com/go118/netip"

	"golang.zx2c4.com/wireguard/device"
	"golang.zx2c4.com/wireguard/tun"
)

type natEntry struct {
	// intAddr         netip.Addr
	// extAddr         netip.Addr
        ipVersion       int
        loop            bool
	srcAddr         netip.Addr
        dstPort         int
	proxyAddr       netip.Addr
	proxyPort       int
}

type natTun struct {
        conntrack       *Conntrack
	tun             tun.Device
        translations    []natEntry
        l               *device.Logger
}

func NewNatTun(t tun.Device, ct *Conntrack, logger *device.Logger) (dev *natTun, err error) {
	dev = &natTun{
                conntrack: ct,
		tun: t,
                translations: make([]natEntry, 4),
                l: logger,
	}

	err = nil
	return
}

func (tun *natTun) addTranslation(ipv6 bool, dstPort int, proxyPort int) {
        var ipVersion int
        loop := true
        var srcStr string
        var proxyStr string
        var srcAddr netip.Addr

        if ipv6 {
                ipVersion = IPV6_VERSION
                //srcStr = "fdba:4b51:1606:a61d::1"
                //proxyStr = "fdba:4b51:1606:a61d::2"
                loop = false
                proxyStr = "2001:470:de6f:2fff::f780"
                proxyPort = 8888
        } else {
                ipVersion = IPV4_VERSION
                //srcStr = "169.254.0.1"
                //proxyStr = "169.254.0.2"
                loop = false
                proxyStr = "10.49.32.1"
                proxyPort = 8888
                
        }

        if srcStr != "" {
                var err error
                srcAddr, err = netip.ParseAddr(srcStr)
                if err != nil {
                        return
                }
        }

        proxyAddr, err := netip.ParseAddr(proxyStr)
        if err != nil {
                return
        }

        // proxyPort = 8888
        // proxyAddr, err = netip.ParseAddr("10.49.124.115")
        // if err != nil {
        //         return
        // }

	translation := natEntry{
		// intAddr: intAddr,
		// extAddr: extAddr,
                ipVersion: ipVersion,
                loop: loop,
		srcAddr: srcAddr,
                dstPort: dstPort,
		proxyAddr: proxyAddr,
		proxyPort: proxyPort,
	}

        tun.translations = append(tun.translations, translation)
}

func (tun *natTun) Name() (string, error) {
	return tun.tun.Name()
}

func (tun *natTun) File() *os.File {
	return nil
}

func (tun *natTun) Events() chan tun.Event {
	return tun.tun.Events()
}

const (
	IPV4_VERSION = 4
	IPV4_HEADER_LEN = 20	// TODO support option headers, refer to https://cs.opensource.google/go/x/net/+/fe4d6282:ipv4/header.go
	IPV4_HEADER_PROTOCOL = 9
	IPV4_HEADER_CHECKSUM = 10
	IPV4_HEADER_SRC_ADDR = 12
	IPV4_HEADER_DST_ADDR = 16

	IPV6_VERSION = 6
	IPV6_HEADER_LEN = 40	// TODO support additional headers
	IPV6_HEADER_NEXT_HEADER = 6
	IPV6_HEADER_SRC_ADDR = 8
	IPV6_HEADER_DST_ADDR = 24

	TCP_HEADER_LEN = 40
	TCP_HEADER_SRC_PORT = 0
	TCP_HEADER_DST_PORT = 2
        TCP_HEADER_FLAGS = 12
	TCP_HEADER_CHECKSUM = 16

        TCP_FLAGS_FIN = 0x01
        TCP_FLAGS_SYN = 0x02
        TCP_FLAGS_RST = 0x04

	UDP_HEADER_SRC_PORT = 0
	UDP_HEADER_DST_PORT = 2
	UDP_HEADER_CHECKSUM = 6

	PROTO_TCP = 6
	PROTO_UDP = 17
)

func getUint16(header []byte, offset int) int {
	return int(binary.BigEndian.Uint16(header[offset:offset+2]))
}

func putUint16(header []byte, offset int, value int) {
	binary.BigEndian.PutUint16(header[offset:offset+2], uint16(value))
}

func getIPv4SrcAddr(header []byte) netip.Addr {
	src, _ := netip.AddrFromSlice(header[IPV4_HEADER_SRC_ADDR:IPV4_HEADER_SRC_ADDR+4])
	return src
}

func putIPv4SrcAddr(header []byte, addr netip.Addr) {
	copy(header[IPV4_HEADER_SRC_ADDR:IPV4_HEADER_SRC_ADDR+4], addr.AsSlice())
}

func getIPv4DstAddr(header []byte) netip.Addr {
	dst, _ := netip.AddrFromSlice(header[IPV4_HEADER_DST_ADDR:IPV4_HEADER_DST_ADDR+4])
	return dst
}

func putIPv4DstAddr(header []byte, addr netip.Addr) {
	copy(header[IPV4_HEADER_DST_ADDR:IPV4_HEADER_DST_ADDR+4], addr.AsSlice())
}

func getIPv6SrcAddr(header []byte) netip.Addr {
	src, _ := netip.AddrFromSlice(header[IPV6_HEADER_SRC_ADDR:IPV6_HEADER_SRC_ADDR+16])
	return src
}

func putIPv6SrcAddr(header []byte, addr netip.Addr) {
	copy(header[IPV6_HEADER_SRC_ADDR:IPV6_HEADER_SRC_ADDR+16], addr.AsSlice())
}

func getIPv6DstAddr(header []byte) netip.Addr {
	dst, _ := netip.AddrFromSlice(header[IPV6_HEADER_DST_ADDR:IPV6_HEADER_DST_ADDR+16])
	return dst
}

func putIPv6DstAddr(header []byte, addr netip.Addr) {
	copy(header[IPV6_HEADER_DST_ADDR:IPV6_HEADER_DST_ADDR+16], addr.AsSlice())
}

func getSrcAddr(header []byte, version int) netip.Addr {
	if version == IPV4_VERSION {
		return getIPv4SrcAddr(header)
	} else if version == IPV6_VERSION {
		return getIPv6SrcAddr(header)
	} else {
                return netip.IPv6Unspecified()
        }
}

func getDstAddr(header []byte, version int) netip.Addr {
	if version == IPV4_VERSION {
		return getIPv4DstAddr(header)
	} else if version == IPV6_VERSION {
		return getIPv6DstAddr(header)
	} else {
                return netip.IPv6Unspecified()
	}
}

func updateIPv4Checksum(header []byte, updateFunc func (checksum int) int) {
	checksum := getUint16(header, IPV4_HEADER_CHECKSUM)
	putUint16(header, IPV4_HEADER_CHECKSUM, updateFunc(checksum))
}

func getTransport(header []byte, version int) (int, []byte) {
	if version == IPV4_VERSION {
		protocol := int(header[IPV4_HEADER_PROTOCOL])
		return protocol, header[IPV4_HEADER_LEN:] // TODO take option headers into account
	} else if version == IPV6_VERSION {
		nextHeader := int(header[IPV6_HEADER_NEXT_HEADER])
		return nextHeader, header[IPV6_HEADER_LEN:] // TODO take additional headers into account
	} else {
		return 0, nil
	}
}

func updateTransportChecksum(header []byte, version int, updateFunc func (checksum int) int) {
	transport, transportPayload := getTransport(header, version)

	if transport == PROTO_UDP {
		// UDP checksum should be calculated with by swinging the 17th carry bit around. Use no checksum.
		//checksum := getUint16(transportPayload, UDP_HEADER_CHECKSUM)
		//putUint16(transportPayload, UDP_HEADER_CHECKSUM, updateFunc(checksum))
		putUint16(transportPayload, UDP_HEADER_CHECKSUM, 0)
	} else if transport == PROTO_TCP {
		checksum := getUint16(transportPayload, TCP_HEADER_CHECKSUM)
		putUint16(transportPayload, TCP_HEADER_CHECKSUM, updateFunc(checksum))
	}
}

func onesComplementUint16(value int) int {
	return (0x10000 + ^(value & 0xffff)) & 0xffff
}

func onesSumUint16(value int) int {
	for value > 0xffff {
		value = (value & 0xffff) + (value >> 16)
	}
	return value
}

func calculateChecksum(hc int, m, mPrim int) int {
	return onesComplementUint16(onesSumUint16(onesComplementUint16(hc) + onesComplementUint16(m) + mPrim))
}

func updateAddr(header []byte, version int, getFunc func (header []byte) netip.Addr, putFunc func (header []byte, addr netip.Addr), updateFunc func (netip.Addr) netip.Addr) {
        addr := getFunc(header)
        newAddr := updateFunc(addr)

        if newAddr != addr {
                putFunc(header, newAddr)
                if version == IPV4_VERSION {
                        updateIPv4Checksum(header, func (checksum int) int {
                                addr4 := addr.As4()
                                newAddr4 := newAddr.As4()
                                checksum = calculateChecksum(checksum, int(binary.BigEndian.Uint16(addr4[0:2])), int(binary.BigEndian.Uint16(newAddr4[0:2])))
                                checksum = calculateChecksum(checksum, int(binary.BigEndian.Uint16(addr4[2:4])), int(binary.BigEndian.Uint16(newAddr4[2:4])))
                                return checksum
                        })
                }
                updateTransportChecksum(header, version, func (checksum int) int {
                        if checksum == 0 {
                                // if version == IPV6_VERSION {
                                //         // Log error
                                // }
                                return 0
                        } else if version == IPV4_VERSION {
                                addr4 := addr.As4()
                                newAddr4 := newAddr.As4()
                                checksum = calculateChecksum(checksum, int(binary.BigEndian.Uint16(addr4[0:2])), int(binary.BigEndian.Uint16(newAddr4[0:2])))
                                checksum = calculateChecksum(checksum, int(binary.BigEndian.Uint16(addr4[2:4])), int(binary.BigEndian.Uint16(newAddr4[2:4])))
                                return checksum
                        } else if version == IPV6_VERSION {
                                addr6 := addr.As16()
                                newAddr6 := newAddr.As16()
                                for i := 0; i < 8; i++ {
                                        checksum = calculateChecksum(checksum, int(binary.BigEndian.Uint16(addr6[i*2:(i+1)*2])), int(binary.BigEndian.Uint16(newAddr6[i*2:(i+1)*2])))
                                }
                                return checksum
                        } else {
                                return checksum
                        }
                })
	}
}

func updateSrcAddr(header []byte, version int, updateFunc func (netip.Addr) netip.Addr ) {
	if version == IPV4_VERSION {
		updateAddr(header, version, getIPv4SrcAddr, putIPv4SrcAddr, updateFunc)
	} else if version == IPV6_VERSION {
		updateAddr(header, version, getIPv6SrcAddr, putIPv6SrcAddr, updateFunc)
        }
}

func updateDstAddr(header []byte, version int, updateFunc func (netip.Addr) netip.Addr ) {
	if version == IPV4_VERSION {
		updateAddr(header, version, getIPv4DstAddr, putIPv4DstAddr, updateFunc)
	} else if version == IPV6_VERSION {
		updateAddr(header, version, getIPv6DstAddr, putIPv6DstAddr, updateFunc)
	}
}

func updateDstPort(header []byte, version int, updateFunc func (int) int ) {
	transport, transportPayload := getTransport(header, version)

	if transport == PROTO_TCP {
		port := getUint16(transportPayload, TCP_HEADER_DST_PORT)
		newPort := updateFunc(port)

		if newPort != port {
			putUint16(transportPayload, TCP_HEADER_DST_PORT, newPort)

			updateTransportChecksum(header, version, func (checksum int) int {
				if checksum == 0 {
                                        // if version == IPV6_VERSION {
                                        //         // Log error
                                        // }
					return 0
				} else {
					checksum = calculateChecksum(checksum, int(port), int(newPort))
					return checksum
				}
			})
		}
	}
}

func updateSrcPort(header []byte, version int, updateFunc func (int) int ) {
	transport, transportPayload := getTransport(header, version)

	if transport == PROTO_TCP {
		port := getUint16(transportPayload, TCP_HEADER_SRC_PORT)
		newPort := updateFunc(port)

		if newPort != port {
			putUint16(transportPayload, TCP_HEADER_SRC_PORT, newPort)

			updateTransportChecksum(header, version, func (checksum int) int {
				if checksum == 0 {
                                        // if version == IPV6_VERSION {
                                        //         // Log error
                                        // }
					return 0
				} else {
					checksum = calculateChecksum(checksum, int(port), int(newPort))
					return checksum
				}
			})
		}
	}
}

func getSrcPort(header []byte, version int) int {
	transport, transportPayload := getTransport(header, version)

	if transport == PROTO_TCP {
		return getUint16(transportPayload, TCP_HEADER_SRC_PORT)
	} else if transport == PROTO_UDP {
		return getUint16(transportPayload, UDP_HEADER_SRC_PORT)
	} else {
		return -1
	}
}

func getDstPort(header []byte, version int) int {
	transport, transportPayload := getTransport(header, version)

	if transport == PROTO_TCP {
		return getUint16(transportPayload, TCP_HEADER_DST_PORT)
	} else if transport == PROTO_UDP {
		return getUint16(transportPayload, UDP_HEADER_DST_PORT)
	} else {
		return -1
	}
}

func getTcpSrcPort(tcpHeader []byte) int {
	return getUint16(tcpHeader, TCP_HEADER_SRC_PORT)
}

func putTcpSrcPort(tcpHeader []byte, port int) {
	putUint16(tcpHeader, TCP_HEADER_SRC_PORT, port)
}

func getTcpDstPort(tcpHeader []byte) int {
	return getUint16(tcpHeader, TCP_HEADER_DST_PORT)
}

func putTcpDstPort(tcpHeader []byte, port int) {
	putUint16(tcpHeader, TCP_HEADER_DST_PORT, port)
}

func updateTcpSrcPort(tcpHeader []byte, updateFunc func (int) int ) {
	port := getTcpSrcPort(tcpHeader)
	putTcpSrcPort(tcpHeader, updateFunc(port))
}

func updateTcpDstPort(tcpHeader []byte, updateFunc func (int) int ) {
	port := getTcpDstPort(tcpHeader)
	putTcpDstPort(tcpHeader, updateFunc(port))
}

func getIPVersion(header []byte, len int) (int, error) {
	version := int(header[0]) >> 4

	if version == IPV4_VERSION {
		if len >= IPV4_HEADER_LEN {
			return version, nil
		} else {
			return -1, fmt.Errorf("IPv4 header length: %d < %d", len, IPV4_HEADER_LEN)
		}		
	} else if version == IPV6_VERSION {
		if len >= IPV6_HEADER_LEN {
			return version, nil
		} else {
			return -1, fmt.Errorf("IPv6 header length: %d < %d", len, IPV6_HEADER_LEN)
		}
	} else {
		return -1, fmt.Errorf("Unknown IP version: %d", version)
	}
}

func testFlag(flags int, flag int) bool {
        return (flags & flag) != 0
}

func (tun *natTun) Read(buf []byte, offset int) (int, error) {
	len, err := tun.tun.Read(buf, offset)
	if err != nil || len <= 0 {
                return len, err
        }

        header := buf[offset:]
        version, err := getIPVersion(header, len)

	if err != nil {
                return len, nil
        }

        if version == IPV4_VERSION || version == IPV6_VERSION {
                transport, transportPayload := getTransport(header, version)
                if (transport != PROTO_TCP) {
                        return len, nil
                }
                
                flags := getUint16(transportPayload, TCP_HEADER_FLAGS) & 0x1f
                srcAddr := getSrcAddr(header, version)
                srcPort := getSrcPort(header, version)
                dstAddr := getDstAddr(header, version)
                dstPort := getDstPort(header, version)

                src := netip.AddrPortFrom(srcAddr, uint16(srcPort))
                dst := netip.AddrPortFrom(dstAddr, uint16(dstPort))
                conn := Connection(src, dst)

                m, ok := tun.conntrack.lookupConnection(conn)
                if ok {
                        tun.l.Verbosef("Found mapping: [%v] %v", conn, m)
                        var event Event
                        var newConn connection
                        if m.orig == conn {
                                // Client
                                newConn = m.nat
                                tun.l.Verbosef("Client: %v", newConn)
                                if testFlag(flags, TCP_FLAGS_SYN) {
                                        event = ClientSyn
                                } else if testFlag(flags, TCP_FLAGS_FIN) {
                                        event = ClientFin
                                } else if testFlag(flags, TCP_FLAGS_RST) {
                                        event = ClientRst
                                }
                        } else {
                                // Server
                                newConn = m.orig
                                tun.l.Verbosef("Server: %v", newConn)
                                if testFlag(flags, TCP_FLAGS_SYN) {
                                        event = ServerSyn
                                } else if testFlag(flags, TCP_FLAGS_FIN) {
                                        event = ServerFin
                                } else if testFlag(flags, TCP_FLAGS_RST) {
                                        event = ServerRst
                                } else {
                                        event = DataPkt
                                }
                        }                                

                        if event != None {
                                tun.l.Verbosef("Event: %v", event)
                                tun.conntrack.event(conn, event)
                        }

                        if m.state == Closed {
                                // TODO log
                                return 0, fmt.Errorf("Mapping closed")
                        }

                        updateSrcAddr(header, version, func(netip.Addr) netip.Addr { return newConn.dst.Addr() })
                        updateSrcPort(header, version, func(int) int { return int(newConn.dst.Port()) })
                        updateDstAddr(header, version, func(netip.Addr) netip.Addr { return newConn.src.Addr() })
                        updateDstPort(header, version, func(int) int { return int(newConn.src.Port()) })
                        // Write back to tunnel
                        writeBuf := buf[:offset+len]
                        written, err := tun.Write(writeBuf, offset)
                        tun.l.Verbosef("Written: %v %v", written, err)
                        if err != nil {
                                return 0, err
                        } else if written != len {
                                return 0, fmt.Errorf("NAT rev buffer partly written %v != %v", written, len)
                        } else {
                                return 0, nil
                        }
                } else {
                        if !testFlag(flags, TCP_FLAGS_SYN) {
                                return len, nil
                        }

                        for _, entry := range(tun.translations) {
                                if version != entry.ipVersion {
                                        continue
                                }

                                if dstPort != entry.dstPort {
                                        continue
                                }

                                newDstPort := 0
                                var newSrcAddr netip.Addr
                                var newDstAddr netip.Addr
                                
                                updateDstPort(header, version, func(port int) int {
                                        newDstPort = entry.proxyPort
                                        return newDstPort
                                })

                                if entry.srcAddr.IsValid() {
                                        updateSrcAddr(header, version, func(addr netip.Addr) netip.Addr {
                                                newSrcAddr = entry.srcAddr
                                                return newSrcAddr
                                        })
                                } else {
                                        newSrcAddr = srcAddr
                                }

                                updateDstAddr(header, version, func(addr netip.Addr) netip.Addr {
                                        newDstAddr = entry.proxyAddr
                                        return newDstAddr
                                })

                                orig := Connection(netip.AddrPortFrom(srcAddr, uint16(srcPort)),
                                        netip.AddrPortFrom(dstAddr, uint16(dstPort)))
                                nat := Connection(netip.AddrPortFrom(newDstAddr, uint16(newDstPort)),
                                        netip.AddrPortFrom(newSrcAddr, uint16(srcPort)))
                                tun.l.Verbosef("New conn: %v -> %v", orig, nat)
                                
                                tun.conntrack.addConnection(orig, nat)

                                if !entry.loop {
                                        return len, nil
                                }

                                // Write back to tunnel
                                writeBuf := buf[:offset+len]
                                written, err := tun.Write(writeBuf, offset)
                                tun.l.Verbosef("Written: %v %v", written, err)
                                if err != nil {
                                        return 0, err
                                } else if written != len {
                                        return 0, fmt.Errorf("NAT buffer partly written %v != %v", written, len)
                                } else {
                                        return 0, nil
                                }
                        }
                }
        }

	return len, err
}

func (tun *natTun) Write(buf []byte, offset int) (int, error) {
	len := len(buf)

        if len <= 0 {
                return len, nil
        }
        
        header := buf[offset:]
        version, err := getIPVersion(header, len)

        if err == nil && (version == IPV4_VERSION || version == IPV6_VERSION) {
                transport, transportPayload := getTransport(header, version)
                if (transport == PROTO_TCP) {
                        flags := getUint16(transportPayload, TCP_HEADER_FLAGS) & 0x1f
                        // TODO remove unused code
                        srcAddr := getSrcAddr(header, version)
                        srcPort := getSrcPort(header, version)
                        dstAddr := getDstAddr(header, version)
                        dstPort := getDstPort(header, version)

                        src := netip.AddrPortFrom(srcAddr, uint16(srcPort))
                        dst := netip.AddrPortFrom(dstAddr, uint16(dstPort))
                        conn := Connection(src, dst)

                        m, ok := tun.conntrack.lookupConnection(conn)
                        if ok {
                                tun.l.Verbosef("Write: Found mapping: [%v] %v", conn, m)
                                var event Event
                                var newConn connection

                                if m.orig == conn {
                                        // Client
                                        newConn = m.nat
                                        tun.l.Verbosef("Client: %v", newConn)
                                        if testFlag(flags, TCP_FLAGS_SYN) {
                                                event = ClientSyn
                                        } else if testFlag(flags, TCP_FLAGS_FIN) {
                                                event = ClientFin
                                        } else if testFlag(flags, TCP_FLAGS_RST) {
                                                event = ClientRst
                                        }
                                } else {
                                        // Server
                                        newConn = m.orig
                                        tun.l.Verbosef("Server: %v", newConn)
                                        if testFlag(flags, TCP_FLAGS_SYN) {
                                                event = ServerSyn
                                        } else if testFlag(flags, TCP_FLAGS_FIN) {
                                                event = ServerFin
                                        } else if testFlag(flags, TCP_FLAGS_RST) {
                                                event = ServerRst
                                        } else {
                                                event = DataPkt
                                        }
                                }                

                                if event != None {
                                        tun.l.Verbosef("Event: %v", event)
                                        tun.conntrack.event(conn, event)
                                }

                                if m.state == Closed {
                                        // TODO log
                                        return 0, fmt.Errorf("Mapping closed")
                                }

                                updateSrcAddr(header, version, func(netip.Addr) netip.Addr { return newConn.dst.Addr() })
                                updateSrcPort(header, version, func(int) int { return int(newConn.dst.Port()) })
                                updateDstAddr(header, version, func(netip.Addr) netip.Addr { return newConn.src.Addr() })
                                updateDstPort(header, version, func(int) int { return int(newConn.src.Port()) })
                        }
                }
        }

	return tun.tun.Write(buf, offset)
}

func (tun *natTun) Flush() error {
	return tun.tun.Flush()
}

func (tun *natTun) Close() error {
	return tun.tun.Close()
}

func (tun *natTun) MTU() (int, error) {
	return tun.tun.MTU()
}