1package addrmgr
2
3import (
4	"io/ioutil"
5	"math/rand"
6	"net"
7	"os"
8	"testing"
9
10	"github.com/btcsuite/btcd/wire"
11)
12
13// randAddr generates a *wire.NetAddress backed by a random IPv4/IPv6 address.
14func randAddr(t *testing.T) *wire.NetAddress {
15	t.Helper()
16
17	ipv4 := rand.Intn(2) == 0
18	var ip net.IP
19	if ipv4 {
20		var b [4]byte
21		if _, err := rand.Read(b[:]); err != nil {
22			t.Fatal(err)
23		}
24		ip = b[:]
25	} else {
26		var b [16]byte
27		if _, err := rand.Read(b[:]); err != nil {
28			t.Fatal(err)
29		}
30		ip = b[:]
31	}
32
33	return &wire.NetAddress{
34		Services: wire.ServiceFlag(rand.Uint64()),
35		IP:       ip,
36		Port:     uint16(rand.Uint32()),
37	}
38}
39
40// assertAddr ensures that the two addresses match. The timestamp is not
41// checked as it does not affect uniquely identifying a specific address.
42func assertAddr(t *testing.T, got, expected *wire.NetAddress) {
43	if got.Services != expected.Services {
44		t.Fatalf("expected address services %v, got %v",
45			expected.Services, got.Services)
46	}
47	if !got.IP.Equal(expected.IP) {
48		t.Fatalf("expected address IP %v, got %v", expected.IP, got.IP)
49	}
50	if got.Port != expected.Port {
51		t.Fatalf("expected address port %d, got %d", expected.Port,
52			got.Port)
53	}
54}
55
56// assertAddrs ensures that the manager's address cache matches the given
57// expected addresses.
58func assertAddrs(t *testing.T, addrMgr *AddrManager,
59	expectedAddrs map[string]*wire.NetAddress) {
60
61	t.Helper()
62
63	addrs := addrMgr.getAddresses()
64
65	if len(addrs) != len(expectedAddrs) {
66		t.Fatalf("expected to find %d addresses, found %d",
67			len(expectedAddrs), len(addrs))
68	}
69
70	for _, addr := range addrs {
71		addrStr := NetAddressKey(addr)
72		expectedAddr, ok := expectedAddrs[addrStr]
73		if !ok {
74			t.Fatalf("expected to find address %v", addrStr)
75		}
76
77		assertAddr(t, addr, expectedAddr)
78	}
79}
80
81// TestAddrManagerSerialization ensures that we can properly serialize and
82// deserialize the manager's current address cache.
83func TestAddrManagerSerialization(t *testing.T) {
84	t.Parallel()
85
86	// We'll start by creating our address manager backed by a temporary
87	// directory.
88	tempDir, err := ioutil.TempDir("", "addrmgr")
89	if err != nil {
90		t.Fatalf("unable to create temp dir: %v", err)
91	}
92	defer os.RemoveAll(tempDir)
93
94	addrMgr := New(tempDir, nil)
95
96	// We'll be adding 5 random addresses to the manager.
97	const numAddrs = 5
98
99	expectedAddrs := make(map[string]*wire.NetAddress, numAddrs)
100	for i := 0; i < numAddrs; i++ {
101		addr := randAddr(t)
102		expectedAddrs[NetAddressKey(addr)] = addr
103		addrMgr.AddAddress(addr, randAddr(t))
104	}
105
106	// Now that the addresses have been added, we should be able to retrieve
107	// them.
108	assertAddrs(t, addrMgr, expectedAddrs)
109
110	// Then, we'll persist these addresses to disk and restart the address
111	// manager.
112	addrMgr.savePeers()
113	addrMgr = New(tempDir, nil)
114
115	// Finally, we'll read all of the addresses from disk and ensure they
116	// match as expected.
117	addrMgr.loadPeers()
118	assertAddrs(t, addrMgr, expectedAddrs)
119}
120
121// TestAddrManagerV1ToV2 ensures that we can properly upgrade the serialized
122// version of the address manager from v1 to v2.
123func TestAddrManagerV1ToV2(t *testing.T) {
124	t.Parallel()
125
126	// We'll start by creating our address manager backed by a temporary
127	// directory.
128	tempDir, err := ioutil.TempDir("", "addrmgr")
129	if err != nil {
130		t.Fatalf("unable to create temp dir: %v", err)
131	}
132	defer os.RemoveAll(tempDir)
133
134	addrMgr := New(tempDir, nil)
135
136	// As we're interested in testing the upgrade path from v1 to v2, we'll
137	// override the manager's current version.
138	addrMgr.version = 1
139
140	// We'll be adding 5 random addresses to the manager. Since this is v1,
141	// each addresses' services will not be stored.
142	const numAddrs = 5
143
144	expectedAddrs := make(map[string]*wire.NetAddress, numAddrs)
145	for i := 0; i < numAddrs; i++ {
146		addr := randAddr(t)
147		expectedAddrs[NetAddressKey(addr)] = addr
148		addrMgr.AddAddress(addr, randAddr(t))
149	}
150
151	// Then, we'll persist these addresses to disk and restart the address
152	// manager - overriding its version back to v1.
153	addrMgr.savePeers()
154	addrMgr = New(tempDir, nil)
155	addrMgr.version = 1
156
157	// When we read all of the addresses back from disk, we should expect to
158	// find all of them, but their services will be set to a default of
159	// SFNodeNetwork since they were not previously stored. After ensuring
160	// that this default is set, we'll override each addresses' services
161	// with the original value from when they were created.
162	addrMgr.loadPeers()
163	addrs := addrMgr.getAddresses()
164	if len(addrs) != len(expectedAddrs) {
165		t.Fatalf("expected to find %d adddresses, found %d",
166			len(expectedAddrs), len(addrs))
167	}
168	for _, addr := range addrs {
169		addrStr := NetAddressKey(addr)
170		expectedAddr, ok := expectedAddrs[addrStr]
171		if !ok {
172			t.Fatalf("expected to find address %v", addrStr)
173		}
174
175		if addr.Services != wire.SFNodeNetwork {
176			t.Fatalf("expected address services to be %v, got %v",
177				wire.SFNodeNetwork, addr.Services)
178		}
179
180		addrMgr.SetServices(addr, expectedAddr.Services)
181	}
182
183	// We'll also bump up the manager's version to v2, which should signal
184	// that it should include the address services when persisting its
185	// state.
186	addrMgr.version = 2
187	addrMgr.savePeers()
188
189	// Finally, we'll recreate the manager and ensure that the services were
190	// persisted correctly.
191	addrMgr = New(tempDir, nil)
192	addrMgr.loadPeers()
193	assertAddrs(t, addrMgr, expectedAddrs)
194}
195