// Copyright 2020 The gVisor Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package stack

import (
	"encoding/binary"
	"fmt"
	"math"
	"sync"
	"time"

	"github.com/sagernet/gvisor/pkg/atomicbitops"
	"github.com/sagernet/gvisor/pkg/tcpip"
	"github.com/sagernet/gvisor/pkg/tcpip/hash/jenkins"
	"github.com/sagernet/gvisor/pkg/tcpip/header"
	"github.com/sagernet/gvisor/pkg/tcpip/transport/tcpconntrack"
)

// Connection tracking is used to track and manipulate packets for NAT rules.
// The connection is created for a packet if it does not exist. Every
// connection contains two tuples (original and reply). The tuples are
// manipulated if there is a matching NAT rule. The packet is modified by
// looking at the tuples in each hook.
//
// Currently, only TCP tracking is supported.

// Our hash table has 16K buckets.
const numBuckets = 1 << 14

const (
	establishedTimeout   time.Duration = 5 * 24 * time.Hour
	unestablishedTimeout time.Duration = 120 * time.Second
)

// ConnTrackState represents the state of a connection.
type ConnTrackState int

const (
	// ConnTrackStateInvalid is the invalid connection tracking state.
	ConnTrackStateInvalid ConnTrackState = -1
	// ConnTrackStateEstablished represents an established connection.
	ConnTrackStateEstablished ConnTrackState = 0
	// ConnTrackStateNew represents a new connection.
	ConnTrackStateNew ConnTrackState = 2
	// ConnTrackStateEstablishedReply represents an established connection
	// in the reply direction.
	ConnTrackStateEstablishedReply ConnTrackState = 3
)

// ConnTrackDirection represents the direction of a connection.
type ConnTrackDirection uint8

const (
	// ConnTrackDirectionOriginal represents the original direction.
	ConnTrackDirectionOriginal ConnTrackDirection = 0
	// ConnTrackDirectionReply represents the reply direction.
	ConnTrackDirectionReply ConnTrackDirection = 1
)

// tuple holds a connection's identifying and manipulating data in one
// direction. It is immutable.
//
// +stateify savable
type tuple struct {
	// tupleEntry is used to build an intrusive list of tuples.
	tupleEntry

	// conn is the connection tracking entry this tuple belongs to.
	conn *conn

	// reply is true iff the tuple's direction is opposite that of the first
	// packet seen on the connection.
	reply bool

	// tupleID is set at initialization and is immutable.
	tupleID tupleID
}

// tupleID uniquely identifies a trackable connection in one direction.
//
// +stateify savable
type tupleID struct {
	srcAddr tcpip.Address
	// The source port of a packet in the original direction is overloaded with
	// the ident of an Echo Request packet.
	//
	// This also matches the behaviour of sending packets on Linux where the
	// socket's source port value is used for the source port of outgoing packets
	// for TCP/UDP and the ident field for outgoing Echo Requests on Ping sockets:
	//
	//   IPv4: https://github.com/torvalds/linux/blob/c5c17547b778975b3d83a73c8d84e8fb5ecf3ba5/net/ipv4/ping.c#L810
	//   IPv6: https://github.com/torvalds/linux/blob/c5c17547b778975b3d83a73c8d84e8fb5ecf3ba5/net/ipv6/ping.c#L133
	srcPortOrEchoRequestIdent uint16
	dstAddr                   tcpip.Address
	// The opposite of srcPortOrEchoRequestIdent; the destination port of a packet
	// in the reply direction is overloaded with the ident of an Echo Reply.
	dstPortOrEchoReplyIdent uint16
	transProto              tcpip.TransportProtocolNumber
	netProto                tcpip.NetworkProtocolNumber
}

// reply creates the reply tupleID.
func (ti tupleID) reply() tupleID {
	return tupleID{
		srcAddr:                   ti.dstAddr,
		srcPortOrEchoRequestIdent: ti.dstPortOrEchoReplyIdent,
		dstAddr:                   ti.srcAddr,
		dstPortOrEchoReplyIdent:   ti.srcPortOrEchoRequestIdent,
		transProto:                ti.transProto,
		netProto:                  ti.netProto,
	}
}

type manipType int

const (
	// manipNotPerformed indicates that NAT has not been performed.
	manipNotPerformed manipType = iota

	// manipPerformed indicates that NAT was performed.
	manipPerformed

	// manipPerformedNoop indicates that NAT was performed but it was a no-op.
	manipPerformedNoop
)

type finalizeResult uint32

const (
	// A finalizeResult must be explicitly set so we don't make use of the zero
	// value.
	_ finalizeResult = iota

	finalizeResultSuccess
	finalizeResultConflict
)

