package connectip

import (
        "errors"
        "fmt"
        "net/netip"
        "slices"
        "sync/atomic"
)

var (
        rejectedIPv4Prefix = netip.PrefixFrom(netip.IPv4Unspecified(), 32)
        rejectedIPv6Prefix = netip.PrefixFrom(netip.IPv6Unspecified(), 128)
)

// AddressRequestID identifies an address request. Zero denotes an unsolicited assignment.
type AddressRequestID uint64

// AddressRequest is a request received from the peer.
// The zero value is not valid.
type AddressRequest struct {
        Prefixes []netip.Prefix

        conn      *Conn
        requested *addressRequestCapsule
        // Use a pointer so copies of the request cannot be answered independently.
        responded *atomic.Bool
}

func newAddressRequest(conn *Conn, requested *addressRequestCapsule) *AddressRequest {
        return &AddressRequest{
                Prefixes:  slices.Clone(requested.Prefixes),
                conn:      conn,
                requested: requested,
                responded: &atomic.Bool{},
        }
}

// Respond queues the complete address assignment for the peer.
// Assignments must have one entry per original requested prefix, in the same order.
// A zero-value netip.Prefix rejects that request.
// Additional contains any other valid prefixes to include, with request ID zero.
// All prefixes must have their host bits cleared.
// A request can be answered only once.
func (r *AddressRequest) Respond(assignments, additional []netip.Prefix) error {
        if r.conn == nil {
                return errors.New("connect-ip: invalid address request")
        }
        if len(assignments) != len(r.requested.RequestIDs) {
                return fmt.Errorf(
                        "connect-ip: expected %d address assignments, got %d",
                        len(r.requested.RequestIDs),
                        len(assignments),
                )
        }
        capsule := &addressAssignCapsule{
                AssignedAddresses: make([]AssignedAddress, 0, len(assignments)+len(additional)),
        }
        var zeroPrefix netip.Prefix
        for i, p := range assignments {
                if p == zeroPrefix {
                        if r.requested.Prefixes[i].Addr().Is4() {
                                p = rejectedIPv4Prefix
                        } else {
                                p = rejectedIPv6Prefix
                        }
                } else if !p.IsValid() || p != p.Masked() {
                        return fmt.Errorf("connect-ip: invalid assigned prefix %d: %s", i, p)
                }
                capsule.AssignedAddresses = append(
                        capsule.AssignedAddresses,
                        AssignedAddress{RequestID: r.requested.RequestIDs[i], IPPrefix: p},
                )
        }
        for i, p := range additional {
                if !p.IsValid() || p != p.Masked() {
                        return fmt.Errorf("connect-ip: invalid additional prefix %d: %s", i, p)
                }
                capsule.AssignedAddresses = append(capsule.AssignedAddresses, AssignedAddress{IPPrefix: p})
        }
        // claim the response after validation, so invalid inputs can be retried
        if !r.responded.CompareAndSwap(false, true) {
                return errors.New("connect-ip: address request already answered")
        }
        if err := r.conn.sendAddressAssignment(capsule); err != nil {
                r.responded.Store(false)
                return err
        }
        return nil
}

package connectip

import (
        "encoding/binary"
        "errors"
        "fmt"
        "io"
        "maps"
        "net/netip"
        "slices"

        "golang.org/x/net/dns/dnsmessage"

        "github.com/quic-go/quic-go/http3"
        "github.com/quic-go/quic-go/quicvarint"
)

const (
        capsuleTypeAddressAssign      http3.CapsuleType = 1
        capsuleTypeAddressRequest     http3.CapsuleType = 2
        capsuleTypeRouteAdvertisement http3.CapsuleType = 3
        // draft-ietf-masque-connect-ip-dns-06
        capsuleTypeDNSAssign http3.CapsuleType = 0x1ace79ec
        capsuleTypePREF64    http3.CapsuleType = 0x274c0fbc
)

// Bound the memory used to parse and retain a peer's addresses and routes.
const (
        maxAddressesPerCapsule = 8192
        maxRoutesPerCapsule    = 8192
)

type capsule interface{ append([]byte) []byte }

// parseCapsule parses a capsule of a type known to connect-ip-go.
// Capsules of unknown types are skipped, as required by RFC 9297, and returned as nil.
func parseCapsule(typ http3.CapsuleType, r http3.CapsuleReader) (capsule, error) {
        switch typ {
        case capsuleTypeAddressAssign:
                return parseAddressAssignCapsule(r)
        case capsuleTypeAddressRequest:
                return parseAddressRequestCapsule(r)
        case capsuleTypeRouteAdvertisement:
                return parseRouteAdvertisementCapsule(r)
        case capsuleTypeDNSAssign:
                return parseDNSAssignCapsule(r)
        case capsuleTypePREF64:
                return parsePREF64Capsule(r)
        default:
                return nil, r.Discard()
        }
}

// addressAssignCapsule represents an ADDRESS_ASSIGN capsule
type addressAssignCapsule struct {
        AssignedAddresses []AssignedAddress
}

// AssignedAddress represents an Assigned Address within an ADDRESS_ASSIGN capsule
type AssignedAddress struct {
        // RequestID is zero for an unsolicited assignment.
        RequestID AddressRequestID
        IPPrefix  netip.Prefix
}

// Rejected reports whether the assignment is a refusal (0.0.0.0/32 or ::/128).
func (a AssignedAddress) Rejected() bool {
        return a.IPPrefix == rejectedIPv4Prefix || a.IPPrefix == rejectedIPv6Prefix
}

func (a AssignedAddress) len() int {
        return quicvarint.Len(uint64(a.RequestID)) + 1 + a.IPPrefix.Addr().BitLen()/8 + 1
}

// addressRequestCapsule represents an ADDRESS_REQUEST capsule
type addressRequestCapsule struct {
        // RequestIDs and Prefixes have matching lengths and order.
        RequestIDs []AddressRequestID
        Prefixes   []netip.Prefix
}

func parseAddressAssignCapsule(r http3.CapsuleReader) (*addressAssignCapsule, error) {
        var assignedAddresses []AssignedAddress
        for r.Remaining() > 0 {
                if len(assignedAddresses) >= maxAddressesPerCapsule {
                        return nil, fmt.Errorf("ADDRESS_ASSIGN capsule contains too many addresses (maximum %d)", maxAddressesPerCapsule)
                }
                requestID, prefix, err := parseAddress(r)
                if err != nil {
                        return nil, err
                }
                assignedAddresses = append(assignedAddresses, AssignedAddress{RequestID: AddressRequestID(requestID), IPPrefix: prefix})
        }
        return &addressAssignCapsule{AssignedAddresses: assignedAddresses}, nil
}

func (c *addressAssignCapsule) append(b []byte) []byte {
        totalLen := 0
        for _, addr := range c.AssignedAddresses {
                totalLen += addr.len()
        }

        b = quicvarint.Append(b, uint64(capsuleTypeAddressAssign))
        b = quicvarint.Append(b, uint64(totalLen))

        for _, addr := range c.AssignedAddresses {
                b = quicvarint.Append(b, uint64(addr.RequestID))
                if addr.IPPrefix.Addr().Is4() {
                        b = append(b, 4)
                } else {
                        b = append(b, 6)
                }
                b = append(b, addr.IPPrefix.Addr().AsSlice()...)
                b = append(b, byte(addr.IPPrefix.Bits()))
        }
        return b
}

func parseAddressRequestCapsule(r http3.CapsuleReader) (*addressRequestCapsule, error) {
        if r.Remaining() == 0 {
                return nil, errors.New("ADDRESS_REQUEST capsule contains no addresses")
        }
        capsule := &addressRequestCapsule{}
        for r.Remaining() > 0 {
                if len(capsule.Prefixes) >= maxAddressesPerCapsule {
                        return nil, fmt.Errorf("ADDRESS_REQUEST capsule contains too many addresses (maximum %d)", maxAddressesPerCapsule)
                }
                requestID, prefix, err := parseAddress(r)
                if err != nil {
                        return nil, err
                }
                if requestID == 0 {
                        return nil, errors.New("ADDRESS_REQUEST capsule contains a zero request ID")
                }
                capsule.RequestIDs = append(capsule.RequestIDs, AddressRequestID(requestID))
                capsule.Prefixes = append(capsule.Prefixes, prefix)
        }
        return capsule, nil
}

func (c *addressRequestCapsule) append(b []byte) []byte {
        var totalLen int
        for i, p := range c.Prefixes {
                totalLen += quicvarint.Len(uint64(c.RequestIDs[i])) + 1 + p.Addr().BitLen()/8 + 1
        }

        b = quicvarint.Append(b, uint64(capsuleTypeAddressRequest))
        b = quicvarint.Append(b, uint64(totalLen))

        for i, p := range c.Prefixes {
                b = quicvarint.Append(b, uint64(c.RequestIDs[i]))
                if p.Addr().Is4() {
                        b = append(b, 4)
                } else {
                        b = append(b, 6)
                }
                b = append(b, p.Addr().AsSlice()...)
                b = append(b, byte(p.Bits()))
        }
        return b
}

func parseAddress(r io.Reader) (requestID uint64, prefix netip.Prefix, _ error) {
        vr := quicvarint.NewReader(r)
        requestID, err := quicvarint.Read(vr)
        if err != nil {
                return 0, netip.Prefix{}, err
        }
        ipVersion, err := vr.ReadByte()
        if err != nil {
                return 0, netip.Prefix{}, err
        }
        var ip netip.Addr
        switch ipVersion {
        case 4:
                var ipv4 [4]byte
                if _, err := io.ReadFull(r, ipv4[:]); err != nil {
                        return 0, netip.Prefix{}, err
                }
                ip = netip.AddrFrom4(ipv4)
        case 6:
                var ipv6 [16]byte
                if _, err := io.ReadFull(r, ipv6[:]); err != nil {
                        return 0, netip.Prefix{}, err
                }
                ip = netip.AddrFrom16(ipv6)
        default:
                return 0, netip.Prefix{}, fmt.Errorf("invalid IP version: %d", ipVersion)
        }
        prefixLen, err := vr.ReadByte()
        if err != nil {
                return 0, netip.Prefix{}, err
        }
        if int(prefixLen) > ip.BitLen() {
                return 0, netip.Prefix{}, fmt.Errorf("prefix length %d exceeds IP address length (%d)", prefixLen, ip.BitLen())
        }
        prefix = netip.PrefixFrom(ip, int(prefixLen))
        if prefix != prefix.Masked() {
                return 0, netip.Prefix{}, errors.New("lower bits not covered by prefix length are not all zero")
        }
        return requestID, prefix, nil
}

// routeAdvertisementCapsule represents a ROUTE_ADVERTISEMENT capsule
type routeAdvertisementCapsule struct {
        IPAddressRanges []IPRoute
}

// IPRoute represents an IP Address Range
type IPRoute struct {
        StartIP netip.Addr
        EndIP   netip.Addr
        // IPProtocol is the Internet Protocol Number for traffic that can be sent to this range.
        // If the value is 0, all protocols are allowed.
        IPProtocol uint8
}

