Pull request 2183: AG-27492-client-runtime-index

Squashed commit of the following:

commit d0b37e3
Merge: 025c29b ee619b2
Author: Stanislav Chzhen <s.chzhen@adguard.com>
Date:   Thu Apr 4 18:58:08 2024 +0300

    Merge branch 'master' into AG-27492-client-runtime-index

commit 025c29b
Author: Stanislav Chzhen <s.chzhen@adguard.com>
Date:   Mon Apr 1 17:20:15 2024 +0300

    client: imp code

commit 548a15c
Author: Stanislav Chzhen <s.chzhen@adguard.com>
Date:   Thu Mar 28 13:43:17 2024 +0300

    all: add tests

commit c9015e7
Author: Stanislav Chzhen <s.chzhen@adguard.com>
Date:   Mon Mar 25 16:33:30 2024 +0300

    all: imp docs

commit 81e8b94
Author: Stanislav Chzhen <s.chzhen@adguard.com>
Date:   Mon Mar 25 15:33:17 2024 +0300

    all: imp code

commit 1428d60
Author: Stanislav Chzhen <s.chzhen@adguard.com>
Date:   Mon Mar 25 14:45:01 2024 +0300

    all: client runtime index

internal/client/client.go

@@ -7,6 +7,7 @@ package client
 import (
 	"encoding"
 	"fmt"
+	"net/netip"
 
 	"github.com/AdguardTeam/AdGuardHome/internal/whois"
 )
@@ -56,6 +57,9 @@ func (cs Source) MarshalText() (text []byte, err error) {
 
 // Runtime is a client information from different sources.
 type Runtime struct {
+	// ip is an IP address of a client.
+	ip netip.Addr
+
 	// whois is the filtered WHOIS information of a client.
 	whois *whois.Info
 
@@ -80,6 +84,15 @@ type Runtime struct {
 	hostsFile []string
 }
 
+// NewRuntime constructs a new runtime client.  ip must be valid IP address.
+//
+// TODO(s.chzhen):  Validate IP address.
+func NewRuntime(ip netip.Addr) (r *Runtime) {
+	return &Runtime{
+		ip: ip,
+	}
+}
+
 // Info returns a client information from the highest-priority source.
 func (r *Runtime) Info() (cs Source, host string) {
 	info := []string{}
@@ -133,8 +146,8 @@ func (r *Runtime) SetWHOIS(info *whois.Info) {
 	r.whois = info
 }
 
-// Unset clears a cs information.
-func (r *Runtime) Unset(cs Source) {
+// unset clears a cs information.
+func (r *Runtime) unset(cs Source) {
 	switch cs {
 	case SourceWHOIS:
 		r.whois = nil
@@ -149,11 +162,16 @@ func (r *Runtime) Unset(cs Source) {
 	}
 }
 
-// IsEmpty returns true if there is no information from any source.
-func (r *Runtime) IsEmpty() (ok bool) {
+// isEmpty returns true if there is no information from any source.
+func (r *Runtime) isEmpty() (ok bool) {
 	return r.whois == nil &&
 		r.arp == nil &&
 		r.rdns == nil &&
 		r.dhcp == nil &&
 		r.hostsFile == nil
 }
+
+// Addr returns an IP address of the client.
+func (r *Runtime) Addr() (ip netip.Addr) {
+	return r.ip
+}

internal/client/runtimeindex.go

@@ -0,0 +1,63 @@
+package client
+
+import "net/netip"
+
+// RuntimeIndex stores information about runtime clients.
+type RuntimeIndex struct {
+	// index maps IP address to runtime client.
+	index map[netip.Addr]*Runtime
+}
+
+// NewRuntimeIndex returns initialized runtime index.
+func NewRuntimeIndex() (ri *RuntimeIndex) {
+	return &RuntimeIndex{
+		index: map[netip.Addr]*Runtime{},
+	}
+}
+
+// Client returns the saved runtime client by ip.  If no such client exists,
+// returns nil.
+func (ri *RuntimeIndex) Client(ip netip.Addr) (rc *Runtime) {
+	return ri.index[ip]
+}
+
+// Add saves the runtime client in the index.  IP address of a client must be
+// unique.  See [Runtime.Client].  rc must not be nil.
+func (ri *RuntimeIndex) Add(rc *Runtime) {
+	ip := rc.Addr()
+	ri.index[ip] = rc
+}
+
+// Size returns the number of the runtime clients.
+func (ri *RuntimeIndex) Size() (n int) {
+	return len(ri.index)
+}
+
+// Range calls f for each runtime client in an undefined order.
+func (ri *RuntimeIndex) Range(f func(rc *Runtime) (cont bool)) {
+	for _, rc := range ri.index {
+		if !f(rc) {
+			return
+		}
+	}
+}
+
+// Delete removes the runtime client by ip.
+func (ri *RuntimeIndex) Delete(ip netip.Addr) {
+	delete(ri.index, ip)
+}
+
+// DeleteBySource removes all runtime clients that have information only from
+// the specified source and returns the number of removed clients.
+func (ri *RuntimeIndex) DeleteBySource(src Source) (n int) {
+	for ip, rc := range ri.index {
+		rc.unset(src)
+
+		if rc.isEmpty() {
+			delete(ri.index, ip)
+			n++
+		}
+	}
+
+	return n
+}

internal/client/runtimeindex_test.go

@@ -0,0 +1,85 @@
+package client_test
+
+import (
+	"net/netip"
+	"testing"
+
+	"github.com/AdguardTeam/AdGuardHome/internal/client"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestRuntimeIndex(t *testing.T) {
+	const cliSrc = client.SourceARP
+
+	var (
+		ip1 = netip.MustParseAddr("1.1.1.1")
+		ip2 = netip.MustParseAddr("2.2.2.2")
+		ip3 = netip.MustParseAddr("3.3.3.3")
+	)
+
+	ri := client.NewRuntimeIndex()
+	currentSize := 0
+
+	testCases := []struct {
+		ip    netip.Addr
+		name  string
+		hosts []string
+		src   client.Source
+	}{{
+		src:   cliSrc,
+		ip:    ip1,
+		name:  "1",
+		hosts: []string{"host1"},
+	}, {
+		src:   cliSrc,
+		ip:    ip2,
+		name:  "2",
+		hosts: []string{"host2"},
+	}, {
+		src:   cliSrc,
+		ip:    ip3,
+		name:  "3",
+		hosts: []string{"host3"},
+	}}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			rc := client.NewRuntime(tc.ip)
+			rc.SetInfo(tc.src, tc.hosts)
+
+			ri.Add(rc)
+			currentSize++
+
+			got := ri.Client(tc.ip)
+			assert.Equal(t, rc, got)
+		})
+	}
+
+	t.Run("size", func(t *testing.T) {
+		assert.Equal(t, currentSize, ri.Size())
+	})
+
+	t.Run("range", func(t *testing.T) {
+		s := 0
+
+		ri.Range(func(rc *client.Runtime) (cont bool) {
+			s++
+
+			return true
+		})
+
+		assert.Equal(t, currentSize, s)
+	})
+
+	t.Run("delete", func(t *testing.T) {
+		ri.Delete(ip1)
+		currentSize--
+
+		assert.Equal(t, currentSize, ri.Size())
+	})
+
+	t.Run("delete_by_src", func(t *testing.T) {
+		assert.Equal(t, currentSize, ri.DeleteBySource(cliSrc))
+		assert.Equal(t, 0, ri.Size())
+	})
+}