// conn is a tracked connection.
//
// +stateify savable
type conn struct {
	ct *ConnTrack

	// original is the tuple in original direction. It is immutable.
	original tuple

	// reply is the tuple in reply direction.
	reply tuple

	// TODO(b/341946753): Restore when netstack is savable.
	finalizeOnce sync.Once `state:"nosave"`
	// Holds a finalizeResult.
	finalizeResult atomicbitops.Uint32

	mu connRWMutex `state:"nosave"`
	// sourceManip indicates the source manipulation type.
	//
	// +checklocks:mu
	sourceManip manipType
	// destinationManip indicates the destination's manipulation type.
	//
	// +checklocks:mu
	destinationManip manipType

	stateMu stateConnRWMutex `state:"nosave"`
	// tcb is TCB control block. It is used to keep track of states
	// of tcp connection.
	//
	// +checklocks:stateMu
	tcb tcpconntrack.TCB
	// lastUsed is the last time the connection saw a relevant packet, and
	// is updated by each packet on the connection.
	//
	// +checklocks:stateMu
	lastUsed tcpip.MonotonicTime
	// replySeen indicates whether a packet in the reply direction has been seen.
	//
	// +checklocks:stateMu
	replySeen bool
}

// timedOut returns whether the connection timed out based on its state.
func (cn *conn) timedOut(now tcpip.MonotonicTime) bool {
	cn.stateMu.RLock()
	defer cn.stateMu.RUnlock()
	if cn.tcb.State() == tcpconntrack.ResultAlive {
		// Use the same default as Linux, which doesn't delete
		// established connections for 5(!) days.
		return now.Sub(cn.lastUsed) > establishedTimeout
	}
	// Use the same default as Linux, which lets connections in most states
	// other than established remain for <= 120 seconds.
	return now.Sub(cn.lastUsed) > unestablishedTimeout
}

// expiresIn returns the duration from now until the connection times out.
func (cn *conn) expiresIn() time.Duration {
	var timeout time.Duration
	var lastUsed tcpip.MonotonicTime
	cn.stateMu.RLock()
	state := cn.tcb.State()
	lastUsed = cn.lastUsed
	cn.stateMu.RUnlock()
	if state == tcpconntrack.ResultAlive {
		timeout = establishedTimeout
	} else {
		timeout = unestablishedTimeout
	}
	now := cn.ct.clock.NowMonotonic()
	expires := timeout - now.Sub(lastUsed)
	if expires < 0 {
		return 0
	}
	return expires
}

// update the connection tracking state.
func (cn *conn) update(pkt *PacketBuffer, reply bool) {
	cn.stateMu.Lock()
	defer cn.stateMu.Unlock()

	// Mark the connection as having been used recently so it isn't reaped.
	cn.lastUsed = cn.ct.clock.NowMonotonic()
	if reply {
		cn.replySeen = true
	}

	if pkt.TransportProtocolNumber != header.TCPProtocolNumber {
		return
	}

	tcpHeader := header.TCP(pkt.TransportHeader().Slice())

	// Update the state of tcb. tcb assumes it's always initialized on the
	// client. However, we only need to know whether the connection is
	// established or not, so the client/server distinction isn't important.
	if cn.tcb.IsEmpty() {
		cn.tcb.Init(tcpHeader, pkt.Data().Size())
		return
	}

	if reply {
		cn.tcb.UpdateStateReply(tcpHeader, pkt.Data().Size())
	} else {
		cn.tcb.UpdateStateOriginal(tcpHeader, pkt.Data().Size())
	}
}

type connTrackRNG interface {
	Uint32() uint32
}

