// Copyright (c) Tailscale Inc & contributors
// SPDX-License-Identifier: BSD-3-Clause

package conn25

import (
	"errors"
	"net/netip"
	"testing"

	"go4.org/netipx"
	"tailscale.com/util/must"
)

func TestNext(t *testing.T) {
	a := ipSetIterator{}
	_, err := a.next()
	if !errors.Is(err, errPoolExhausted) {
		t.Fatalf("expected errPoolExhausted, got %v", err)
	}

	var isb netipx.IPSetBuilder
	ipset := must.Get(isb.IPSet())
	b := newIPPool(ipset)
	_, err = b.next()
	if !errors.Is(err, errPoolExhausted) {
		t.Fatalf("expected errPoolExhausted, got %v", err)
	}

	isb.AddRange(netipx.IPRangeFrom(netip.MustParseAddr("192.168.0.0"), netip.MustParseAddr("192.168.0.2")))
	isb.AddRange(netipx.IPRangeFrom(netip.MustParseAddr("200.0.0.0"), netip.MustParseAddr("200.0.0.0")))
	isb.AddRange(netipx.IPRangeFrom(netip.MustParseAddr("201.0.0.0"), netip.MustParseAddr("201.0.0.1")))
	ipset = must.Get(isb.IPSet())
	c := newIPPool(ipset)
	expected := []string{
		"192.168.0.0",
		"192.168.0.1",
		"192.168.0.2",
		"200.0.0.0",
		"201.0.0.0",
		"201.0.0.1",
	}
	for i, want := range expected {
		addr, err := c.next()
		if err != nil {
			t.Fatal(err)
		}
		if addr != netip.MustParseAddr(want) {
			t.Fatalf("next call %d want: %s, got: %v", i, want, addr)
		}
	}
	_, err = c.next()
	if !errors.Is(err, errPoolExhausted) {
		t.Fatalf("expected errPoolExhausted, got %v", err)
	}
	_, err = c.next()
	if !errors.Is(err, errPoolExhausted) {
		t.Fatalf("expected errPoolExhausted, got %v", err)
	}
}

// TestReturnAddr tests that if a pool is exhausted, an address can be returned to the
// pool, and then that address will be handed out again.
func TestReturnAddr(t *testing.T) {
	addrString := "192.168.0.0"
	// There's an IPPool with one address in it.
	var isb netipx.IPSetBuilder
	isb.AddRange(netipx.IPRangeFrom(netip.MustParseAddr(addrString), netip.MustParseAddr(addrString)))
	ipset := must.Get(isb.IPSet())
	ipp := newIPPool(ipset)
	// The first time we call next we get the address.
	addr, err := ipp.next()
	if err != nil {
		t.Fatalf("expected nil error, got: %v", err)
	}
	if addr != netip.MustParseAddr(addrString) {
		t.Fatalf("want %v, got %v", addrString, addr)
	}
	// The second time we call next we get errPoolExhausted
	_, err = ipp.next()
	if !errors.Is(err, errPoolExhausted) {
		t.Fatalf("expected errPoolExhausted, got %v", err)
	}
	// Return the addr to the pool
	err = ipp.returnAddr(netip.MustParseAddr(addrString))
	if err != nil {
		t.Fatal(err)
	}
	// It's not possible to return addresses that are already in the pool.
	err = ipp.returnAddr(netip.MustParseAddr(addrString))
	if !errors.Is(err, errAddrExists) {
		t.Fatalf("want errAddrExists, got: %v", err)
	}
	// When we call next we get the returned addr
	addrAfterReturn, err := ipp.next()
	if err != nil {
		t.Fatalf("expected nil error, got: %v", err)
	}
	if addrAfterReturn != netip.MustParseAddr(addrString) {
		t.Fatalf("want %v, got %v", addrString, addrAfterReturn)
	}
	// You can't return addresses that aren't from the pool.
	err = ipp.returnAddr(netip.MustParseAddr("100.100.100.0"))
	if !errors.Is(err, errNotOurAddress) {
		t.Fatalf("want errNotOurAddress, got: %v", err)
	}
}

