AdGuardHome/internal/client/storage_test.go

package client_test

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

	"github.com/AdguardTeam/AdGuardHome/internal/client"
	"github.com/AdguardTeam/AdGuardHome/internal/whois"
	"github.com/AdguardTeam/golibs/testutil"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

// newStorage is a helper function that returns a client storage filled with
// persistent clients from the m.  It also generates a UID for each client.
func newStorage(tb testing.TB, m []*client.Persistent) (s *client.Storage) {
	tb.Helper()

	s = client.NewStorage(&client.Config{
		AllowedTags: nil,
	})

	for _, c := range m {
		c.UID = client.MustNewUID()
		require.NoError(tb, s.Add(c))
	}

	require.Equal(tb, len(m), s.Size())

	return s
}

// newRuntimeClient is a helper function that returns a new runtime client.
func newRuntimeClient(ip netip.Addr, source client.Source, host string) (rc *client.Runtime) {
	rc = client.NewRuntime(ip)
	rc.SetInfo(source, []string{host})

	return rc
}

// mustParseMAC is wrapper around [net.ParseMAC] that panics if there is an
// error.
func mustParseMAC(s string) (mac net.HardwareAddr) {
	mac, err := net.ParseMAC(s)
	if err != nil {
		panic(err)
	}

	return mac
}

func TestStorage_Add(t *testing.T) {
	const (
		existingName     = "existing_name"
		existingClientID = "existing_client_id"

		allowedTag    = "tag"
		notAllowedTag = "not_allowed_tag"
	)

	var (
		existingClientUID = client.MustNewUID()
		existingIP        = netip.MustParseAddr("1.2.3.4")
		existingSubnet    = netip.MustParsePrefix("1.2.3.0/24")
	)

	existingClient := &client.Persistent{
		Name:      existingName,
		IPs:       []netip.Addr{existingIP},
		Subnets:   []netip.Prefix{existingSubnet},
		ClientIDs: []string{existingClientID},
		UID:       existingClientUID,
	}

	s := client.NewStorage(&client.Config{
		AllowedTags: []string{allowedTag},
	})
	err := s.Add(existingClient)
	require.NoError(t, err)

	testCases := []struct {
		name       string
		cli        *client.Persistent
		wantErrMsg string
	}{{
		name: "basic",
		cli: &client.Persistent{
			Name: "basic",
			IPs:  []netip.Addr{netip.MustParseAddr("1.1.1.1")},
			UID:  client.MustNewUID(),
		},
		wantErrMsg: "",
	}, {
		name: "duplicate_uid",
		cli: &client.Persistent{
			Name: "no_uid",
			IPs:  []netip.Addr{netip.MustParseAddr("2.2.2.2")},
			UID:  existingClientUID,
		},
		wantErrMsg: `adding client: another client "existing_name" uses the same uid`,
	}, {
		name: "duplicate_name",
		cli: &client.Persistent{
			Name: existingName,
			IPs:  []netip.Addr{netip.MustParseAddr("3.3.3.3")},
			UID:  client.MustNewUID(),
		},
		wantErrMsg: `adding client: another client uses the same name "existing_name"`,
	}, {
		name: "duplicate_ip",
		cli: &client.Persistent{
			Name: "duplicate_ip",
			IPs:  []netip.Addr{existingIP},
			UID:  client.MustNewUID(),
		},
		wantErrMsg: `adding client: another client "existing_name" uses the same IP "1.2.3.4"`,
	}, {
		name: "duplicate_subnet",
		cli: &client.Persistent{
			Name:    "duplicate_subnet",
			Subnets: []netip.Prefix{existingSubnet},
			UID:     client.MustNewUID(),
		},
		wantErrMsg: `adding client: another client "existing_name" ` +
			`uses the same subnet "1.2.3.0/24"`,
	}, {
		name: "duplicate_client_id",
		cli: &client.Persistent{
			Name:      "duplicate_client_id",
			ClientIDs: []string{existingClientID},
			UID:       client.MustNewUID(),
		},
		wantErrMsg: `adding client: another client "existing_name" ` +
			`uses the same ClientID "existing_client_id"`,
	}, {
		name: "not_allowed_tag",
		cli: &client.Persistent{
			Name: "nont_allowed_tag",
			Tags: []string{notAllowedTag},
			IPs:  []netip.Addr{netip.MustParseAddr("4.4.4.4")},
			UID:  client.MustNewUID(),
		},
		wantErrMsg: `adding client: invalid tag: "not_allowed_tag"`,
	}}

	for _, tc := range testCases {
		t.Run(tc.name, func(t *testing.T) {
			err = s.Add(tc.cli)

			testutil.AssertErrorMsg(t, tc.wantErrMsg, err)
		})
	}
}