// ConnTrack tracks all connections created for NAT rules. Most users are
// expected to only call handlePacket, insertRedirectConn, and maybeInsertNoop.
//
// ConnTrack keeps all connections in a slice of buckets, each of which holds a
// linked list of tuples. This gives us some desirable properties:
//   - Each bucket has its own lock, lessening lock contention.
//   - The slice is large enough that lists stay short (<10 elements on average).
//     Thus traversal is fast.
//   - During linked list traversal we reap expired connections. This amortizes
//     the cost of reaping them and makes reapUnused faster.
//
// Locks are ordered by their location in the buckets slice. That is, a
// goroutine that locks buckets[i] can only lock buckets[j] s.t. i < j.
//
// +stateify savable
type ConnTrack struct {
	// seed is a one-time random value initialized at stack startup
	// and is used in the calculation of hash keys for the list of buckets.
	// It is immutable.
	//
	// TODO(gvisor.dev/issue/4595): When Stack.tables becomes savable and
	// ConnTrack flows into checkpoint state, this seed must be redrawn
	// from secureRNG during restore AND the entries in buckets must be
	// rehashed under the new seed. bucket_index = jenkins.Sum32(seed) %
	// len(buckets) couples the seed value to bucket layout; redrawing the
	// seed without rehashing leaves restored entries unreachable by
	// Lookup. Persisting the pre-checkpoint seed extends the brute-force
	// window across save boundaries.
	seed uint32

	// nftIDSeed is a one-time random value initialized at stack startup
	// and is used in the calculation of tuple IDs for nftables.
	// It is immutable.
	nftIDSeed uint32

	// clock provides timing used to determine conntrack reapings.
	clock tcpip.Clock
	// TODO(b/341946753): Restore when netstack is savable.
	rng connTrackRNG `state:"nosave"`

	mu connTrackRWMutex `state:"nosave"`
	// mu protects the buckets slice, but not buckets' contents. Only take
	// the write lock if you are modifying the slice or saving for S/R.
	//
	// +checklocks:mu
	buckets []bucket
}

// +stateify savable
type bucket struct {
	mu bucketRWMutex `state:"nosave"`
	// +checklocks:mu
	tuples tupleList
}

// A netAndTransHeadersFunc returns the network and transport headers found
// in an ICMP payload. The transport layer's payload will not be returned.
//
// May panic if the packet does not hold the transport header.
type netAndTransHeadersFunc func(icmpPayload []byte, minTransHdrLen int) (netHdr header.Network, transHdrBytes []byte)

func v4NetAndTransHdr(icmpPayload []byte, minTransHdrLen int) (header.Network, []byte) {
	netHdr := header.IPv4(icmpPayload)
	// Do not use netHdr.Payload() as we might not hold the full packet
	// in the ICMP error; Payload() panics if the buffer is smaller than
	// the total length specified in the IPv4 header.
	transHdr := icmpPayload[netHdr.HeaderLength():]
	return netHdr, transHdr[:minTransHdrLen]
}

func v6NetAndTransHdr(icmpPayload []byte, minTransHdrLen int) (header.Network, []byte) {
	netHdr := header.IPv6(icmpPayload)
	// Do not use netHdr.Payload() as we might not hold the full packet
	// in the ICMP error; Payload() panics if the IP payload is smaller than
	// the payload length specified in the IPv6 header.
	transHdr := icmpPayload[header.IPv6MinimumSize:]
	return netHdr, transHdr[:minTransHdrLen]
}

func getTupleIDForRegularPacket(netHdr header.Network, netProto tcpip.NetworkProtocolNumber, transHdr header.Transport, transProto tcpip.TransportProtocolNumber) tupleID {
	return tupleID{
		srcAddr:                   netHdr.SourceAddress(),
		srcPortOrEchoRequestIdent: transHdr.SourcePort(),
		dstAddr:                   netHdr.DestinationAddress(),
		dstPortOrEchoReplyIdent:   transHdr.DestinationPort(),
		transProto:                transProto,
		netProto:                  netProto,
	}
}

func getTupleIDForPacketInICMPError(pkt *PacketBuffer, getNetAndTransHdr netAndTransHeadersFunc, netProto tcpip.NetworkProtocolNumber, netLen int, transProto tcpip.TransportProtocolNumber) (tupleID, bool) {
	if netHdr, transHdr, ok := pkt.GetEmbeddedNetAndTransHeaders(netLen, getNetAndTransHdr, transProto); ok {
		return tupleID{
			srcAddr:                   netHdr.DestinationAddress(),
			srcPortOrEchoRequestIdent: transHdr.DestinationPort(),
			dstAddr:                   netHdr.SourceAddress(),
			dstPortOrEchoReplyIdent:   transHdr.SourcePort(),
			transProto:                transProto,
			netProto:                  netProto,
		}, true
	}

	return tupleID{}, false
}

type getTupleIDDisposition int

const (
	getTupleIDNotOK getTupleIDDisposition = iota
	getTupleIDOKAndAllowNewConn
	getTupleIDOKAndDontAllowNewConn
)

func getTupleIDForEchoPacket(pkt *PacketBuffer, ident uint16, request bool) tupleID {
	netHdr := pkt.Network()
	tid := tupleID{
		srcAddr:    netHdr.SourceAddress(),
		dstAddr:    netHdr.DestinationAddress(),
		transProto: pkt.TransportProtocolNumber,
		netProto:   pkt.NetworkProtocolNumber,
	}

	if request {
		tid.srcPortOrEchoRequestIdent = ident
	} else {
		tid.dstPortOrEchoReplyIdent = ident
	}

	return tid
}