func (r IPRoute) len() int { return 1 + r.StartIP.BitLen()/8 + r.EndIP.BitLen()/8 + 1 }

func (r IPRoute) contains(addr netip.Addr) bool {
        return r.StartIP.Compare(addr) <= 0 && addr.Compare(r.EndIP) <= 0
}

// Prefixes returns the prefixes that this IP address range covers.
// Note that depending on the start and end addresses,
// this conversion can result in a large number of prefixes.
func (r IPRoute) Prefixes() []netip.Prefix { return rangeToPrefixes(r.StartIP, r.EndIP) }

// validateRouteAdvertisement checks the endpoints, ordering, and overlap rules in RFC 9484, Section 4.7.3.
// https://www.rfc-editor.org/rfc/rfc9484.html#section-4.7.3
func validateRouteAdvertisement(routes []IPRoute) error {
        var numIPv4, numIPv4Protocol0, numIPv6Protocol0 int
        for _, route := range routes {
                if !route.StartIP.IsValid() || !route.EndIP.IsValid() {
                        return errors.New("invalid route: IP addresses must be valid")
                }
                if route.StartIP.Is4() != route.EndIP.Is4() {
                        return errors.New("invalid route: IP addresses must have the same address family")
                }
                if route.StartIP.Zone() != "" || route.EndIP.Zone() != "" {
                        return errors.New("invalid route: IP addresses must not have zones")
                }
                if route.StartIP.Compare(route.EndIP) > 0 {
                        return fmt.Errorf("invalid route: start IP %s is greater than end IP %s", route.StartIP, route.EndIP)
                }
                if route.StartIP.Is4() {
                        numIPv4++
                        if route.IPProtocol == 0 {
                                numIPv4Protocol0++
                        }
                } else if route.IPProtocol == 0 {
                        numIPv6Protocol0++
                }
        }

        // The RFC defines three rules for A preceding B. Checking adjacent pairs suffices.
        for i := 1; i < len(routes); i++ {
                a := routes[i-1]
                b := routes[i]
                // 1. "The IP Version of A MUST be less than or equal to the IP Version of B."
                if a.StartIP.BitLen() > b.StartIP.BitLen() {
                        return errors.New("route IP versions must be in increasing order")
                }
                if a.StartIP.BitLen() != b.StartIP.BitLen() {
                        continue
                }
                // 2. Within one IP version, protocol(A) <= protocol(B).
                if a.IPProtocol > b.IPProtocol {
                        return errors.New("route IP protocols must be in increasing order")
                }
                // 3. With matching versions and protocols, end(A) < start(B).
                if a.IPProtocol == b.IPProtocol && a.EndIP.Compare(b.StartIP) >= 0 {
                        return errors.New("route address ranges must be disjoint and in increasing order")
                }
        }

        // Rule 3 only compares routes with equal IPProtocol values.
        // For example, advertising 192.0.2.0-192.0.2.255 with both IPProtocol 0 and 6 passes that check.
        // The RFC forbids this overlap too: IPProtocol == 0 allows all protocols, including 6.
        // Rule 1 groups routes by address family; rule 2 puts protocol 0 first in each family.
        ipv4 := routes[:numIPv4]
        if err := validateNonOverlappingRoutes(ipv4[:numIPv4Protocol0], ipv4[numIPv4Protocol0:]); err != nil {
                return err
        }
        ipv6 := routes[numIPv4:]
        return validateNonOverlappingRoutes(ipv6[:numIPv6Protocol0], ipv6[numIPv6Protocol0:])
}

// validateNonOverlappingRoutes checks routes allowing all protocols against routes
// allowing a specific protocol. Both slices must belong to the same address family
// and already be sorted by protocol and address.
func validateNonOverlappingRoutes(all, specific []IPRoute) error {
        var j int // cursor into all
        for i, b := range specific {
                if i > 0 && specific[i-1].IPProtocol != b.IPProtocol {
                        j = 0 // address ordering restarts for each protocol
                }
                for j < len(all) {
                        a := all[j]
                        if a.EndIP.Compare(b.StartIP) < 0 {
                                // a ends before b. Skip a for the rest of this protocol
                                j++
                                continue
                        }
                        if b.EndIP.Compare(a.StartIP) < 0 {
                                // b ends before a and every later route allowing all protocols
                                break
                        }
                        // neither range ends before the other starts: they overlap
                        return errors.New("route overlaps a route for all IP protocols")
                }
        }
        return nil
}

func parseRouteAdvertisementCapsule(r http3.CapsuleReader) (*routeAdvertisementCapsule, error) {
        var ranges []IPRoute
        for r.Remaining() > 0 {
                if len(ranges) >= maxRoutesPerCapsule {
                        return nil, fmt.Errorf("ROUTE_ADVERTISEMENT capsule contains too many routes (maximum %d)", maxRoutesPerCapsule)
                }
                ipRange, err := parseIPAddressRange(r)
                if err != nil {
                        return nil, err
                }
                ranges = append(ranges, ipRange)
        }
        if err := validateRouteAdvertisement(ranges); err != nil {
                return nil, err
        }
        return &routeAdvertisementCapsule{IPAddressRanges: ranges}, nil
}

func (c *routeAdvertisementCapsule) append(b []byte) []byte {
        var totalLen int
        for _, ipRange := range c.IPAddressRanges {
                totalLen += ipRange.len()
        }

        b = quicvarint.Append(b, uint64(capsuleTypeRouteAdvertisement))
        b = quicvarint.Append(b, uint64(totalLen))

        for _, ipRange := range c.IPAddressRanges {
                if ipRange.StartIP.Is4() {
                        b = append(b, 4)
                } else {
                        b = append(b, 6)
                }
                b = append(b, ipRange.StartIP.AsSlice()...)
                b = append(b, ipRange.EndIP.AsSlice()...)
                b = append(b, ipRange.IPProtocol)
        }
        return b
}

func parseIPAddressRange(r io.Reader) (IPRoute, error) {
        var ipVersion uint8
        if err := binary.Read(r, binary.LittleEndian, &ipVersion); err != nil {
                return IPRoute{}, err
        }

        var startIP, endIP netip.Addr
        switch ipVersion {
        case 4:
                var start, end [4]byte
                if _, err := io.ReadFull(r, start[:]); err != nil {
                        return IPRoute{}, err
                }
                if _, err := io.ReadFull(r, end[:]); err != nil {
                        return IPRoute{}, err
                }
                startIP = netip.AddrFrom4(start)
                endIP = netip.AddrFrom4(end)
        case 6:
                var start, end [16]byte
                if _, err := io.ReadFull(r, start[:]); err != nil {
                        return IPRoute{}, err
                }
                if _, err := io.ReadFull(r, end[:]); err != nil {
                        return IPRoute{}, err
                }
                startIP = netip.AddrFrom16(start)
                endIP = netip.AddrFrom16(end)
        default:
                return IPRoute{}, fmt.Errorf("invalid IP version: %d", ipVersion)
        }

        var ipProtocol uint8
        if err := binary.Read(r, binary.LittleEndian, &ipProtocol); err != nil {
                return IPRoute{}, err
        }
        return IPRoute{
                StartIP:    startIP,
                EndIP:      endIP,
                IPProtocol: ipProtocol,
        }, nil
}

// dnsAssignCapsule represents a DNS_ASSIGN capsule defined by
// draft-ietf-masque-connect-ip-dns-06.
type dnsAssignCapsule struct {
        DNSConfigurations []DNSConfiguration
}

// This limits the wire size, not memory use. Parsing can amplify memory use by
// roughly 50x, and the limit is chosen accordingly.
const maxDNSAssignCapsuleSize = 32 << 10

func parseCounted[T any](r http3.CapsuleReader, parse func(http3.CapsuleReader) (T, error)) ([]T, error) {
        count, err := quicvarint.Read(r)
        if err != nil {
                return nil, err
        }
        var values []T
        for range count {
                value, err := parse(r)
                if err != nil {
                        return nil, err
                }
                values = append(values, value)
        }
        return values, nil
}

func parseDNSAssignCapsule(r http3.CapsuleReader) (*dnsAssignCapsule, error) {
        if r.Remaining() > maxDNSAssignCapsuleSize {
                return nil, fmt.Errorf("DNS_ASSIGN capsule too large: %d bytes (maximum %d)", r.Remaining(), maxDNSAssignCapsuleSize)
        }
        var configurations []DNSConfiguration
        for r.Remaining() > 0 {
                nameservers, err := parseCounted(r, parseDNSNameserver)
                if err != nil {
                        return nil, err
                }
                internalDomains, err := parseCounted(r, parseDomain)
                if err != nil {
                        return nil, err
                }
                searchDomains, err := parseCounted(r, parseDomain)
                if err != nil {
                        return nil, err
                }
                configuration := DNSConfiguration{
                        Nameservers:     nameservers,
                        InternalDomains: internalDomains,
                        SearchDomains:   searchDomains,
                }
                if err := configuration.validate(); err != nil {
                        return nil, fmt.Errorf("invalid DNS configuration: %w", err)
                }
                configurations = append(configurations, configuration)
        }
        return &dnsAssignCapsule{DNSConfigurations: configurations}, nil
}

func parseDNSNameserver(r http3.CapsuleReader) (DNSNameserver, error) {
        var priorityBytes [2]byte
        if _, err := io.ReadFull(r, priorityBytes[:]); err != nil {
                return DNSNameserver{}, err
        }
        servicePriority := binary.BigEndian.Uint16(priorityBytes[:])
        ipv4Count, err := quicvarint.Read(r)
        if err != nil {
                return DNSNameserver{}, err
        }
        var ipv4Addresses []netip.Addr
        for range ipv4Count {
                var addr [4]byte
                if _, err := io.ReadFull(r, addr[:]); err != nil {
                        return DNSNameserver{}, err
                }
                ipv4Addresses = append(ipv4Addresses, netip.AddrFrom4(addr))
        }

        ipv6Count, err := quicvarint.Read(r)
        if err != nil {
                return DNSNameserver{}, err
        }
        var ipv6Addresses []netip.Addr
        for range ipv6Count {
                var addr [16]byte
                if _, err := io.ReadFull(r, addr[:]); err != nil {
                        return DNSNameserver{}, err
                }
                ipv6Addresses = append(ipv6Addresses, netip.AddrFrom16(addr))
        }

        authenticationDomainName, err := parseDomain(r)
        if err != nil {
                return DNSNameserver{}, err
        }
        paramsLen, err := quicvarint.Read(r)
        if err != nil {
                return DNSNameserver{}, err
        }
        if paramsLen > maxServiceParametersLen {
                return DNSNameserver{}, fmt.Errorf("service parameters too long: %d bytes", paramsLen)
        }
        if paramsLen > uint64(r.Remaining()) {
                return DNSNameserver{}, io.ErrUnexpectedEOF
        }
        var serviceParameters map[dnsmessage.SVCParamKey][]byte
        if paramsLen > 0 {
                b := make([]byte, int(paramsLen))
                if _, err := io.ReadFull(r, b); err != nil {
                        return DNSNameserver{}, err
                }
                serviceParameters = make(map[dnsmessage.SVCParamKey][]byte)
                var previousKey dnsmessage.SVCParamKey
                for len(b) > 0 {
                        if len(b) < 4 {
                                return DNSNameserver{}, fmt.Errorf("invalid service parameter header: %w", io.ErrUnexpectedEOF)
                        }
                        key := dnsmessage.SVCParamKey(binary.BigEndian.Uint16(b))
                        valueLen := int(binary.BigEndian.Uint16(b[2:]))
                        b = b[4:]
                        if valueLen > len(b) {
                                return DNSNameserver{}, fmt.Errorf("invalid service parameter value: %w", io.ErrUnexpectedEOF)
                        }
                        if len(serviceParameters) > 0 && key <= previousKey {
                                return DNSNameserver{}, errors.New("service parameter keys must be in strictly increasing order")
                        }
                        // The second index sets capacity so append on the value can't overwrite later parameters.
                        serviceParameters[key] = b[:valueLen:valueLen]
                        previousKey = key
                        b = b[valueLen:]
                }
        }
        return DNSNameserver{
                ServicePriority:          servicePriority,
                IPv4Addresses:            ipv4Addresses,
                IPv6Addresses:            ipv6Addresses,
                AuthenticationDomainName: authenticationDomainName,
                ServiceParameters:        serviceParameters,
        }, nil
}