func TestStorage_RemoveByName(t *testing.T) {
	const (
		existingName = "existing_name"
	)

	existingClient := &client.Persistent{
		Name: existingName,
		IPs:  []netip.Addr{netip.MustParseAddr("1.2.3.4")},
		UID:  client.MustNewUID(),
	}

	s := client.NewStorage(&client.Config{
		AllowedTags: nil,
	})
	err := s.Add(existingClient)
	require.NoError(t, err)

	testCases := []struct {
		want    assert.BoolAssertionFunc
		name    string
		cliName string
	}{{
		name:    "existing_client",
		cliName: existingName,
		want:    assert.True,
	}, {
		name:    "non_existing_client",
		cliName: "non_existing_client",
		want:    assert.False,
	}}

	for _, tc := range testCases {
		t.Run(tc.name, func(t *testing.T) {
			tc.want(t, s.RemoveByName(tc.cliName))
		})
	}

	t.Run("duplicate_remove", func(t *testing.T) {
		s = client.NewStorage(&client.Config{
			AllowedTags: nil,
		})
		err = s.Add(existingClient)
		require.NoError(t, err)

		assert.True(t, s.RemoveByName(existingName))
		assert.False(t, s.RemoveByName(existingName))
	})
}

func TestStorage_Find(t *testing.T) {
	const (
		cliIPNone = "1.2.3.4"
		cliIP1    = "1.1.1.1"
		cliIP2    = "2.2.2.2"

		cliIPv6 = "1:2:3::4"

		cliSubnet   = "2.2.2.0/24"
		cliSubnetIP = "2.2.2.222"

		cliID  = "client-id"
		cliMAC = "11:11:11:11:11:11"

		linkLocalIP     = "fe80::abcd:abcd:abcd:ab%eth0"
		linkLocalSubnet = "fe80::/16"
	)

	var (
		clientWithBothFams = &client.Persistent{
			Name: "client1",
			IPs: []netip.Addr{
				netip.MustParseAddr(cliIP1),
				netip.MustParseAddr(cliIPv6),
			},
		}

		clientWithSubnet = &client.Persistent{
			Name:    "client2",
			IPs:     []netip.Addr{netip.MustParseAddr(cliIP2)},
			Subnets: []netip.Prefix{netip.MustParsePrefix(cliSubnet)},
		}

		clientWithMAC = &client.Persistent{
			Name: "client_with_mac",
			MACs: []net.HardwareAddr{mustParseMAC(cliMAC)},
		}

		clientWithID = &client.Persistent{
			Name:      "client_with_id",
			ClientIDs: []string{cliID},
		}

		clientLinkLocal = &client.Persistent{
			Name:    "client_link_local",
			Subnets: []netip.Prefix{netip.MustParsePrefix(linkLocalSubnet)},
		}
	)

	clients := []*client.Persistent{
		clientWithBothFams,
		clientWithSubnet,
		clientWithMAC,
		clientWithID,
		clientLinkLocal,
	}
	s := newStorage(t, clients)

	testCases := []struct {
		want *client.Persistent
		name string
		ids  []string
	}{{
		name: "ipv4_ipv6",
		ids:  []string{cliIP1, cliIPv6},
		want: clientWithBothFams,
	}, {
		name: "ipv4_subnet",
		ids:  []string{cliIP2, cliSubnetIP},
		want: clientWithSubnet,
	}, {
		name: "mac",
		ids:  []string{cliMAC},
		want: clientWithMAC,
	}, {
		name: "client_id",
		ids:  []string{cliID},
		want: clientWithID,
	}, {
		name: "client_link_local_subnet",
		ids:  []string{linkLocalIP},
		want: clientLinkLocal,
	}}

	for _, tc := range testCases {
		t.Run(tc.name, func(t *testing.T) {
			for _, id := range tc.ids {
				c, ok := s.Find(id)
				require.True(t, ok)

				assert.Equal(t, tc.want, c)
			}
		})
	}

	t.Run("not_found", func(t *testing.T) {
		_, ok := s.Find(cliIPNone)
		assert.False(t, ok)
	})
}