func getTupleID(pkt *PacketBuffer) (tupleID, getTupleIDDisposition) {
	switch pkt.TransportProtocolNumber {
	case header.TCPProtocolNumber:
		if transHeader := header.TCP(pkt.TransportHeader().Slice()); len(transHeader) >= header.TCPMinimumSize {
			return getTupleIDForRegularPacket(pkt.Network(), pkt.NetworkProtocolNumber, transHeader, pkt.TransportProtocolNumber), getTupleIDOKAndAllowNewConn
		}
	case header.UDPProtocolNumber:
		if transHeader := header.UDP(pkt.TransportHeader().Slice()); len(transHeader) >= header.UDPMinimumSize {
			return getTupleIDForRegularPacket(pkt.Network(), pkt.NetworkProtocolNumber, transHeader, pkt.TransportProtocolNumber), getTupleIDOKAndAllowNewConn
		}
	case header.ICMPv4ProtocolNumber:
		icmp := header.ICMPv4(pkt.TransportHeader().Slice())
		if len(icmp) < header.ICMPv4MinimumSize {
			return tupleID{}, getTupleIDNotOK
		}

		switch icmp.Type() {
		case header.ICMPv4Echo:
			return getTupleIDForEchoPacket(pkt, icmp.Ident(), true /* request */), getTupleIDOKAndAllowNewConn
		case header.ICMPv4EchoReply:
			// Do not create a new connection in response to a reply packet as only
			// the first packet of a connection should create a conntrack entry but
			// a reply is never the first packet sent for a connection.
			return getTupleIDForEchoPacket(pkt, icmp.Ident(), false /* request */), getTupleIDOKAndDontAllowNewConn
		case header.ICMPv4DstUnreachable, header.ICMPv4TimeExceeded, header.ICMPv4ParamProblem:
		default:
			// Unsupported ICMP type for NAT-ing.
			return tupleID{}, getTupleIDNotOK
		}

		h, ok := pkt.Data().PullUp(header.IPv4MinimumSize)
		if !ok {
			return tupleID{}, getTupleIDNotOK
		}

		ipv4 := header.IPv4(h)
		if ipv4.HeaderLength() > header.IPv4MinimumSize {
			// TODO(https://gvisor.dev/issue/6765): Handle IPv4 options.
			return tupleID{}, getTupleIDNotOK
		}

		if tid, ok := getTupleIDForPacketInICMPError(pkt, v4NetAndTransHdr, header.IPv4ProtocolNumber, header.IPv4MinimumSize, ipv4.TransportProtocol()); ok {
			// Do not create a new connection in response to an ICMP error.
			return tid, getTupleIDOKAndDontAllowNewConn
		}
	case header.ICMPv6ProtocolNumber:
		icmp := header.ICMPv6(pkt.TransportHeader().Slice())
		if len(icmp) < header.ICMPv6MinimumSize {
			return tupleID{}, getTupleIDNotOK
		}

		switch icmp.Type() {
		case header.ICMPv6EchoRequest:
			return getTupleIDForEchoPacket(pkt, icmp.Ident(), true /* request */), getTupleIDOKAndAllowNewConn
		case header.ICMPv6EchoReply:
			// Do not create a new connection in response to a reply packet as only
			// the first packet of a connection should create a conntrack entry but
			// a reply is never the first packet sent for a connection.
			return getTupleIDForEchoPacket(pkt, icmp.Ident(), false /* request */), getTupleIDOKAndDontAllowNewConn
		case header.ICMPv6DstUnreachable, header.ICMPv6PacketTooBig, header.ICMPv6TimeExceeded, header.ICMPv6ParamProblem:
		default:
			return tupleID{}, getTupleIDNotOK
		}

		h, ok := pkt.Data().PullUp(header.IPv6MinimumSize)
		if !ok {
			return tupleID{}, getTupleIDNotOK
		}

		// TODO(https://gvisor.dev/issue/6789): Handle extension headers.
		if tid, ok := getTupleIDForPacketInICMPError(pkt, v6NetAndTransHdr, header.IPv6ProtocolNumber, header.IPv6MinimumSize, header.IPv6(h).TransportProtocol()); ok {
			// Do not create a new connection in response to an ICMP error.
			return tid, getTupleIDOKAndDontAllowNewConn
		}
	}

	return tupleID{}, getTupleIDNotOK
}

func (ct *ConnTrack) init() {
	ct.mu.Lock()
	defer ct.mu.Unlock()
	ct.buckets = make([]bucket, numBuckets)
}