func expectAddrNext(t *testing.T, ipp *ippool, addrString string) {
	t.Helper()
	got, err := ipp.next()
	if err != nil {
		t.Fatalf("expected nil error, got: %v", err)
	}
	want := netip.MustParseAddr(addrString)
	if want != got {
		t.Fatalf("want %v; got %v", want, got)
	}
}

func expectErrPoolExhaustedNext(t *testing.T, ipp *ippool) {
	t.Helper()
	_, err := ipp.next()
	if !errors.Is(err, errPoolExhausted) {
		t.Fatalf("expected errPoolExhausted; got %v", err)
	}
}

// TestGettingReturnedAddresses tests that when addresses are returned to the IP Pool
// they are then handed out in the order they were returned.
func TestGettingReturnedAddresses(t *testing.T) {
	var isb netipx.IPSetBuilder
	isb.AddRange(netipx.IPRangeFrom(netip.MustParseAddr("192.168.0.0"), netip.MustParseAddr("192.168.0.4")))
	ipset := must.Get(isb.IPSet())
	ipp := newIPPool(ipset)
	expectAddrNext(t, ipp, "192.168.0.0")
	expectAddrNext(t, ipp, "192.168.0.1")
	expectAddrNext(t, ipp, "192.168.0.2")
	expectAddrNext(t, ipp, "192.168.0.3")
	expectAddrNext(t, ipp, "192.168.0.4")
	expectErrPoolExhaustedNext(t, ipp)
	ipp.returnAddr(netip.MustParseAddr("192.168.0.2"))
	ipp.returnAddr(netip.MustParseAddr("192.168.0.4"))
	expectAddrNext(t, ipp, "192.168.0.2")
	expectAddrNext(t, ipp, "192.168.0.4")
	expectErrPoolExhaustedNext(t, ipp)
}

func TestIPPoolReconfig(t *testing.T) {
	var isb netipx.IPSetBuilder
	isb.AddRange(netipx.IPRangeFrom(netip.MustParseAddr("192.168.0.0"), netip.MustParseAddr("192.168.0.4")))
	ipsetOne := must.Get(isb.IPSet())
	ipsetOneClone := must.Get(isb.IPSet())
	isb = netipx.IPSetBuilder{}
	isb.AddRange(netipx.IPRangeFrom(netip.MustParseAddr("192.168.0.7"), netip.MustParseAddr("192.168.0.10")))
	ipsetTwo := must.Get(isb.IPSet())

	var ipp *ippool
	ipp = ipp.reconfig(ipsetOne)
	if ipp.ipSet != ipsetOne {
		t.Fatalf("want %v, got %v", ipsetOne, ipp.ipSet)
	}
	expectAddrNext(t, ipp, "192.168.0.0")

	// check that we don't lose iterator state when we reconfig with the same ranges
	expectAddrNext(t, ipp, "192.168.0.1")
	ipp.returnAddr(netip.MustParseAddr("192.168.0.1"))
	ipp = ipp.reconfig(ipsetOneClone)
	expectAddrNext(t, ipp, "192.168.0.2")

	// when we reconfig with different ranges, we only hand out addresses from the new ranges
	ipp = ipp.reconfig(ipsetTwo)
	if ipp.ipSet != ipsetTwo {
		t.Fatalf("want %v, got %v", ipsetTwo, ipp.ipSet)
	}
	expectAddrNext(t, ipp, "192.168.0.7")
	expectAddrNext(t, ipp, "192.168.0.8")
	expectAddrNext(t, ipp, "192.168.0.9")
	expectAddrNext(t, ipp, "192.168.0.10")
	expectErrPoolExhaustedNext(t, ipp)

	// but we have not lost track of the fact that the old addresses are in use
	if !ipp.inUse.Contains(netip.MustParseAddr("192.168.0.0")) {
		t.Fatalf("expected inUse to still have the address")
	}

	// old addresses can be returned
	ipp.returnAddr(netip.MustParseAddr("192.168.0.0"))

	// but they are not handed out again
	expectErrPoolExhaustedNext(t, ipp)
	if ipp.inUse.Contains(netip.MustParseAddr("192.168.0.0")) {
		t.Fatalf("expected inUse to no longer have the address")
	}

	// returning addresses from the new ranges works as normal
	ipp.returnAddr(netip.MustParseAddr("192.168.0.9"))
	expectAddrNext(t, ipp, "192.168.0.9")
}