func TestStorage_FindLoose(t *testing.T) {
	const (
		nonExistingClientID = "client_id"
	)

	var (
		ip         = netip.MustParseAddr("fe80::a098:7654:32ef:ff1")
		ipWithZone = netip.MustParseAddr("fe80::1ff:fe23:4567:890a%eth2")
	)

	var (
		clientNoZone = &client.Persistent{
			Name: "client",
			IPs:  []netip.Addr{ip},
		}

		clientWithZone = &client.Persistent{
			Name: "client_with_zone",
			IPs:  []netip.Addr{ipWithZone},
		}
	)

	s := newStorage(
		t,
		[]*client.Persistent{
			clientNoZone,
			clientWithZone,
		},
	)

	testCases := []struct {
		ip      netip.Addr
		want    assert.BoolAssertionFunc
		wantCli *client.Persistent
		name    string
	}{{
		name:    "without_zone",
		ip:      ip,
		wantCli: clientNoZone,
		want:    assert.True,
	}, {
		name:    "with_zone",
		ip:      ipWithZone,
		wantCli: clientWithZone,
		want:    assert.True,
	}, {
		name:    "zero_address",
		ip:      netip.Addr{},
		wantCli: nil,
		want:    assert.False,
	}}

	for _, tc := range testCases {
		t.Run(tc.name, func(t *testing.T) {
			c, ok := s.FindLoose(tc.ip.WithZone(""), nonExistingClientID)
			assert.Equal(t, tc.wantCli, c)
			tc.want(t, ok)
		})
	}
}

func TestStorage_FindByName(t *testing.T) {
	const (
		cliIP1 = "1.1.1.1"
		cliIP2 = "2.2.2.2"
	)

	const (
		clientExistingName        = "client_existing"
		clientAnotherExistingName = "client_another_existing"
		nonExistingClientName     = "client_non_existing"
	)

	var (
		clientExisting = &client.Persistent{
			Name: clientExistingName,
			IPs:  []netip.Addr{netip.MustParseAddr(cliIP1)},
		}

		clientAnotherExisting = &client.Persistent{
			Name: clientAnotherExistingName,
			IPs:  []netip.Addr{netip.MustParseAddr(cliIP2)},
		}
	)

	clients := []*client.Persistent{
		clientExisting,
		clientAnotherExisting,
	}
	s := newStorage(t, clients)

	testCases := []struct {
		want       *client.Persistent
		name       string
		clientName string
	}{{
		name:       "existing",
		clientName: clientExistingName,
		want:       clientExisting,
	}, {
		name:       "another_existing",
		clientName: clientAnotherExistingName,
		want:       clientAnotherExisting,
	}, {
		name:       "non_existing",
		clientName: nonExistingClientName,
		want:       nil,
	}}

	for _, tc := range testCases {
		t.Run(tc.name, func(t *testing.T) {
			c, ok := s.FindByName(tc.clientName)
			if tc.want == nil {
				assert.False(t, ok)

				return
			}

			assert.True(t, ok)
			assert.Equal(t, tc.want, c)
		})
	}
}

func TestStorage_FindByMAC(t *testing.T) {
	var (
		cliMAC               = mustParseMAC("11:11:11:11:11:11")
		cliAnotherMAC        = mustParseMAC("22:22:22:22:22:22")
		nonExistingClientMAC = mustParseMAC("33:33:33:33:33:33")
	)

	var (
		clientExisting = &client.Persistent{
			Name: "client",
			MACs: []net.HardwareAddr{cliMAC},
		}

		clientAnotherExisting = &client.Persistent{
			Name: "another_client",
			MACs: []net.HardwareAddr{cliAnotherMAC},
		}
	)

	clients := []*client.Persistent{
		clientExisting,
		clientAnotherExisting,
	}
	s := newStorage(t, clients)

	testCases := []struct {
		want      *client.Persistent
		name      string
		clientMAC net.HardwareAddr
	}{{
		name:      "existing",
		clientMAC: cliMAC,
		want:      clientExisting,
	}, {
		name:      "another_existing",
		clientMAC: cliAnotherMAC,
		want:      clientAnotherExisting,
	}, {
		name:      "non_existing",
		clientMAC: nonExistingClientMAC,
		want:      nil,
	}}

	for _, tc := range testCases {
		t.Run(tc.name, func(t *testing.T) {
			c, ok := s.FindByMAC(tc.clientMAC)
			if tc.want == nil {
				assert.False(t, ok)

				return
			}

			assert.True(t, ok)
			assert.Equal(t, tc.want, c)
		})
	}
}