// getConnAndUpdate attempts to get a connection or creates one if no
// connection exists for the packet and packet's protocol is trackable.
//
// If the packet's protocol is trackable, the connection's state is updated to
// match the contents of the packet.
func (ct *ConnTrack) getConnAndUpdate(pkt *PacketBuffer, skipChecksumValidation bool) *tuple {
	// Get or (maybe) create a connection.
	t := func() *tuple {
		var allowNewConn bool
		tid, res := getTupleID(pkt)
		switch res {
		case getTupleIDNotOK:
			return nil
		case getTupleIDOKAndAllowNewConn:
			allowNewConn = true
		case getTupleIDOKAndDontAllowNewConn:
			allowNewConn = false
		default:
			panic(fmt.Sprintf("unhandled %[1]T = %[1]d", res))
		}

		// Just skip bad packets. They'll be rejected later by the appropriate
		// protocol package.
		switch pkt.TransportProtocolNumber {
		case header.TCPProtocolNumber:
			_, csumValid, ok := header.TCPValid(
				header.TCP(pkt.TransportHeader().Slice()),
				func() uint16 { return pkt.Data().Checksum() },
				uint16(pkt.Data().Size()),
				tid.srcAddr,
				tid.dstAddr,
				pkt.RXChecksumValidated || skipChecksumValidation)
			if !csumValid || !ok {
				return nil
			}
		case header.UDPProtocolNumber:
			lengthValid, csumValid := header.UDPValid(
				header.UDP(pkt.TransportHeader().Slice()),
				func() uint16 { return pkt.Data().Checksum() },
				uint16(pkt.Data().Size()),
				pkt.NetworkProtocolNumber,
				tid.srcAddr,
				tid.dstAddr,
				pkt.RXChecksumValidated || skipChecksumValidation)
			if !lengthValid || !csumValid {
				return nil
			}
		}

		ct.mu.RLock()
		bkt := &ct.buckets[ct.bucket(tid)]
		ct.mu.RUnlock()

		now := ct.clock.NowMonotonic()
		if t := bkt.connForTID(tid, now); t != nil {
			return t
		}

		if !allowNewConn {
			return nil
		}

		bkt.mu.Lock()
		defer bkt.mu.Unlock()

		// Make sure a connection wasn't added between when we last checked the
		// bucket and acquired the bucket's write lock.
		if t := bkt.connForTIDRLocked(tid, now); t != nil {
			return t
		}

		// This is the first packet we're seeing for the connection. Create an entry
		// for this new connection.
		conn := &conn{
			ct:       ct,
			original: tuple{tupleID: tid},
			reply:    tuple{tupleID: tid.reply(), reply: true},
			lastUsed: now,
		}
		conn.original.conn = conn
		conn.reply.conn = conn

		// For now, we only map an entry for the packet's original tuple as NAT may be
		// performed on this connection. Until the packet goes through all the hooks
		// and its final address/port is known, we cannot know what the response
		// packet's addresses/ports will look like.
		//
		// This is okay because the destination cannot send its response until it
		// receives the packet; the packet will only be received once all the hooks
		// have been performed.
		//
		// See (*conn).finalize.
		bkt.tuples.PushFront(&conn.original)
		return &conn.original
	}()
	if t != nil {
		t.conn.update(pkt, t.reply)
	}
	return t
}

// GetConnAndUpdatePkt gets the connection for the packet and also sets the packet's tuple.
func (ct *ConnTrack) GetConnAndUpdatePkt(pkt *PacketBuffer, skipChecksumValidation bool) {
	pkt.tuple = ct.getConnAndUpdate(pkt, skipChecksumValidation)
}

func (ct *ConnTrack) connForTID(tid tupleID) *tuple {
	ct.mu.RLock()
	bkt := &ct.buckets[ct.bucket(tid)]
	ct.mu.RUnlock()

	return bkt.connForTID(tid, ct.clock.NowMonotonic())
}

// ConnTrackInfo holds connection tracking information for a packet.
type ConnTrackInfo struct {
	State      ConnTrackState
	Direction  ConnTrackDirection
	SrcAddr    tcpip.Address
	DstAddr    tcpip.Address
	SrcPort    uint16
	DstPort    uint16
	NetProto   tcpip.NetworkProtocolNumber
	TransProto tcpip.TransportProtocolNumber
	Expiration time.Duration
	PseudoID   uint32
	Bytes      uint64
	Packets    uint64
}

// ConnTrackInfoOpts holds options for GetConnTrackInfo.
type ConnTrackInfoOpts struct {
	FillState      bool
	UseReplyDir    bool
	FillPseudoID   bool
	FillExpiration bool
}