func parseDomain(r http3.CapsuleReader) (string, error) {
        l, err := quicvarint.Read(r)
        if err != nil {
                return "", err
        }
        if l > maxDomainNameLen {
                return "", fmt.Errorf("domain name too long: %d bytes", l)
        }
        if l == 0 {
                return "", nil
        }
        b := make([]byte, int(l))
        if _, err := io.ReadFull(r, b); err != nil {
                return "", err
        }
        return string(b), nil
}

func (c *dnsAssignCapsule) append(b []byte) []byte {
        payload := make([]byte, 0, 256)
        for _, cfg := range c.DNSConfigurations {
                payload = quicvarint.Append(payload, uint64(len(cfg.Nameservers)))
                for _, ns := range cfg.Nameservers {
                        payload = binary.BigEndian.AppendUint16(payload, ns.ServicePriority)
                        payload = quicvarint.Append(payload, uint64(len(ns.IPv4Addresses)))
                        for _, addr := range ns.IPv4Addresses {
                                payload = append(payload, addr.AsSlice()...)
                        }
                        payload = quicvarint.Append(payload, uint64(len(ns.IPv6Addresses)))
                        for _, addr := range ns.IPv6Addresses {
                                payload = append(payload, addr.AsSlice()...)
                        }
                        payload = appendDomain(payload, ns.AuthenticationDomainName)
                        keys := slices.Sorted(maps.Keys(ns.ServiceParameters))
                        var l int
                        for _, key := range keys {
                                l += 4 + len(ns.ServiceParameters[key])
                        }
                        payload = quicvarint.Append(payload, uint64(l))
                        for _, key := range keys {
                                value := ns.ServiceParameters[key]
                                payload = binary.BigEndian.AppendUint16(payload, uint16(key))
                                payload = binary.BigEndian.AppendUint16(payload, uint16(len(value)))
                                payload = append(payload, value...)
                        }
                }
                payload = quicvarint.Append(payload, uint64(len(cfg.InternalDomains)))
                for _, domain := range cfg.InternalDomains {
                        payload = appendDomain(payload, domain)
                }
                payload = quicvarint.Append(payload, uint64(len(cfg.SearchDomains)))
                for _, domain := range cfg.SearchDomains {
                        payload = appendDomain(payload, domain)
                }
        }
        b = quicvarint.Append(b, uint64(capsuleTypeDNSAssign))
        b = quicvarint.Append(b, uint64(len(payload)))
        return append(b, payload...)
}

func appendDomain(b []byte, domain string) []byte {
        b = quicvarint.Append(b, uint64(len(domain)))
        return append(b, domain...)
}

// pref64Capsule represents a PREF64 capsule defined by
// draft-ietf-masque-connect-ip-dns-06.
type pref64Capsule struct {
        Prefixes []netip.Prefix
}

const maxPREF64Prefixes = 256

func parsePREF64Capsule(r http3.CapsuleReader) (*pref64Capsule, error) {
        // each prefix consists of a 1-byte prefix length and 12 bytes of address
        if r.Remaining()%13 != 0 {
                return nil, errors.New("PREF64 capsule length is not a multiple of 13")
        }
        numPrefixes := r.Remaining() / 13
        if numPrefixes > maxPREF64Prefixes {
                return nil, fmt.Errorf("PREF64 capsule contains too many prefixes: %d (maximum %d)", numPrefixes, maxPREF64Prefixes)
        }
        prefixes := make([]netip.Prefix, 0, numPrefixes)
        for r.Remaining() > 0 {
                prefixLen, err := r.ReadByte()
                if err != nil {
                        return nil, err
                }
                switch prefixLen {
                case 32, 40, 48, 56, 64, 96:
                default:
                        return nil, fmt.Errorf("invalid NAT64 prefix length: %d", prefixLen)
                }
                var addrBytes [16]byte
                if _, err := io.ReadFull(r, addrBytes[:12]); err != nil {
                        return nil, err
                }
                prefix := netip.PrefixFrom(netip.AddrFrom16(addrBytes), int(prefixLen))
                if prefix.Addr().Is4In6() {
                        return nil, errors.New("IPv4-mapped IPv6 addresses are not valid NAT64 prefixes")
                }
                prefixes = append(prefixes, prefix)
        }
        return &pref64Capsule{Prefixes: prefixes}, nil
}

func (c *pref64Capsule) append(b []byte) []byte {
        b = quicvarint.Append(b, uint64(capsuleTypePREF64))
        b = quicvarint.Append(b, uint64(13*len(c.Prefixes)))
        for _, prefix := range c.Prefixes {
                b = append(b, byte(prefix.Bits()))
                addr := prefix.Addr().As16()
                b = append(b, addr[:12]...)
        }
        return b
}

package connectip

import (
        "context"
        "errors"
        "fmt"
        "net/http"

        "github.com/quic-go/quic-go"
        "github.com/quic-go/quic-go/http3"
)

// A ClientConn represents a connection to a single proxy server.
// Multiple proxied connections can be established over a single ClientConn.
type ClientConn struct {
        clientConn *http3.ClientConn
}

// NewClientConn creates a client connection using an existing HTTP/3 connection.
// The HTTP/3 connection and its underlying QUIC connection must have datagrams enabled.
// The caller owns the HTTP/3 connection and can also use it for ordinary HTTP requests.
// Closing a [Conn] returned by [ClientConn.Dial] does not close the HTTP/3 connection.
func NewClientConn(conn *http3.ClientConn) *ClientConn {
        return &ClientConn{clientConn: conn}
}

// Dial dials a proxied connection over the proxy connection.
func (c *ClientConn) Dial(req *Request) (*Conn, *http.Response, error) {
        return c.dial(req, nil)
}

func (c *ClientConn) dial(req *Request, closeConn func() error) (*Conn, *http.Response, error) {
        httpReq := req.httpRequest()
        if httpReq.URL == nil {
                return nil, nil, errors.New("connect-ip: request URL is nil")
        }
        if httpReq.Host == "" && httpReq.URL.Host == "" {
                return nil, nil, errors.New("connect-ip: request needs a host")
        }

        select {
        case <-httpReq.Context().Done():
                return nil, nil, context.Cause(httpReq.Context())
        case <-c.clientConn.Context().Done():
                return nil, nil, context.Cause(c.clientConn.Context())
        case <-c.clientConn.ReceivedSettings():
        }

        settings := c.clientConn.Settings()
        if !settings.EnableExtendedConnect {
                return nil, nil, errors.New("connect-ip: server didn't enable Extended CONNECT")
        }
        if !settings.EnableDatagrams {
                return nil, nil, errors.New("connect-ip: server didn't enable datagrams")
        }

        rstr, err := c.clientConn.OpenRequestStream(httpReq.Context())
        if err != nil {
                return nil, nil, fmt.Errorf("connect-ip: failed to open request stream: %w", err)
        }
        var keepStream bool
        defer func() {
                if !keepStream {
                        rstr.CancelRead(quic.StreamErrorCode(http3.ErrCodeNoError))
                        rstr.CancelWrite(quic.StreamErrorCode(http3.ErrCodeNoError))
                }
        }()
        if err := rstr.SendRequestHeader(httpReq); err != nil {
                return nil, nil, fmt.Errorf("connect-ip: failed to send request: %w", err)
        }
        // TODO: optimistically return the connection
        rsp, err := rstr.ReadResponse()
        if err != nil {
                return nil, nil, fmt.Errorf("connect-ip: failed to read response: %w", err)
        }
        if rsp.StatusCode < 200 || rsp.StatusCode > 299 {
                return nil, rsp, fmt.Errorf("connect-ip: server responded with %d", rsp.StatusCode)
        }
        keepStream = true
        return newProxiedConn(rstr, closeConn), rsp, nil
}

package connectip

import (
        "context"
        "encoding/binary"
        "errors"
        "fmt"
        "io"
        "log"
        "net"
        "net/netip"
        "slices"
        "sync"
        "time"

        "github.com/quic-go/quic-go"
        "github.com/quic-go/quic-go/http3"
        "github.com/quic-go/quic-go/quicvarint"
)

type CloseError struct {
        Remote bool
}

func (e *CloseError) Error() string        { return net.ErrClosed.Error() }
func (e *CloseError) Is(target error) bool { return target == net.ErrClosed }

type http3Stream interface {
        io.ReadWriteCloser
        ReceiveDatagram(context.Context) ([]byte, error)
        SendDatagram([]byte) error
        CancelRead(quic.StreamErrorCode)
        CancelWrite(quic.StreamErrorCode)
        SetWriteDeadline(time.Time) error
}

var (
        _ http3Stream = &http3.Stream{}
        _ http3Stream = &http3.RequestStream{}
)

// If a packet is too large to fit into a QUIC datagram,
// we send an ICMP Packet Too Big packet.
// On IPv6, the minimum MTU of a link is 1280 bytes.
const minMTU = 1280

// Bound pending capsules in both directions. Queues only grow when the peer
// falls behind reading the stream or the application falls behind receiving updates.
const maxQueuedCapsules = 128

type streamWrite struct {
        Data []byte
        Fin  bool
}