func TestStorage_Update(t *testing.T) {
	const (
		clientName          = "client_name"
		obstructingName     = "obstructing_name"
		obstructingClientID = "obstructing_client_id"
	)

	var (
		obstructingIP     = netip.MustParseAddr("1.2.3.4")
		obstructingSubnet = netip.MustParsePrefix("1.2.3.0/24")
	)

	obstructingClient := &client.Persistent{
		Name:      obstructingName,
		IPs:       []netip.Addr{obstructingIP},
		Subnets:   []netip.Prefix{obstructingSubnet},
		ClientIDs: []string{obstructingClientID},
	}

	clientToUpdate := &client.Persistent{
		Name: clientName,
		IPs:  []netip.Addr{netip.MustParseAddr("1.1.1.1")},
	}

	testCases := []struct {
		name       string
		cli        *client.Persistent
		wantErrMsg string
	}{{
		name: "basic",
		cli: &client.Persistent{
			Name: "basic",
			IPs:  []netip.Addr{netip.MustParseAddr("1.1.1.1")},
			UID:  client.MustNewUID(),
		},
		wantErrMsg: "",
	}, {
		name: "duplicate_name",
		cli: &client.Persistent{
			Name: obstructingName,
			IPs:  []netip.Addr{netip.MustParseAddr("3.3.3.3")},
			UID:  client.MustNewUID(),
		},
		wantErrMsg: `updating client: another client uses the same name "obstructing_name"`,
	}, {
		name: "duplicate_ip",
		cli: &client.Persistent{
			Name: "duplicate_ip",
			IPs:  []netip.Addr{obstructingIP},
			UID:  client.MustNewUID(),
		},
		wantErrMsg: `updating client: another client "obstructing_name" uses the same IP "1.2.3.4"`,
	}, {
		name: "duplicate_subnet",
		cli: &client.Persistent{
			Name:    "duplicate_subnet",
			Subnets: []netip.Prefix{obstructingSubnet},
			UID:     client.MustNewUID(),
		},
		wantErrMsg: `updating client: another client "obstructing_name" ` +
			`uses the same subnet "1.2.3.0/24"`,
	}, {
		name: "duplicate_client_id",
		cli: &client.Persistent{
			Name:      "duplicate_client_id",
			ClientIDs: []string{obstructingClientID},
			UID:       client.MustNewUID(),
		},
		wantErrMsg: `updating client: another client "obstructing_name" ` +
			`uses the same ClientID "obstructing_client_id"`,
	}}

	for _, tc := range testCases {
		t.Run(tc.name, func(t *testing.T) {
			s := newStorage(
				t,
				[]*client.Persistent{
					clientToUpdate,
					obstructingClient,
				},
			)

			err := s.Update(clientName, tc.cli)
			testutil.AssertErrorMsg(t, tc.wantErrMsg, err)
		})
	}
}

func TestStorage_RangeByName(t *testing.T) {
	sortedClients := []*client.Persistent{{
		Name:      "clientA",
		ClientIDs: []string{"A"},
	}, {
		Name:      "clientB",
		ClientIDs: []string{"B"},
	}, {
		Name:      "clientC",
		ClientIDs: []string{"C"},
	}, {
		Name:      "clientD",
		ClientIDs: []string{"D"},
	}, {
		Name:      "clientE",
		ClientIDs: []string{"E"},
	}}

	testCases := []struct {
		name string
		want []*client.Persistent
	}{{
		name: "basic",
		want: sortedClients,
	}, {
		name: "nil",
		want: nil,
	}, {
		name: "one_element",
		want: sortedClients[:1],
	}, {
		name: "two_elements",
		want: sortedClients[:2],
	}}

	for _, tc := range testCases {
		t.Run(tc.name, func(t *testing.T) {
			s := newStorage(t, tc.want)

			var got []*client.Persistent
			s.RangeByName(func(c *client.Persistent) (cont bool) {
				got = append(got, c)

				return true
			})

			assert.Equal(t, tc.want, got)
		})
	}
}

