Skip to content

Commit

Permalink
Merge branch 'master' into AGDNS-1982-upd-proxy
Browse files Browse the repository at this point in the history
  • Loading branch information
EugeneOne1 committed Apr 4, 2024
2 parents 9feeba5 + fd25dca commit a79298d
Show file tree
Hide file tree
Showing 6 changed files with 219 additions and 59 deletions.
26 changes: 22 additions & 4 deletions internal/client/client.go
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@ package client
import (
"encoding"
"fmt"
"net/netip"

"github.com/AdguardTeam/AdGuardHome/internal/whois"
)
Expand Down Expand Up @@ -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

Expand All @@ -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{}
Expand Down Expand Up @@ -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
Expand All @@ -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
}
63 changes: 63 additions & 0 deletions internal/client/runtimeindex.go
Original file line number Diff line number Diff line change
@@ -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
}
85 changes: 85 additions & 0 deletions internal/client/runtimeindex_test.go
Original file line number Diff line number Diff line change
@@ -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())
})
}
Loading

0 comments on commit a79298d

Please sign in to comment.