// Conn is a connection that proxies IP packets over HTTP/3.
type Conn struct {
        str         http3Stream
        closeConn   func() error
        writeNotify chan struct{}
        writeDone   chan error

        assignedAddressUpdates  chan []AssignedAddress
        addressRequests         chan *addressRequestCapsule
        availableRouteUpdates   chan []IPRoute
        dnsConfigurationUpdates chan []DNSConfiguration
        pref64Updates           chan []netip.Prefix

        mu                   sync.Mutex
        queuedWrites         []streamWrite
        peerAddresses        []netip.Prefix // IP prefixes that we assigned to the peer
        localRoutes          []IPRoute      // IP routes that we advertised to the peer
        assignedAddresses    []netip.Prefix // IP prefixes that the peer assigned to us
        peerRoutes           []IPRoute      // IP routes that the peer advertised to us
        lastAddressRequestID AddressRequestID

        closeChan chan struct{}
        closeErr  error
}

func newProxiedConn(str http3Stream, closeConn func() error) *Conn {
        c := &Conn{
                str:                     str,
                closeConn:               closeConn,
                writeNotify:             make(chan struct{}, 1),
                writeDone:               make(chan error, 1),
                assignedAddressUpdates:  make(chan []AssignedAddress, maxQueuedCapsules),
                addressRequests:         make(chan *addressRequestCapsule, maxQueuedCapsules),
                availableRouteUpdates:   make(chan []IPRoute, 1),
                dnsConfigurationUpdates: make(chan []DNSConfiguration, 1),
                pref64Updates:           make(chan []netip.Prefix, 1),
                closeChan:               make(chan struct{}),
        }
        go func() {
                err := c.readFromStream()
                if err != nil {
                        log.Printf("handling stream failed: %v", err)
                }
                c.mu.Lock()
                if c.closeErr == nil {
                        c.closeErr = &CloseError{Remote: true}
                        close(c.closeChan)
                        if err != nil {
                                // Abort without consuming the remaining capsule payload.
                                c.str.CancelRead(quic.StreamErrorCode(http3.ErrCodeExcessiveLoad))
                                c.str.CancelWrite(quic.StreamErrorCode(http3.ErrCodeExcessiveLoad))
                                close(c.writeNotify)
                        } else {
                                _ = c.queueWrite(streamWrite{Fin: true})
                        }
                }
                c.mu.Unlock()
        }()
        go func() {
                err := c.writeToStream()
                if err != nil {
                        log.Printf("writing to stream failed: %v", err)
                        c.mu.Lock()
                        if c.closeErr == nil {
                                c.closeErr = &CloseError{Remote: true}
                                close(c.closeChan)
                                c.str.CancelRead(quic.StreamErrorCode(http3.ErrCodeExcessiveLoad))
                                c.str.CancelWrite(quic.StreamErrorCode(http3.ErrCodeExcessiveLoad))
                        } else {
                                // A write can time out while graceful shutdown is pending.
                                c.str.CancelWrite(quic.StreamErrorCode(http3.ErrCodeNoError))
                        }
                        c.mu.Unlock()
                }
                c.writeDone <- err
                close(c.writeDone)
        }()
        return c
}

// AdvertiseRoute schedules an advertisement of the available routes to the peer.
// It returns once the advertisement has been queued.
func (c *Conn) AdvertiseRoute(routes []IPRoute) error {
        if err := validateRouteAdvertisement(routes); err != nil {
                return err
        }

        c.mu.Lock()
        if c.closeErr != nil {
                err := c.closeErr
                c.mu.Unlock()
                return err
        }
        routes = slices.Clone(routes)
        err := c.queueWrite(streamWrite{Data: (&routeAdvertisementCapsule{IPAddressRanges: routes}).append(nil)})
        if err == nil {
                c.localRoutes = routes
        }
        c.mu.Unlock()
        if err != nil {
                _ = c.Close()
                return err
        }
        return nil
}

// RequestAddresses requests address prefixes from the peer and returns once the
// request is queued. It allocates a unique, nonzero ID for each prefix in input order.
// The IDs have no semantic meaning and can be used to correlate assignments
// returned by [Conn.ReceiveAddressAssignment] with the requested prefixes.
//
// Prefixes must be valid, with all bits outside the prefix set to zero.
// An unspecified address (0.0.0.0 or ::) requests any address of that family,
// with the prefix length indicating the preferred size.
func (c *Conn) RequestAddresses(prefixes []netip.Prefix) ([]AddressRequestID, error) {
        if len(prefixes) == 0 {
                return nil, errors.New("connect-ip: address request must contain at least one prefix")
        }
        for i, p := range prefixes {
                if !p.IsValid() || p != p.Masked() {
                        return nil, fmt.Errorf("connect-ip: invalid requested prefix %d: %s", i, p)
                }
        }

        c.mu.Lock()
        if c.closeErr != nil {
                err := c.closeErr
                c.mu.Unlock()
                return nil, err
        }
        ids := make([]AddressRequestID, len(prefixes))
        for i := range ids {
                ids[i] = c.lastAddressRequestID + AddressRequestID(i) + 1
        }
        capsule := &addressRequestCapsule{RequestIDs: ids, Prefixes: prefixes}
        err := c.queueWrite(streamWrite{Data: capsule.append(nil)})
        if err == nil {
                c.lastAddressRequestID = ids[len(ids)-1]
        }
        c.mu.Unlock()
        if err != nil {
                _ = c.Close()
                return nil, err
        }
        return ids, nil
}

// ReceiveAddressAssignment waits for the next complete address assignment,
// Each assignment replaces the preceding one; request IDs only provide optional correlation.
// Call this method in a loop from one goroutine.
func (c *Conn) ReceiveAddressAssignment(ctx context.Context) ([]AssignedAddress, error) {
        select {
        case <-ctx.Done():
                return nil, ctx.Err()
        case assignment := <-c.assignedAddressUpdates:
                return assignment, nil
        case <-c.closeChan:
                // Deliver queued responses before reporting closure.
                select {
                case assignment := <-c.assignedAddressUpdates:
                        return assignment, nil
                default:
                        return nil, c.closeErr
                }
        }
}

// ReceiveAddressRequest waits for the next address request from the peer.
// This method should be called in a loop, and each request should be answered with [AddressRequest.Respond].
func (c *Conn) ReceiveAddressRequest(ctx context.Context) (*AddressRequest, error) {
        var requested *addressRequestCapsule
        select {
        case <-ctx.Done():
                return nil, ctx.Err()
        case requested = <-c.addressRequests:
        case <-c.closeChan:
                select {
                case requested = <-c.addressRequests:
                default:
                        return nil, c.closeErr
                }
        }
        return newAddressRequest(c, requested), nil
}

// AssignAddresses schedules an assignment of address prefixes to the peer.
// A nil or empty slice removes all assigned addresses.
func (c *Conn) AssignAddresses(prefixes []netip.Prefix) error {
        capsule := &addressAssignCapsule{
                AssignedAddresses: make([]AssignedAddress, len(prefixes)),
        }
        for i, p := range prefixes {
                capsule.AssignedAddresses[i] = AssignedAddress{IPPrefix: p}
        }
        return c.sendAddressAssignment(capsule)
}

func (c *Conn) sendAddressAssignment(capsule *addressAssignCapsule) error {
        for i, addr := range capsule.AssignedAddresses {
                p := addr.IPPrefix
                if !p.IsValid() || p != p.Masked() {
                        return fmt.Errorf("connect-ip: invalid assigned prefix %d: %s", i, p)
                }
        }

        c.mu.Lock()
        if c.closeErr != nil {
                err := c.closeErr
                c.mu.Unlock()
                return err
        }
        if err := c.queueWrite(streamWrite{Data: capsule.append(nil)}); err != nil {
                c.mu.Unlock()
                _ = c.Close()
                return err
        }

        // Keep an empty assignment distinct from not having sent one.
        prefixes := make([]netip.Prefix, 0, len(capsule.AssignedAddresses))
        for _, assigned := range capsule.AssignedAddresses {
                if !assigned.Rejected() {
                        prefixes = append(prefixes, assigned.IPPrefix)
                }
        }
        c.peerAddresses = prefixes
        c.mu.Unlock()
        return nil
}

// SendDNSConfiguration schedules a DNS configuration update to the peer.
// The update supersedes the DNS configuration previously sent on this connection.
//
// To avoid leaking DNS traffic outside the tunnel, the application is responsible
// for advertising the corresponding routes before calling this method. See
// [Section 5 of draft-ietf-masque-connect-ip-dns-06].
//
// [Section 5 of draft-ietf-masque-connect-ip-dns-06]: https://datatracker.ietf.org/doc/html/draft-ietf-masque-connect-ip-dns-06#section-5
func (c *Conn) SendDNSConfiguration(configurations []DNSConfiguration) error {
        for _, config := range configurations {
                if err := config.validate(); err != nil {
                        return fmt.Errorf("invalid DNS configuration: %w", err)
                }
        }
        return c.sendCapsule((&dnsAssignCapsule{DNSConfigurations: configurations}).append(nil))
}

// ReceiveDNSConfiguration waits for the next DNS configuration update from the peer.
// Each update supersedes the preceding one.
func (c *Conn) ReceiveDNSConfiguration(ctx context.Context) ([]DNSConfiguration, error) {
        select {
        case <-ctx.Done():
                return nil, ctx.Err()
        case <-c.closeChan:
                return nil, c.closeErr
        case configurations := <-c.dnsConfigurationUpdates:
                return configurations, nil
        }
}

// SendPREF64Configuration schedules an update of the NAT64 prefixes to use for
// IPv6/IPv4 address synthesis. It returns once the update has been queued. An
// empty slice clears the previously sent configuration.
func (c *Conn) SendPREF64Configuration(prefixes []netip.Prefix) error {
        for i, prefix := range prefixes {
                if !prefix.IsValid() || !prefix.Addr().Is6() || prefix.Addr().Is4In6() {
                        return fmt.Errorf("invalid NAT64 prefix %d: not an IPv6 prefix", i)
                }
                switch prefix.Bits() {
                case 32, 40, 48, 56, 64, 96:
                default:
                        return fmt.Errorf("invalid NAT64 prefix %d: invalid prefix length %d", i, prefix.Bits())
                }
        }
        return c.sendCapsule((&pref64Capsule{Prefixes: prefixes}).append(nil))
}

// ReceivePREF64Configuration waits for the next NAT64 prefix update from the
// peer. An empty slice means that NAT64 prefixes are not available.
func (c *Conn) ReceivePREF64Configuration(ctx context.Context) ([]netip.Prefix, error) {
        select {
        case <-ctx.Done():
                return nil, ctx.Err()
        case <-c.closeChan:
                return nil, c.closeErr
        case prefixes := <-c.pref64Updates:
                return prefixes, nil
        }
}

func (c *Conn) sendCapsule(capsuleData []byte) error {
        c.mu.Lock()
        if c.closeErr != nil {
                err := c.closeErr
                c.mu.Unlock()
                return err
        }
        err := c.queueWrite(streamWrite{Data: capsuleData})
        c.mu.Unlock()
        if err != nil {
                _ = c.Close()
                return err
        }
        return nil
}