// getTCPConnTrackState converts the TCB state to ConnTrackState.
func (cn *conn) getTCPConnTrackState(useReplyDir bool) ConnTrackState {
	state := ConnTrackStateInvalid
	cn.stateMu.RLock()
	tcbState := cn.tcb.State()
	cn.stateMu.RUnlock()
	switch tcbState {
	case tcpconntrack.ResultConnecting:
		state = ConnTrackStateNew

	case tcpconntrack.ResultAlive, tcpconntrack.ResultReset,
		tcpconntrack.ResultClosedByOriginator, tcpconntrack.ResultClosedByResponder:

		if useReplyDir {
			state = ConnTrackStateEstablishedReply
		} else {
			state = ConnTrackStateEstablished
		}
	case tcpconntrack.ResultDrop:
		state = ConnTrackStateInvalid
	}
	return state
}

// getConnTrackState returns the connection tracking state for the connection.
func (cn *conn) getConnTrackState(useReplyDir bool) ConnTrackState {
	state := ConnTrackStateInvalid
	// TCP connections have their own state machine in the TCB.
	if cn.original.tupleID.transProto == header.TCPProtocolNumber {
		return cn.getTCPConnTrackState(useReplyDir)
	}
	// For non-TCP connections, fill the info based on the reply.
	cn.stateMu.RLock()
	replySeen := cn.replySeen
	cn.stateMu.RUnlock()
	if useReplyDir {
		state = ConnTrackStateEstablishedReply
	} else if replySeen {
		state = ConnTrackStateEstablished
	} else {
		state = ConnTrackStateNew
	}
	return state
}

// FillConnTrackInfo fills connection tracking information for the connection.
func (cn *conn) FillConnTrackInfo(opts ConnTrackInfoOpts, info *ConnTrackInfo) bool {
	state := ConnTrackStateInvalid
	if opts.FillState {
		state = cn.getConnTrackState(opts.UseReplyDir)
	}

	dir := ConnTrackDirectionOriginal
	t := &cn.original
	if opts.UseReplyDir {
		t = &cn.reply
		dir = ConnTrackDirectionReply
	}
	tID := t.tupleID

	pID := uint32(0)
	if opts.FillPseudoID {
		// Generate a pseudo-ID similar to Linux nf_ct_get_id
		pID = tupleHash(cn.original.tupleID, cn.ct.nftIDSeed)
	}

	var expires time.Duration
	if opts.FillExpiration {
		expires = cn.expiresIn()
	}

	info.State = state
	info.Direction = dir
	info.SrcAddr = tID.srcAddr
	info.DstAddr = tID.dstAddr
	info.SrcPort = tID.srcPortOrEchoRequestIdent
	info.DstPort = tID.dstPortOrEchoReplyIdent
	info.NetProto = tID.netProto
	info.TransProto = tID.transProto
	info.Expiration = expires
	info.PseudoID = pID
	return true
}

func (bkt *bucket) connForTID(tid tupleID, now tcpip.MonotonicTime) *tuple {
	bkt.mu.RLock()
	defer bkt.mu.RUnlock()
	return bkt.connForTIDRLocked(tid, now)
}

// +checklocksread:bkt.mu
func (bkt *bucket) connForTIDRLocked(tid tupleID, now tcpip.MonotonicTime) *tuple {
	for other := bkt.tuples.Front(); other != nil; other = other.Next() {
		if tid == other.tupleID && !other.conn.timedOut(now) {
			return other
		}
	}
	return nil
}

func (ct *ConnTrack) finalize(cn *conn) finalizeResult {
	ct.mu.RLock()
	buckets := ct.buckets
	ct.mu.RUnlock()

	{
		tid := cn.reply.tupleID
		id := ct.bucketWithTableLength(tid, len(buckets))

		bkt := &buckets[id]
		bkt.mu.Lock()
		t := bkt.connForTIDRLocked(tid, ct.clock.NowMonotonic())
		if t == nil {
			bkt.tuples.PushFront(&cn.reply)
			bkt.mu.Unlock()
			return finalizeResultSuccess
		}
		bkt.mu.Unlock()

		if t.conn == cn {
			// We already have an entry for the reply tuple.
			//
			// This can occur when the source address/port is the same as the
			// destination address/port. In this scenario, tid == tid.reply().
			return finalizeResultSuccess
		}
	}

	// Another connection for the reply already exists. Remove the original and
	// let the caller know we failed.
	//
	// TODO(https://gvisor.dev/issue/6850): Investigate handling this clash
	// better.

	tid := cn.original.tupleID
	id := ct.bucketWithTableLength(tid, len(buckets))
	bkt := &buckets[id]
	bkt.mu.Lock()
	defer bkt.mu.Unlock()
	bkt.tuples.Remove(&cn.original)
	return finalizeResultConflict
}