func TestStorage_UpdateRuntime(t *testing.T) {
	const (
		addedARP       = "added_arp"
		addedSecondARP = "added_arp"

		updatedARP = "updated_arp"

		cliCity    = "City"
		cliCountry = "Country"
		cliOrgname = "Orgname"
	)

	var (
		ip  = netip.MustParseAddr("1.1.1.1")
		ip2 = netip.MustParseAddr("2.2.2.2")
	)

	updated := client.NewRuntime(ip)
	updated.SetInfo(client.SourceARP, []string{updatedARP})

	info := &whois.Info{
		City:    cliCity,
		Country: cliCountry,
		Orgname: cliOrgname,
	}
	updated.SetWHOIS(info)

	s := client.NewStorage(&client.Config{
		AllowedTags: nil,
	})

	t.Run("add_arp_client", func(t *testing.T) {
		added := client.NewRuntime(ip)
		added.SetInfo(client.SourceARP, []string{addedARP})

		require.True(t, s.UpdateRuntime(added))
		require.Equal(t, 1, s.SizeRuntime())

		got := s.ClientRuntime(ip)
		source, host := got.Info()
		assert.Equal(t, client.SourceARP, source)
		assert.Equal(t, addedARP, host)
	})

	t.Run("add_second_arp_client", func(t *testing.T) {
		added := client.NewRuntime(ip2)
		added.SetInfo(client.SourceARP, []string{addedSecondARP})

		require.True(t, s.UpdateRuntime(added))
		require.Equal(t, 2, s.SizeRuntime())

		got := s.ClientRuntime(ip2)
		source, host := got.Info()
		assert.Equal(t, client.SourceARP, source)
		assert.Equal(t, addedSecondARP, host)
	})

	t.Run("update_first_client", func(t *testing.T) {
		require.False(t, s.UpdateRuntime(updated))
		got := s.ClientRuntime(ip)
		require.Equal(t, 2, s.SizeRuntime())

		source, host := got.Info()
		assert.Equal(t, client.SourceARP, source)
		assert.Equal(t, updatedARP, host)
	})

	t.Run("remove_arp_info", func(t *testing.T) {
		n := s.DeleteBySource(client.SourceARP)
		require.Equal(t, 1, n)
		require.Equal(t, 1, s.SizeRuntime())

		got := s.ClientRuntime(ip)
		source, _ := got.Info()
		assert.Equal(t, client.SourceWHOIS, source)
		assert.Equal(t, info, got.WHOIS())
	})

	t.Run("remove_whois_info", func(t *testing.T) {
		n := s.DeleteBySource(client.SourceWHOIS)
		require.Equal(t, 1, n)
		require.Equal(t, 0, s.SizeRuntime())
	})
}

func TestStorage_BatchUpdateBySource(t *testing.T) {
	const (
		defSrc = client.SourceARP

		cliFirstHost1   = "host1"
		cliFirstHost2   = "host2"
		cliUpdatedHost3 = "host3"
		cliUpdatedHost4 = "host4"
		cliUpdatedHost5 = "host5"
	)

	var (
		cliFirstIP1   = netip.MustParseAddr("1.1.1.1")
		cliFirstIP2   = netip.MustParseAddr("2.2.2.2")
		cliUpdatedIP3 = netip.MustParseAddr("3.3.3.3")
		cliUpdatedIP4 = netip.MustParseAddr("4.4.4.4")
		cliUpdatedIP5 = netip.MustParseAddr("5.5.5.5")
	)

	firstClients := []*client.Runtime{
		newRuntimeClient(cliFirstIP1, defSrc, cliFirstHost1),
		newRuntimeClient(cliFirstIP2, defSrc, cliFirstHost2),
	}

	updatedClients := []*client.Runtime{
		newRuntimeClient(cliUpdatedIP3, defSrc, cliUpdatedHost3),
		newRuntimeClient(cliUpdatedIP4, defSrc, cliUpdatedHost4),
		newRuntimeClient(cliUpdatedIP5, defSrc, cliUpdatedHost5),
	}

	s := client.NewStorage(&client.Config{
		AllowedTags: nil,
	})

	t.Run("populate_storage_with_first_clients", func(t *testing.T) {
		added, removed := s.BatchUpdateBySource(defSrc, firstClients)
		require.Equal(t, len(firstClients), added)
		require.Equal(t, 0, removed)
		require.Equal(t, len(firstClients), s.SizeRuntime())

		rc := s.ClientRuntime(cliFirstIP1)
		src, host := rc.Info()
		assert.Equal(t, defSrc, src)
		assert.Equal(t, cliFirstHost1, host)
	})

	t.Run("update_storage", func(t *testing.T) {
		added, removed := s.BatchUpdateBySource(defSrc, updatedClients)
		require.Equal(t, len(updatedClients), added)
		require.Equal(t, len(firstClients), removed)
		require.Equal(t, len(updatedClients), s.SizeRuntime())

		rc := s.ClientRuntime(cliUpdatedIP3)
		src, host := rc.Info()
		assert.Equal(t, defSrc, src)
		assert.Equal(t, cliUpdatedHost3, host)

		rc = s.ClientRuntime(cliFirstIP1)
		assert.Nil(t, rc)
	})

	t.Run("remove_all", func(t *testing.T) {
		added, removed := s.BatchUpdateBySource(defSrc, []*client.Runtime{})
		require.Equal(t, 0, added)
		require.Equal(t, len(updatedClients), removed)
		require.Equal(t, 0, s.SizeRuntime())
	})
}