func (c *Conn) queueWrite(w streamWrite) error {
        if w.Fin {
                // Interrupt pending capsule writes so shutdown cannot stall.
                _ = c.str.SetWriteDeadline(time.Now())
        } else if len(c.queuedWrites) >= maxQueuedCapsules {
                c.closeErr = &CloseError{Remote: false}
                close(c.closeChan)
                c.str.CancelRead(quic.StreamErrorCode(http3.ErrCodeExcessiveLoad))
                c.str.CancelWrite(quic.StreamErrorCode(http3.ErrCodeExcessiveLoad))
                close(c.writeNotify)
                return errors.New("connect-ip: capsule queue full")
        }

        c.queuedWrites = append(c.queuedWrites, w)

        select {
        case c.writeNotify <- struct{}{}:
        default:
        }
        return nil
}

func queueLatest[T any](ch chan T, value T) {
        for {
                select {
                case ch <- value:
                        return
                case <-ch:
                }
        }
}

// Routes returns the routes that the peer currently advertised.
// Note that at any point during the connection, the peer can change the advertised routes.
// It is therefore recommended to call this function in a loop.
func (c *Conn) Routes(ctx context.Context) ([]IPRoute, error) {
        select {
        case <-ctx.Done():
                return nil, ctx.Err()
        case <-c.closeChan:
                return nil, c.closeErr
        case routes := <-c.availableRouteUpdates:
                return routes, nil
        }
}

func (c *Conn) readFromStream() error {
        p := http3.NewCapsuleParser(c.str)
        for {
                t, cr, err := p.Next()
                if errors.Is(err, io.EOF) {
                        return nil
                }
                if err != nil {
                        return err
                }
                capsule, err := parseCapsule(t, cr)
                if err != nil {
                        return err
                }
                switch capsule := capsule.(type) {
                case *addressAssignCapsule:
                        prefixes := make([]netip.Prefix, 0, len(capsule.AssignedAddresses))
                        for _, assigned := range capsule.AssignedAddresses {
                                if !assigned.Rejected() {
                                        prefixes = append(prefixes, assigned.IPPrefix)
                                }
                        }
                        c.mu.Lock()
                        c.assignedAddresses = prefixes
                        c.mu.Unlock()
                        select {
                        case c.assignedAddressUpdates <- capsule.AssignedAddresses:
                        default:
                                return errors.New("connect-ip: address assignment queue full")
                        }
                case *addressRequestCapsule:
                        select {
                        case c.addressRequests <- capsule:
                        default:
                                return errors.New("connect-ip: address request queue full")
                        }
                case *routeAdvertisementCapsule:
                        // Store the routes first, so that WritePacket can send to them once Routes returns them.
                        c.mu.Lock()
                        c.peerRoutes = slices.Clone(capsule.IPAddressRanges)
                        c.mu.Unlock()
                        queueLatest(c.availableRouteUpdates, capsule.IPAddressRanges)
                case *dnsAssignCapsule:
                        queueLatest(c.dnsConfigurationUpdates, capsule.DNSConfigurations)
                case *pref64Capsule:
                        queueLatest(c.pref64Updates, capsule.Prefixes)
                }
        }
}

func (c *Conn) writeToStream() error {
        for range c.writeNotify {
                for {
                        c.mu.Lock()
                        if len(c.queuedWrites) == 0 {
                                c.mu.Unlock()
                                break
                        }
                        w := c.queuedWrites[0]
                        c.queuedWrites[0] = streamWrite{}
                        c.queuedWrites = c.queuedWrites[1:]
                        c.mu.Unlock()

                        if w.Fin {
                                return c.str.Close()
                        }
                        if _, err := c.str.Write(w.Data); err != nil {
                                return err
                        }
                }
        }
        return c.closeErr
}

func (c *Conn) ReadPacket(b []byte) (n int, err error) {
start:
        data, err := c.str.ReceiveDatagram(context.Background())
        if err != nil {
                select {
                case <-c.closeChan:
                        return 0, c.closeErr
                default:
                        return 0, err
                }
        }
        contextID, n, err := quicvarint.Parse(data)
        if err != nil {
                // TODO: close connection
                return 0, fmt.Errorf("connect-ip: malformed datagram: %w", err)
        }
        if contextID != 0 {
                // Drop this datagram. We currently only support proxying of IP payloads.
                goto start
        }
        if err := c.handleIncomingProxiedPacket(data[n:]); err != nil {
                log.Printf("dropping proxied packet: %s", err)
                goto start
        }
        return copy(b, data[n:]), nil
}

func (c *Conn) handleIncomingProxiedPacket(data []byte) error {
        src, dst, ipProto, err := parseIPHeader(data)
        if err != nil {
                return err
        }

        c.mu.Lock()
        assignedAddresses := c.assignedAddresses
        localRoutes := c.localRoutes
        peerAddresses := c.peerAddresses
        peerRoutes := c.peerRoutes
        c.mu.Unlock()

        // We don't necessarily assign any addresses to the peer.
        // For example, in the Remote Access VPN use case (RFC 9484, section 8.1),
        // the client accepts incoming traffic from all IPs.
        if !isAllowedSource(src, peerAddresses, nil) {
                // ICMP errors can originate from any router on the path (see Section 7.2.1 of RFC 9484).
                // We accept them if we could have sent the packet that caused the error.
                invoking, ok := icmpInvokingPacket(data, ipProto)
                if !ok {
                        // TODO: send ICMP
                        return fmt.Errorf("connect-ip: datagram source address not allowed: %s", src)
                }
                invokingSrc, invokingDst, invokingProto, err := parseTruncatedIPHeader(invoking)
                if err != nil {
                        return errors.New("connect-ip: malformed ICMP error message")
                }
                // ICMP errors are always sent to the source address of the packet that caused the error.
                if invokingSrc != dst {
                        return fmt.Errorf("connect-ip: ICMP error sent to %s for a packet from %s", dst, invokingSrc)
                }
                if !isAllowedSource(invokingSrc, assignedAddresses, localRoutes) ||
                        !isAllowedDestination(invokingDst, invokingProto, peerAddresses, peerRoutes) {
                        return fmt.Errorf(
                                "connect-ip: ICMP error for a packet that couldn't have been sent through the tunnel: %s -> %s (protocol: %d)",
                                invokingSrc, invokingDst, invokingProto,
                        )
                }
        }

        // The destination IP address is valid if it
        // 1. is within one of the ranges assigned to us, or
        // 2. is within one of the ranges that we advertised to the peer.
        if !isAllowedDestination(dst, ipProto, assignedAddresses, localRoutes) {
                // TODO: send ICMP
                return fmt.Errorf("connect-ip: datagram destination address / protocol not allowed: %s (protocol: %d)", dst, ipProto)
        }
        return nil
}

// WritePacket writes an IP packet to the stream.
// The packet is dropped unless its destination was assigned to the peer or matches one of the peer's routes.
// If the peer assigned addresses to us, the source address needs to be one of these addresses,
// or be covered by one of the routes that we advertised.
// If sending the packet fails, it might return an ICMP packet.
// It is the caller's responsibility to send the ICMP packet to the sender.
func (c *Conn) WritePacket(b []byte) (icmp []byte, err error) {
        // The stream's send side is closed asynchronously,
        // so it might still accept datagrams after the connection was closed.
        select {
        case <-c.closeChan:
                return nil, c.closeErr
        default:
        }

        data, err := c.composeDatagram(b)
        if err != nil {
                log.Printf("dropping proxied packet (%d bytes) that can't be proxied: %s", len(b), err)
                return nil, nil
        }
        if err := c.str.SendDatagram(data); err != nil {
                if _, ok := errors.AsType[*quic.DatagramTooLargeError](err); ok {
                        icmpPacket, err := composeICMPTooLargePacket(b, minMTU)
                        if err != nil {
                                log.Printf("failed to compose ICMP too large packet: %s", err)
                        }
                        return icmpPacket, nil
                }
                select {
                case <-c.closeChan:
                        return nil, c.closeErr
                default:
                        return nil, err
                }
        }
        return nil, nil
}

func (c *Conn) composeDatagram(b []byte) ([]byte, error) {
        src, dst, ipProto, err := parseIPHeader(b)
        if err != nil {
                return nil, err
        }

        c.mu.Lock()
        assignedAddresses := c.assignedAddresses
        localRoutes := c.localRoutes
        peerAddresses := c.peerAddresses
        peerRoutes := c.peerRoutes
        c.mu.Unlock()

        // The source IP address is valid if
        // 1. the peer didn't assign any addresses to us, or
        // 2. it is within one of the ranges assigned to us, or
        // 3. it is within one of the ranges that we advertised to the peer (independent of the IP protocol),
        //    since we're forwarding packets from these networks.
        //
        // ICMP errors can originate from any router on the path, so we don't check their source address.
        // Instead, we check that we would have accepted the packet that caused the error from the peer.
        // This prevents proxies from forwarding ICMP errors to the wrong client
        // when clients share an address (see Section 11 of RFC 9484).
        if invoking, ok := icmpInvokingPacket(b, ipProto); ok {
                invokingSrc, invokingDst, invokingProto, err := parseTruncatedIPHeader(invoking)
                if err != nil {
                        return nil, errors.New("connect-ip: malformed ICMP error message")
                }
                // ICMP errors are always sent to the source address of the packet that caused the error.
                if invokingSrc != dst {
                        return nil, fmt.Errorf("connect-ip: ICMP error sent to %s for a packet from %s", dst, invokingSrc)
                }
                if !isAllowedSource(invokingSrc, peerAddresses, nil) ||
                        !isAllowedDestination(invokingDst, invokingProto, assignedAddresses, localRoutes) {
                        return nil, fmt.Errorf(
                                "connect-ip: ICMP error for a packet that couldn't have been received through the tunnel: %s -> %s (protocol: %d)",
                                invokingSrc, invokingDst, invokingProto,
                        )
                }
        } else if !isAllowedSource(src, assignedAddresses, localRoutes) {
                return nil, fmt.Errorf("connect-ip: source address not allowed: %s", src)
        }
        // The destination IP address is valid if it
        // 1. is within one of the ranges that we assigned to the peer, or
        // 2. is within one of the ranges that the peer advertised to us.
        if !isAllowedDestination(dst, ipProto, peerAddresses, peerRoutes) {
                return nil, fmt.Errorf("connect-ip: destination address / protocol not allowed: %s (protocol: %d)", dst, ipProto)
        }

        switch ipVersion(b) {
        case 4:
                ttl := b[8]
                if ttl <= 1 {
                        return nil, fmt.Errorf("connect-ip: datagram TTL too small: %d", ttl)
                }
                b[8]-- // decrement TTL
                // recalculate the checksum
                hdrLen := int(b[0]&0x0f) * 4
                binary.BigEndian.PutUint16(b[10:12], calculateIPv4Checksum(b[:hdrLen]))
        case 6:
                hopLimit := b[7]
                if hopLimit <= 1 {
                        return nil, fmt.Errorf("connect-ip: datagram Hop Limit too small: %d", hopLimit)
                }
                b[7]-- // Decrement Hop Limit
        }
        data := make([]byte, 0, len(contextIDZero)+len(b))
        data = append(data, contextIDZero...)
        data = append(data, b...)
        return data, nil
}