func (cn *conn) getFinalizeResult() finalizeResult {
	return finalizeResult(cn.finalizeResult.Load())
}

// finalize attempts to finalize the connection and returns true iff the
// connection was successfully finalized.
//
// If the connection failed to finalize, the caller should drop the packet
// associated with the connection.
//
// If multiple goroutines attempt to finalize at the same time, only one
// goroutine will perform the work to finalize the connection, but all
// goroutines will block until the finalizing goroutine finishes finalizing.
func (cn *conn) finalize() bool {
	cn.finalizeOnce.Do(func() {
		cn.finalizeResult.Store(uint32(cn.ct.finalize(cn)))
	})

	switch res := cn.getFinalizeResult(); res {
	case finalizeResultSuccess:
		return true
	case finalizeResultConflict:
		return false
	default:
		panic(fmt.Sprintf("unhandled result = %d", res))
	}
}

// bucket gets the conntrack bucket for a tupleID.
// +checklocksread:ct.mu
func (ct *ConnTrack) bucket(id tupleID) int {
	return ct.bucketWithTableLength(id, len(ct.buckets))
}

func tupleHash(id tupleID, seed uint32) uint32 {
	h := jenkins.Sum32(seed)
	h.Write(id.srcAddr.AsSlice())
	h.Write(id.dstAddr.AsSlice())
	shortBuf := make([]byte, 2)
	binary.LittleEndian.PutUint16(shortBuf, id.srcPortOrEchoRequestIdent)
	h.Write([]byte(shortBuf))
	binary.LittleEndian.PutUint16(shortBuf, id.dstPortOrEchoReplyIdent)
	h.Write([]byte(shortBuf))
	binary.LittleEndian.PutUint16(shortBuf, uint16(id.transProto))
	h.Write([]byte(shortBuf))
	binary.LittleEndian.PutUint16(shortBuf, uint16(id.netProto))
	h.Write([]byte(shortBuf))
	return h.Sum32()
}

func (ct *ConnTrack) bucketWithTableLength(id tupleID, tableLength int) int {
	h := tupleHash(id, ct.seed)
	return int(h) % tableLength
}

// reapUnused deletes timed out entries from the conntrack map. The rules for
// reaping are:
//   - Each call to reapUnused traverses a fraction of the conntrack table.
//     Specifically, it traverses len(ct.buckets)/fractionPerReaping.
//   - After reaping, reapUnused decides when it should next run based on the
//     ratio of expired connections to examined connections. If the ratio is
//     greater than maxExpiredPct, it schedules the next run quickly. Otherwise it
//     slightly increases the interval between runs.
//   - maxFullTraversal caps the time it takes to traverse the entire table.
//
// reapUnused returns the next bucket that should be checked and the time after
// which it should be called again.
func (ct *ConnTrack) reapUnused(start int, prevInterval time.Duration) (int, time.Duration) {
	const fractionPerReaping = 128
	const maxExpiredPct = 50
	const maxFullTraversal = 60 * time.Second
	const minInterval = 10 * time.Millisecond
	const maxInterval = maxFullTraversal / fractionPerReaping

	now := ct.clock.NowMonotonic()
	checked := 0
	expired := 0
	var idx int
	ct.mu.RLock()
	defer ct.mu.RUnlock()
	for i := 0; i < len(ct.buckets)/fractionPerReaping; i++ {
		idx = (i + start) % len(ct.buckets)
		bkt := &ct.buckets[idx]
		bkt.mu.Lock()
		for tuple := bkt.tuples.Front(); tuple != nil; {
			// reapTupleLocked updates tuple's next pointer so we grab it here.
			nextTuple := tuple.Next()

			checked++
			if ct.reapTupleLocked(tuple, idx, bkt, now) {
				expired++
			}

			tuple = nextTuple
		}
		bkt.mu.Unlock()
	}
	// We already checked buckets[idx].
	idx++

	// If half or more of the connections are expired, the table has gotten
	// stale. Reschedule quickly.
	expiredPct := 0
	if checked != 0 {
		expiredPct = expired * 100 / checked
	}
	if expiredPct > maxExpiredPct {
		return idx, minInterval
	}
	if interval := prevInterval + minInterval; interval <= maxInterval {
		// Increment the interval between runs.
		return idx, interval
	}
	// We've hit the maximum interval.
	return idx, maxInterval
}