func (c *Conn) Close() error {
        c.mu.Lock()
        if c.closeErr == nil {
                c.closeErr = &CloseError{Remote: false}
                close(c.closeChan)
                _ = c.queueWrite(streamWrite{Fin: true})
        }
        closeConn := c.closeConn
        c.closeConn = nil
        c.mu.Unlock()
        err := <-c.writeDone
        c.str.CancelRead(quic.StreamErrorCode(http3.ErrCodeNoError))
        if closeConn != nil {
                return errors.Join(err, closeConn())
        }
        return err
}

package connectip

import (
        "errors"
        "fmt"
        "net/netip"
        "strings"
        "unicode/utf8"

        "golang.org/x/net/dns/dnsmessage"
        "golang.org/x/net/idna"
)

const (
        // RFC 1035, Section 2.3.4 limits a domain name to 255 octets.
        maxDomainNameLen = 255
        // SvcParams use DNS RDATA encoding, whose RDLENGTH field is 16 bits
        // (RFC 1035, Section 4.1.3).
        maxServiceParametersLen = 1<<16 - 1
)

// DNSNameserver describes the Nameserver structure defined by
// draft-ietf-masque-connect-ip-dns-06.
type DNSNameserver struct {
        ServicePriority uint16
        IPv4Addresses   []netip.Addr
        // IPv6Addresses must not contain scoped addressing zones.
        // Zone identifiers are local to an endpoint and are not part of the wire format.
        IPv6Addresses []netip.Addr
        // AuthenticationDomainName is the fully qualified domain name of the
        // nameserver, without the trailing root dot and using IDNA A-labels. It may be
        // empty only when the nameserver supports unencrypted DNS exclusively.
        AuthenticationDomainName string
        // ServiceParameters is the set of SVCB parameters that apply to this nameserver.
        // Each map value is the wire-format SvcParamValue for its key.
        ServiceParameters map[dnsmessage.SVCParamKey][]byte
}

// DNSConfiguration describes the DNS Configuration structure defined by
// draft-ietf-masque-connect-ip-dns-06.
type DNSConfiguration struct {
        // Nameservers must not be empty.
        Nameservers []DNSNameserver
        // InternalDomains contains fully qualified domain names, without the trailing
        // root dot and using IDNA A-labels. An empty string represents the DNS root.
        InternalDomains []string
        // SearchDomains contains fully qualified domain names, without the trailing
        // root dot and using IDNA A-labels.
        SearchDomains []string
}

var dnsNameProfile = idna.New(
        idna.MapForLookup(),
        idna.StrictDomainName(false),
        idna.BidiRule(),
        idna.VerifyDNSLength(true),
)

func validateDomainName(name string, allowEmpty bool) error {
        if name == "" {
                if allowEmpty {
                        return nil
                }
                return errors.New("must not be empty")
        }
        for i := range len(name) {
                if name[i] >= utf8.RuneSelf {
                        return errors.New("must use IDNA A-label form")
                }
        }
        if strings.HasSuffix(name, ".") {
                return errors.New("must omit the trailing root dot")
        }
        ascii, err := dnsNameProfile.ToASCII(name)
        if err != nil {
                return fmt.Errorf("must be a valid IDNA A-label: %w", err)
        }
        if !strings.EqualFold(ascii, name) {
                return errors.New("must use IDNA A-label form")
        }
        return nil
}

func (c DNSConfiguration) validate() error {
        if len(c.Nameservers) == 0 {
                return errors.New("nameserver count must not be zero")
        }
        for _, ns := range c.Nameservers {
                if ns.ServicePriority == 0 {
                        return errors.New("service priority must not be zero")
                }
                var serviceParametersLen int
                var hasALPN, hasNoDefaultALPN bool
                for k, v := range ns.ServiceParameters {
                        if len(v) > maxServiceParametersLen {
                                return errors.New("service parameter value too long")
                        }
                        serviceParametersLen += 4 + len(v)
                        switch k {
                        case dnsmessage.SVCParamALPN:
                                hasALPN = true
                        case dnsmessage.SVCParamNoDefaultALPN:
                                hasNoDefaultALPN = true
                        case dnsmessage.SVCParamIPv4Hint, dnsmessage.SVCParamIPv6Hint:
                                return fmt.Errorf("service parameter %s is not allowed", k)
                        }
                }
                if serviceParametersLen > maxServiceParametersLen {
                        return errors.New("service parameters too long")
                }
                if err := validateDomainName(ns.AuthenticationDomainName, true); err != nil {
                        return fmt.Errorf("invalid authentication domain name: %w", err)
                }
                if ns.AuthenticationDomainName == "" && (hasALPN || hasNoDefaultALPN) {
                        return errors.New("ALPN service parameters require an authentication domain name")
                }
                if !hasNoDefaultALPN && len(ns.IPv4Addresses)+len(ns.IPv6Addresses) == 0 {
                        return errors.New("nameserver must have an address when no-default-alpn is omitted")
                }
                for _, addr := range ns.IPv4Addresses {
                        if !addr.Is4() {
                                return fmt.Errorf("non-IPv4 address in IPv4 address list: %s", addr)
                        }
                }
                for _, addr := range ns.IPv6Addresses {
                        if !addr.Is6() || addr.Is4In6() {
                                return fmt.Errorf("non-IPv6 address in IPv6 address list: %s", addr)
                        }
                        if addr.Zone() != "" {
                                return fmt.Errorf("IPv6 address with zone in IPv6 address list: %s", addr)
                        }
                }
        }
        for _, domain := range c.InternalDomains {
                if err := validateDomainName(domain, true); err != nil {
                        return fmt.Errorf("invalid internal domain name: %w", err)
                }
        }
        for _, domain := range c.SearchDomains {
                if err := validateDomainName(domain, false); err != nil {
                        return fmt.Errorf("invalid search domain name: %w", err)
                }
        }
        return nil
}

package connectip

import (
        "encoding/binary"
        "errors"
        "fmt"

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

// icmpInvokingPacket returns the invoking packet of an ICMP error message, i.e. the packet that caused the error.
// It returns false if b isn't an ICMP error message. b must have been accepted by parseIPHeader.
//
// The invoking packet is truncated if the ICMP error message would otherwise exceed 576 bytes
// for IPv4 (Section 4.3.2.3 of RFC 1812), or the minimum MTU for IPv6 (Section 2.4 of RFC 4443).
// It is empty if the ICMP header itself is truncated.
func icmpInvokingPacket(b []byte, ipProto uint8) (invoking []byte, ok bool) {
        var msg []byte
        switch {
        case ipVersion(b) == 4 && ipProto == ipProtoICMP:
                // fragments at non-zero fragment offsets don't contain the ICMP header
                if binary.BigEndian.Uint16(b[6:8])&0x1fff != 0 {
                        return nil, false
                }
                msg = b[int(b[0]&0x0f)*4:]
                if len(msg) == 0 {
                        return nil, false
                }
                // Unlike ICMPv6, ICMPv4 doesn't reserve a range of message types for errors.
                // Redirect and Source Quench messages also quote the packet that caused them,
                // but we treat them like any other ICMP message: Redirects only apply to hosts
                // on the same link as the router, and Source Quench is deprecated (see RFC 6633).
                switch ipv4.ICMPType(msg[0]) {
                case ipv4.ICMPTypeDestinationUnreachable, ipv4.ICMPTypeTimeExceeded, ipv4.ICMPTypeParameterProblem:
                default:
                        return nil, false
                }
        case ipVersion(b) == 6 && ipProto == ipProtoICMPv6:
                _, msg, _ = ipv6UpperLayerProtocol(b)
                // ICMPv6 error messages have types from 0 to 127, see Section 2.1 of RFC 4443
                if len(msg) == 0 || msg[0] >= 128 {
                        return nil, false
                }
        default:
                return nil, false
        }
        // The invoking packet follows the 8 byte ICMP header.
        // For ICMPv6, this also applies to unknown error message types, see Section 2.4 and Appendix A of RFC 4443.
        if len(msg) < 8 {
                return nil, true
        }
        return msg[8:], true
}

func composeICMPTooLargePacket(b []byte, mtu int) ([]byte, error) {
        if len(b) == 0 {
                return nil, errors.New("connect-ip: empty packet")
        }

        var icmpMessage *icmp.Message
        var psh []byte
        switch v := ipVersion(b); v {
        case 4:
                if len(b) < ipv4.HeaderLen {
                        return nil, errors.New("connect-ip: IPv4 packet too short")
                }
                // RFC 1812, section 4.3.2.3 recommends quoting as much of the original
                // packet as possible without exceeding 576 bytes, including the IPv4 and ICMP headers.
                const maxDataLen = 576 - ipv4.HeaderLen - 8
                icmpMessage = &icmp.Message{
                        Type: ipv4.ICMPTypeDestinationUnreachable,
                        Code: 4, // Fragmentation Needed and Don't Fragment was Set
                        Body: &icmp.PacketTooBig{
                                MTU:  mtu,
                                Data: b[:min(len(b), maxDataLen)],
                        },
                }
        case 6:
                if len(b) < ipv6.HeaderLen {
                        return nil, errors.New("connect-ip: IPv6 packet too short")
                }
                icmpMessage = &icmp.Message{
                        Type: ipv6.ICMPTypePacketTooBig,
                        Body: &icmp.PacketTooBig{
                                MTU:  mtu,
                                Data: b[:min(len(b), 1232)],
                        },
                }
                psh = icmp.IPv6PseudoHeader(b[24:40], b[8:24])
        default:
                return nil, fmt.Errorf("connect-ip: unknown IP version: %d", v)
        }

        icmp, err := icmpMessage.Marshal(psh)
        if err != nil {
                return nil, fmt.Errorf("connect-ip: failed to marshal ICMP message: %w", err)
        }

        if ipVersion(b) == 4 {
                var header [ipv4.HeaderLen]byte
                header[0] = 4<<4 | ipv4.HeaderLen>>2 // Version and IHL
                ipLen := ipv4.HeaderLen + len(icmp)
                binary.BigEndian.PutUint16(header[2:4], uint16(ipLen)) // Total Length
                header[8] = 64                                         // TTL
                header[9] = 1                                          // Protocol (ICMP)
                copy(header[12:16], b[16:20])                          // Source IP from original packet
                copy(header[16:20], b[12:16])                          // Dest IP from original packet (swapped)
                binary.BigEndian.PutUint16(header[10:12], calculateIPv4Checksum(header[:]))
                return append(header[:], icmp...), nil
        }

        var header [ipv6.HeaderLen]byte
        header[0] = 6 << 4                                         // Version 6
        binary.BigEndian.PutUint16(header[4:6], uint16(len(icmp))) // Payload Length
        header[6] = 58                                             // Next Header (ICMPv6)
        header[7] = 64                                             // Hop Limit
        copy(header[8:24], b[24:40])                               // Source IP from original packet
        copy(header[24:40], b[8:24])                               // Dest IP from original packet (swapped)
        return append(header[:], icmp...), nil
}

package connectip

import (
        "errors"
        "fmt"
        "net/netip"
        "slices"

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

const (
        ipProtoICMP   = 1
        ipProtoICMPv6 = 58
)

// parseIPHeader returns the source and destination addresses and the upper-layer protocol of an IP packet.
// It validates the length fields, as well as the header checksum for IPv4.
func parseIPHeader(b []byte) (src, dst netip.Addr, ipProto uint8, err error) {
        if len(b) == 0 {
                return netip.Addr{}, netip.Addr{}, 0, errors.New("connect-ip: empty packet")
        }
        switch ipVersion(b) {
        case 4:
                if len(b) < ipv4.HeaderLen {
                        return netip.Addr{}, netip.Addr{}, 0, errors.New("connect-ip: IPv4 packet too short")
                }
                if err := validateIPv4Checksum(b); err != nil {
                        return netip.Addr{}, netip.Addr{}, 0, err
                }
                if err := validateIPv4TotalLength(b); err != nil {
                        return netip.Addr{}, netip.Addr{}, 0, err
                }
        case 6:
                if len(b) < ipv6.HeaderLen {
                        return netip.Addr{}, netip.Addr{}, 0, errors.New("connect-ip: IPv6 packet too short")
                }
                if err := validateIPv6PayloadLength(b); err != nil {
                        return netip.Addr{}, netip.Addr{}, 0, err
                }
        }
        return parseTruncatedIPHeader(b)
}

// parseTruncatedIPHeader is like parseIPHeader, but for packets that might be truncated,
// for example the packet quoted in an ICMP error message.
// It doesn't validate the length fields or the IPv4 header checksum.
func parseTruncatedIPHeader(b []byte) (src, dst netip.Addr, ipProto uint8, err error) {
        if len(b) == 0 {
                return netip.Addr{}, netip.Addr{}, 0, errors.New("connect-ip: empty packet")
        }
        switch v := ipVersion(b); v {
        default:
                return netip.Addr{}, netip.Addr{}, 0, fmt.Errorf("connect-ip: unknown IP versions: %d", v)
        case 4:
                if len(b) < ipv4.HeaderLen {
                        return netip.Addr{}, netip.Addr{}, 0, errors.New("connect-ip: IPv4 packet too short")
                }
                src = netip.AddrFrom4([4]byte(b[12:16]))
                dst = netip.AddrFrom4([4]byte(b[16:20]))
                ipProto = b[9]
        case 6:
                if len(b) < ipv6.HeaderLen {
                        return netip.Addr{}, netip.Addr{}, 0, errors.New("connect-ip: IPv6 packet too short")
                }
                src = netip.AddrFrom16([16]byte(b[8:24]))
                dst = netip.AddrFrom16([16]byte(b[24:40]))
                if ipProto, _, err = ipv6UpperLayerProtocol(b); err != nil {
                        return netip.Addr{}, netip.Addr{}, 0, err
                }
        }
        return src, dst, ipProto, nil
}

// isAllowedSource checks if src is within one of the prefixes, or within one of the routes.
// If prefixes is nil, no addresses were assigned, and all source addresses are allowed.
func isAllowedSource(src netip.Addr, prefixes []netip.Prefix, routes []IPRoute) bool {
        if prefixes == nil {
                return true
        }
        return slices.ContainsFunc(prefixes, func(p netip.Prefix) bool { return p.Contains(src) }) ||
                slices.ContainsFunc(routes, func(r IPRoute) bool { return r.contains(src) })
}

// isAllowedDestination checks if dst is within one of the prefixes,
// or within one of the routes that allow the IP protocol.
func isAllowedDestination(dst netip.Addr, ipProto uint8, prefixes []netip.Prefix, routes []IPRoute) bool {
        for _, p := range prefixes {
                if p.Contains(dst) {
                        return true
                }
        }
        // ICMP is always allowed
        isICMP := (dst.Is4() && ipProto == ipProtoICMP) || (dst.Is6() && ipProto == ipProtoICMPv6)
        // Fragments at non-zero fragment offsets don't contain the upper-layer header.
        // For a stateless firewall, there's no way to reassemble the original packet.
        // In that case, we'd only drop the first fragment, but forward the rest.
        // Without the first fragment, the receiver will not be able to reassemble the original packet.
        // See Section 4 of RFC 7112.
        isFragment := dst.Is6() && ipProto == ipProtoFragment
        for _, r := range routes {
                if !r.contains(dst) {
                        continue
                }
                if isICMP || isFragment || r.IPProtocol == 0 || r.IPProtocol == ipProto {
                        return true
                }
        }
        return false
}

func ipVersion(b []byte) uint8 { return b[0] >> 4 }

package connectip

import "net/netip"

// rangeToPrefixes converts an IP range defined by start and end addresses
// into a slice of CIDR prefixes that exactly cover the range.
// It assumes that the start and end addresses are of the same IP version,
// and that the start address is less than or equal to the end address.
func rangeToPrefixes(start, end netip.Addr) []netip.Prefix {
        var prefixes []netip.Prefix
        for current := start; current.Compare(end) <= 0; {
                // find the largest prefix that fits within our remaining range
                prefix := findLargestPrefix(current, end)
                prefixes = append(prefixes, prefix)

                // move to the next IP address after this prefix
                lastIP := lastIPInPrefix(prefix)
                if lastIP.Compare(end) >= 0 {
                        break
                }
                current = lastIP.Next()
        }
        return prefixes
}

// findLargestPrefix finds the largest prefix starting at 'start' that doesn't exceed 'end'
func findLargestPrefix(start, end netip.Addr) netip.Prefix {
        if start == end {
                return netip.PrefixFrom(start, start.BitLen())
        }

        // Start with the smallest possible prefix (/32 for IPv4, /128 for IPv6),
        // and try progressively larger prefixes until we find one that exceeds our range.
        var prefixLen int
        for prefixLen = start.BitLen(); prefixLen > 0; prefixLen-- {
                prefix := netip.PrefixFrom(start, prefixLen-1) // Try one bit larger
                if lastIPInPrefix(prefix).Compare(end) > 0 || !isAligned(start, prefixLen-1) {
                        break
                }
        }
        return netip.PrefixFrom(start, prefixLen)
}

// lastIPInPrefix returns the last IP address in a prefix
func lastIPInPrefix(prefix netip.Prefix) netip.Addr {
        addr := prefix.Addr()
        bits := addr.As16()

        hostBits := addr.BitLen() - prefix.Bits()

        // Set all host bits to 1
        for i := len(bits) - 1; i >= 0 && hostBits > 0; i-- {
                bitsInThisByte := min(8, hostBits)
                mask := byte((1 << bitsInThisByte) - 1)
                bits[i] |= mask
                hostBits -= bitsInThisByte
        }

        if addr.Is4() {
                return netip.AddrFrom4([4]byte(bits[12:16]))
        }
        return netip.AddrFrom16(bits)
}

// isAligned checks if an IP address is aligned to the given prefix length
func isAligned(addr netip.Addr, prefixLen int) bool {
        bits := addr.As16()

        hostBits := addr.BitLen() - prefixLen
        for i := len(bits) - 1; i >= 0 && hostBits > 0; i-- {
                bitsInThisByte := min(8, hostBits)
                mask := byte((1 << bitsInThisByte) - 1)
                if bits[i]&mask != 0 {
                        return false
                }
                hostBits -= bitsInThisByte
        }
        return true
}

package connectip

import (
        "encoding/binary"
        "errors"
        "fmt"

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

// calculateIPv4Checksum calculates the checksum of an IPv4 header, including options.
func calculateIPv4Checksum(header []byte) uint16 {
        // add every 16-bit word in the header, skipping the checksum field (bytes 10 and 11)
        var sum uint32
        for i := 0; i < len(header); i += 2 {
                if i == 10 {
                        continue // skip checksum field
                }
                sum += uint32(binary.BigEndian.Uint16(header[i : i+2]))
        }
        for (sum >> 16) > 0 {
                sum = (sum & 0xffff) + (sum >> 16)
        }
        return ^uint16(sum)
}

// validateIPv4Checksum checks the header checksum of an IPv4 packet, including options.
// The packet must be at least ipv4.HeaderLen bytes long.
func validateIPv4Checksum(b []byte) error {
        hdrLen := int(b[0]&0x0f) * 4
        if hdrLen < ipv4.HeaderLen || hdrLen > len(b) {
                return fmt.Errorf("connect-ip: invalid IPv4 header length: %d", hdrLen)
        }
        var sum uint32
        for i := 0; i < hdrLen; i += 2 {
                sum += uint32(binary.BigEndian.Uint16(b[i : i+2]))
        }
        for (sum >> 16) > 0 {
                sum = (sum & 0xffff) + (sum >> 16)
        }
        if sum != 0xffff {
                return errors.New("connect-ip: invalid IPv4 header checksum")
        }
        return nil
}

// validateIPv4TotalLength checks that the Total Length field matches the packet size.
// An HTTP Datagram contains exactly one IP packet, with no trailing bytes (Section 6 of RFC 9484).
// Since validateIPv4Checksum checks that the header fits into the packet,
// this also ensures that the Total Length covers the header (Section 5.2.2 of RFC 1812).
func validateIPv4TotalLength(b []byte) error {
        if l := int(binary.BigEndian.Uint16(b[2:4])); l != len(b) {
                return fmt.Errorf("connect-ip: IPv4 total length (%d) doesn't match packet size (%d)", l, len(b))
        }
        return nil
}

package connectip

import (
        "encoding/binary"
        "errors"
        "fmt"

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

// IPv6 extension headers listed in Section 4 of RFC 8200, except for ESP.
const (
        ipProtoHopByHop = 0
        ipProtoRouting  = 43
        ipProtoFragment = 44
        ipProtoAH       = 51
        ipProtoDestOpts = 60
)

var errTruncatedExtensionHeader = errors.New("connect-ip: malformed datagram: truncated IPv6 extension header")

// validateIPv6PayloadLength checks that the Payload Length field matches the packet size.
// An HTTP Datagram contains exactly one IP packet, with no trailing bytes (Section 6 of RFC 9484).
func validateIPv6PayloadLength(b []byte) error {
        // A Payload Length of zero is also used by jumbograms (RFC 2675), which are rejected here.
        // Jumbograms are larger than 65535 bytes, and therefore never fit into a QUIC DATAGRAM frame.
        if l := int(binary.BigEndian.Uint16(b[4:6])); ipv6.HeaderLen+l != len(b) {
                return fmt.Errorf("connect-ip: IPv6 payload length (%d) doesn't match packet size (%d)", l, len(b)-ipv6.HeaderLen)
        }
        return nil
}

// ipv6UpperLayerProtocol walks the extension header chain of an IPv6 packet
// and returns the upper-layer protocol number, as well as the payload starting with the upper-layer header.
//
// Only the extension headers listed in Section 4 of RFC 8200 are walked, except for ESP,
// which encrypts everything following it. To keep the implementation simple, all other
// extension headers are returned as the upper-layer protocol.
//
// Fragments at non-zero fragment offsets don't contain the upper-layer header.
// For these, ipProtoFragment and a nil payload are returned.
func ipv6UpperLayerProtocol(b []byte) (uint8, []byte, error) {
        next := b[6]
        pos := ipv6.HeaderLen
        for {
                switch next {
                case ipProtoHopByHop, ipProtoRouting, ipProtoFragment, ipProtoAH, ipProtoDestOpts:
                default:
                        return next, b[pos:], nil
                }
                // the Hop-by-Hop Options header must immediately follow the IPv6 header
                if next == ipProtoHopByHop && pos != ipv6.HeaderLen {
                        return 0, nil, errors.New("connect-ip: malformed datagram: misplaced Hop-by-Hop Options header")
                }
                // each of these headers is at least 8 bytes long
                if len(b) < pos+8 {
                        return 0, nil, errTruncatedExtensionHeader
                }
                var hdrLen int
                switch next {
                case ipProtoFragment:
                        // the 13 bit fragment offset is non-zero for all but the first fragment
                        if binary.BigEndian.Uint16(b[pos+2:pos+4])>>3 != 0 {
                                return ipProtoFragment, nil, nil
                        }
                        hdrLen = 8
                case ipProtoAH:
                        // the length is encoded in 4-byte units, minus 2 (see Section 2.2 of RFC 4302)
                        hdrLen = (int(b[pos+1]) + 2) * 4
                default:
                        // the length is encoded in 8-byte units, not including the first 8 bytes
                        hdrLen = (int(b[pos+1]) + 1) * 8
                }
                if len(b) < pos+hdrLen {
                        return 0, nil, errTruncatedExtensionHeader
                }
                next = b[pos]
                pos += hdrLen
        }
}

package connectip

import (
        "net/http"

        "github.com/quic-go/quic-go/http3"
        "github.com/quic-go/quic-go/quicvarint"
)

var contextIDZero = quicvarint.Append([]byte{}, 0)

type Proxy struct{}

func (s *Proxy) Proxy(w http.ResponseWriter, _ *ProxyRequest) (*Conn, error) {
        w.Header().Set(http3.CapsuleProtocolHeader, capsuleProtocolHeaderValue)
        w.WriteHeader(http.StatusOK)

        str := w.(http3.HTTPStreamer).HTTPStream()
        return newProxiedConn(str, nil), nil
}

package connectip

import (
        "context"
        "errors"
        "fmt"
        "net/http"
        "net/url"
        "reflect"

        "github.com/dunglas/httpsfv"
        "github.com/quic-go/quic-go/http3"
        "github.com/yosida95/uritemplate/v3"
)

const requestProtocol = "connect-ip"

var capsuleProtocolHeaderValue string

func init() {
        v, err := httpsfv.Marshal(httpsfv.NewItem(true))
        if err != nil {
                panic(fmt.Sprintf("failed to marshal capsule protocol header value: %v", err))
        }
        capsuleProtocolHeaderValue = v
}

// Request is a CONNECT-IP request created by NewRequest.
// The zero value is not valid.
type Request struct {
        req *http.Request
}

// NewRequest creates a CONNECT-IP request.
func NewRequest(ctx context.Context, proxyTemplate *uritemplate.Template) (*Request, error) {
        // connect-ip-go currently does not support IP flow forwarding,
        // see https://github.com/quic-go/connect-ip-go/issues/31.
        if len(proxyTemplate.Varnames()) > 0 {
                return nil, errors.New("connect-ip: IP flow forwarding not supported")
        }
        req, err := http.NewRequestWithContext(ctx, http.MethodConnect, proxyTemplate.Raw(), nil)
        if err != nil {
                return nil, fmt.Errorf("connect-ip: failed to create request: %w", err)
        }
        req.Proto = requestProtocol
        req.Host = req.URL.Host
        req.Header.Set(http3.CapsuleProtocolHeader, capsuleProtocolHeaderValue)
        return &Request{req: req}, nil
}

// Header returns the HTTP header fields sent with the CONNECT-IP request.
// Callers may add custom headers before dialing.
func (r *Request) Header() http.Header { return r.req.Header }

func (r *Request) httpRequest() *http.Request { return r.req }

// ProxyRequest is the parsed CONNECT-IP request returned from ParseProxyRequest.
// It currently doesn't have any fields, since connect-ip-go doesn't support IP flow forwarding.
type ProxyRequest struct{}

// ProxyRequestParseError is returned from ParseProxyRequest if parsing the CONNECT-IP request fails.
// It is recommended that the request is rejected with the corresponding HTTP status code.
type ProxyRequestParseError struct {
        HTTPStatus int
        Err        error
}

func (e *ProxyRequestParseError) Error() string { return e.Err.Error() }
func (e *ProxyRequestParseError) Unwrap() error { return e.Err }

// ParseProxyRequest parses a CONNECT-IP request.
// The template is the URI template that clients will use to configure this proxy.
func ParseProxyRequest(r *http.Request, template *uritemplate.Template) (*ProxyRequest, error) {
        if len(template.Varnames()) > 0 {
                return nil, errors.New("connect-ip-go currently does not support IP flow forwarding")
        }

        u, err := url.Parse(template.Raw())
        if err != nil {
                return nil, &ProxyRequestParseError{
                        HTTPStatus: http.StatusInternalServerError,
                        Err:        fmt.Errorf("failed to parse template: %w", err),
                }
        }
        if r.Method != http.MethodConnect {
                return nil, &ProxyRequestParseError{
                        HTTPStatus: http.StatusMethodNotAllowed,
                        Err:        fmt.Errorf("expected CONNECT request, got %s", r.Method),
                }
        }
        if r.Proto != requestProtocol {
                return nil, &ProxyRequestParseError{
                        HTTPStatus: http.StatusNotImplemented,
                        Err:        fmt.Errorf("unexpected protocol: %s", r.Proto),
                }
        }
        if r.Host != u.Host {
                return nil, &ProxyRequestParseError{
                        HTTPStatus: http.StatusBadRequest,
                        Err:        fmt.Errorf("host in :authority (%s) does not match template host (%s)", r.Host, u.Host),
                }
        }
        capsuleHeaderValues, ok := r.Header[http3.CapsuleProtocolHeader]
        if !ok {
                return nil, &ProxyRequestParseError{
                        HTTPStatus: http.StatusBadRequest,
                        Err:        fmt.Errorf("missing Capsule-Protocol header"),
                }
        }
        item, err := httpsfv.UnmarshalItem(capsuleHeaderValues)
        if err != nil {
                return nil, &ProxyRequestParseError{
                        HTTPStatus: http.StatusBadRequest,
                        Err:        fmt.Errorf("invalid capsule header value: %s", capsuleHeaderValues),
                }
        }
        if v, ok := item.Value.(bool); !ok {
                return nil, &ProxyRequestParseError{
                        HTTPStatus: http.StatusBadRequest,
                        Err:        fmt.Errorf("incorrect capsule header value type: %s", reflect.TypeOf(item.Value)),
                }
        } else if !v {
                return nil, &ProxyRequestParseError{
                        HTTPStatus: http.StatusBadRequest,
                        Err:        fmt.Errorf("incorrect capsule header value: %t", item.Value),
                }
        }

        return &ProxyRequest{}, nil
}

package connectip

import (
        "context"
        "crypto/tls"
        "errors"
        "fmt"
        "net"
        "net/http"

        "github.com/quic-go/quic-go"
        "github.com/quic-go/quic-go/http3"
)

// defaultInitialPacketSize allows tunneling QUIC connections, which require an MTU of at least 1200 bytes.
const defaultInitialPacketSize = 1350

// A Transport establishes proxied connections to multiple proxy servers.
type Transport struct {
        // TLSClientConfig is the TLS client config used when dialing the QUIC connection to the proxy.
        // It must set the "h3" ALPN.
        TLSClientConfig *tls.Config

        // QUICConfig is the QUIC config used when dialing the QUIC connection.
        QUICConfig *quic.Config

        // DialAddr dials the QUIC connection to the proxy.
        // If unset, quic.DialAddr is used.
        DialAddr func(ctx context.Context, addr string, tlsConf *tls.Config, quicConf *quic.Config) (*quic.Conn, error)
}

// Dial opens a QUIC connection to the proxy and then dials a proxied connection.
// Closing the returned Conn also closes the QUIC connection to the proxy.
// To establish multiple proxied connections over one proxy connection, use [Transport.NewClientConn].
func (t *Transport) Dial(req *Request) (*Conn, *http.Response, error) {
        httpReq := req.httpRequest()
        if httpReq.URL == nil || httpReq.URL.Host == "" {
                return nil, nil, errors.New("connect-ip: request URL needs a host")
        }

        quicConf := t.QUICConfig
        if quicConf == nil {
                quicConf = &quic.Config{EnableDatagrams: true, InitialPacketSize: defaultInitialPacketSize}
        }
        if !quicConf.EnableDatagrams {
                return nil, nil, errors.New("connect-ip: QUICConfig needs to enable datagrams")
        }
        tlsConf := t.TLSClientConfig
        if tlsConf == nil {
                tlsConf = &tls.Config{NextProtos: []string{http3.NextProtoH3}}
        }
        dial := t.DialAddr
        if dial == nil {
                dial = quic.DialAddr
        }
        addr := httpReq.URL.Host
        if httpReq.URL.Port() == "" {
                addr = net.JoinHostPort(httpReq.URL.Hostname(), "443")
        }
        conn, err := dial(httpReq.Context(), addr, tlsConf, quicConf)
        if err != nil {
                return nil, nil, fmt.Errorf("connect-ip: dialing QUIC connection failed: %w", err)
        }
        c, err := t.NewClientConn(conn)
        if err != nil {
                conn.CloseWithError(0, "")
                return nil, nil, err
        }
        pconn, rsp, err := c.dial(req, func() error { return conn.CloseWithError(0, "") })
        if err != nil {
                conn.CloseWithError(0, "")
                return nil, rsp, err
        }
        return pconn, rsp, nil
}

// NewClientConn creates a client connection for an already established QUIC connection.
// It returns an error if the QUIC connection didn't negotiate datagram support.
// The caller owns the QUIC connection and closes it when done.
// To reuse an existing HTTP/3 connection, use [NewClientConn].
func (t *Transport) NewClientConn(conn *quic.Conn) (*ClientConn, error) {
        datagrams := conn.ConnectionState().SupportsDatagrams
        if !datagrams.Local || !datagrams.Remote {
                return nil, errors.New("connect-ip: QUIC connection needs datagram support")
        }
        tr := &http3.Transport{EnableDatagrams: true}
        return NewClientConn(tr.NewClientConn(conn)), nil
}