// reapTupleLocked tries to remove tuple and its reply from the table. It
// returns whether the tuple's connection has timed out.
//
// Precondition: ct.mu is read locked and bkt.mu is write locked.
// +checklocksread:ct.mu
// +checklocks:bkt.mu
func (ct *ConnTrack) reapTupleLocked(reapingTuple *tuple, bktID int, bkt *bucket, now tcpip.MonotonicTime) bool {
	if !reapingTuple.conn.timedOut(now) {
		return false
	}

	var otherTuple *tuple
	if reapingTuple.reply {
		otherTuple = &reapingTuple.conn.original
	} else {
		otherTuple = &reapingTuple.conn.reply
	}

	otherTupleBktID := ct.bucket(otherTuple.tupleID)
	replyTupleInserted := reapingTuple.conn.getFinalizeResult() == finalizeResultSuccess

	// To maintain lock order, we can only reap both tuples if the tuple for the
	// other direction appears later in the table.
	if bktID > otherTupleBktID && replyTupleInserted {
		return true
	}

	bkt.tuples.Remove(reapingTuple)

	if !replyTupleInserted {
		// The other tuple is the reply which has not yet been inserted.
		return true
	}

	// Reap the other connection.
	if bktID == otherTupleBktID {
		// Don't re-lock if both tuples are in the same bucket.
		bkt.tuples.Remove(otherTuple)
	} else {
		otherTupleBkt := &ct.buckets[otherTupleBktID]
		otherTupleBkt.mu.NestedLock(bucketLockOthertuple)
		otherTupleBkt.tuples.Remove(otherTuple)
		otherTupleBkt.mu.NestedUnlock(bucketLockOthertuple)
	}

	return true
}

func (ct *ConnTrack) originalDst(epID TransportEndpointID, netProto tcpip.NetworkProtocolNumber, transProto tcpip.TransportProtocolNumber) (tcpip.Address, uint16, tcpip.Error) {
	// Lookup the connection. The reply's original destination
	// describes the original address.
	tid := tupleID{
		srcAddr:                   epID.LocalAddress,
		srcPortOrEchoRequestIdent: epID.LocalPort,
		dstAddr:                   epID.RemoteAddress,
		dstPortOrEchoReplyIdent:   epID.RemotePort,
		transProto:                transProto,
		netProto:                  netProto,
	}
	t := ct.connForTID(tid)
	if t == nil {
		// Not a tracked connection.
		return tcpip.Address{}, 0, &tcpip.ErrNotConnected{}
	}

	t.conn.mu.RLock()
	defer t.conn.mu.RUnlock()
	if t.conn.destinationManip == manipNotPerformed {
		// Unmanipulated destination.
		return tcpip.Address{}, 0, &tcpip.ErrInvalidOptionValue{}
	}

	id := t.conn.original.tupleID
	return id.dstAddr, id.dstPortOrEchoReplyIdent, nil
}

// NewConnTrack creates and initializes a  new ConnTrack object.
func NewConnTrack(clock tcpip.Clock, rng connTrackRNG, seed *uint32) *ConnTrack {
	if seed == nil {
		r := rng.Uint32()
		seed = &r
	}
	ct := &ConnTrack{
		clock:     clock,
		rng:       rng,
		seed:      *seed,
		nftIDSeed: rng.Uint32(),
	}
	ct.init()
	return ct
}

// NewConnTrackWithReaper creates and initializes a new ConnTrack and reaper.
// Reaper garbage collects unused connections.
func NewConnTrackWithReaper(clock tcpip.Clock, rng connTrackRNG, seed *uint32) (*ConnTrack, tcpip.Timer) {
	ct := NewConnTrack(clock, rng, seed)
	var reaper tcpip.Timer
	bucket := 0
	interval := 1 * time.Second
	reaper = ct.clock.AfterFunc(interval, func() {
		bucket, interval = ct.reapUnused(bucket, interval)
		reaper.Reset(interval)
	})
	return ct, reaper
}

// NfConnTrackPriority returns the priority of the conntrack hook.
// Check `ipv4/ipv6_conntrack_ops` in nf_conntrack_proto.c.
func NfConnTrackPriority(hook NFHook) (int, bool) {
	switch hook {
	case NFPrerouting:
		// NF_IP_PRI_CONNTRACK
		return -200, true
	case NFInput:
		// NF_IP_PRI_CONNTRACK_CONFIRM
		return math.MaxInt32, true
	case NFPostrouting:
		// NF_IP_PRI_CONNTRACK_CONFIRM
		return math.MaxInt32, true
	case NFOutput:
		// NF_IP_PRI_CONNTRACK
		return -200, true
	}
	return 0, false
}