mock_graph_test.go
1 package routing 2 3 import ( 4 "bytes" 5 "context" 6 "fmt" 7 "testing" 8 9 "github.com/btcsuite/btcd/btcec/v2" 10 "github.com/btcsuite/btcd/btcutil" 11 graphdb "github.com/lightningnetwork/lnd/graph/db" 12 "github.com/lightningnetwork/lnd/graph/db/models" 13 "github.com/lightningnetwork/lnd/lnwire" 14 "github.com/lightningnetwork/lnd/routing/route" 15 ) 16 17 // createPubkey return a new test pubkey. 18 func createPubkey(id byte) route.Vertex { 19 _, secpPub := btcec.PrivKeyFromBytes([]byte{id}) 20 21 var bytes [33]byte 22 copy(bytes[:], secpPub.SerializeCompressed()[:33]) 23 24 pubkey := route.Vertex(bytes) 25 return pubkey 26 } 27 28 // mockChannel holds the channel state of a channel in the mock graph. 29 type mockChannel struct { 30 id uint64 31 capacity btcutil.Amount 32 balance lnwire.MilliSatoshi 33 } 34 35 // mockNode holds a set of mock channels and routing policies for a node in the 36 // mock graph. 37 type mockNode struct { 38 channels map[route.Vertex]*mockChannel 39 baseFee lnwire.MilliSatoshi 40 pubkey route.Vertex 41 } 42 43 // newMockNode instantiates a new mock node with a newly generated pubkey. 44 func newMockNode(id byte) *mockNode { 45 pubkey := createPubkey(id) 46 return &mockNode{ 47 channels: make(map[route.Vertex]*mockChannel), 48 pubkey: pubkey, 49 } 50 } 51 52 // fwd simulates an htlc forward through this node. If the from parameter is 53 // nil, this node is considered to be the sender of the payment. The route 54 // parameter describes the remaining route from this node onwards. If route.next 55 // is nil, this node is the final hop. 56 func (m *mockNode) fwd(from *mockNode, route *hop) (htlcResult, error) { 57 next := route.next 58 59 // Get the incoming channel, if any. 60 var inChan *mockChannel 61 if from != nil { 62 inChan = m.channels[from.pubkey] 63 } 64 65 // If there is no next node, this is the final node and we can settle the htlc. 66 if next == nil { 67 // Update the incoming balance. 68 inChan.balance += route.amtToFwd 69 70 return htlcResult{}, nil 71 } 72 73 // Check if the outgoing channel has enough balance. 74 outChan, ok := m.channels[next.node.pubkey] 75 if !ok { 76 return htlcResult{}, 77 fmt.Errorf("%v: unknown next %v", 78 m.pubkey, next.node.pubkey) 79 } 80 if outChan.balance < route.amtToFwd { 81 return htlcResult{ 82 failureSource: m.pubkey, 83 failure: lnwire.NewTemporaryChannelFailure(nil), 84 }, nil 85 } 86 87 // Htlc can be forwarded, update channel balances. 88 outChan.balance -= route.amtToFwd 89 if inChan != nil { 90 inChan.balance += route.amtToFwd 91 } 92 93 // Recursively forward down the given route. 94 result, err := next.node.fwd(m, route.next) 95 if err != nil { 96 return htlcResult{}, err 97 } 98 99 // Revert balances when a failure occurs. 100 if result.failure != nil { 101 outChan.balance += route.amtToFwd 102 if inChan != nil { 103 inChan.balance -= route.amtToFwd 104 } 105 } 106 107 return result, nil 108 } 109 110 // mockGraph contains a set of nodes that together for a mocked graph. 111 type mockGraph struct { 112 t *testing.T 113 nodes map[route.Vertex]*mockNode 114 source *mockNode 115 } 116 117 // newMockGraph instantiates a new mock graph. 118 func newMockGraph(t *testing.T) *mockGraph { 119 return &mockGraph{ 120 nodes: make(map[route.Vertex]*mockNode), 121 t: t, 122 } 123 } 124 125 // addNode adds the given mock node to the network. 126 func (m *mockGraph) addNode(node *mockNode) { 127 m.t.Helper() 128 129 if _, exists := m.nodes[node.pubkey]; exists { 130 m.t.Fatal("node already exists") 131 } 132 m.nodes[node.pubkey] = node 133 } 134 135 // addChannel adds a new channel between two existing nodes on the network. It 136 // sets the channel balance to 50/50%. 137 // 138 // Ignore linter error because addChannel isn't yet called with different 139 // capacities. 140 // nolint:unparam 141 func (m *mockGraph) addChannel(id uint64, node1id, node2id byte, 142 capacity btcutil.Amount) { 143 144 node1pubkey := createPubkey(node1id) 145 node2pubkey := createPubkey(node2id) 146 147 if _, exists := m.nodes[node1pubkey].channels[node2pubkey]; exists { 148 m.t.Fatal("channel already exists") 149 } 150 if _, exists := m.nodes[node2pubkey].channels[node1pubkey]; exists { 151 m.t.Fatal("channel already exists") 152 } 153 154 m.nodes[node1pubkey].channels[node2pubkey] = &mockChannel{ 155 capacity: capacity, 156 id: id, 157 balance: lnwire.NewMSatFromSatoshis(capacity / 2), 158 } 159 m.nodes[node2pubkey].channels[node1pubkey] = &mockChannel{ 160 capacity: capacity, 161 id: id, 162 balance: lnwire.NewMSatFromSatoshis(capacity / 2), 163 } 164 } 165 166 // forEachNodeChannel calls the callback for every channel of the given node. 167 // 168 // NOTE: Part of the Graph interface. 169 func (m *mockGraph) ForEachNodeDirectedChannel(_ context.Context, 170 nodePub route.Vertex, cb func(channel *graphdb.DirectedChannel) error, 171 _ func()) error { 172 173 // Look up the mock node. 174 node, ok := m.nodes[nodePub] 175 if !ok { 176 return graphdb.ErrGraphNodeNotFound 177 } 178 179 // Iterate over all of its channels. 180 for peer, channel := range node.channels { 181 // Lexicographically sort the pubkeys. 182 var node1 route.Vertex 183 if bytes.Compare(nodePub[:], peer[:]) == -1 { 184 node1 = peer 185 } else { 186 node1 = nodePub 187 } 188 189 peerNode := m.nodes[peer] 190 191 // Call the per channel callback. 192 err := cb( 193 &graphdb.DirectedChannel{ 194 ChannelID: channel.id, 195 IsNode1: nodePub == node1, 196 OtherNode: peer, 197 Capacity: channel.capacity, 198 OutPolicySet: true, 199 InPolicy: &models.CachedEdgePolicy{ 200 ChannelID: channel.id, 201 ToNodePubKey: func() route.Vertex { 202 return nodePub 203 }, 204 ToNodeFeatures: lnwire.EmptyFeatureVector(), 205 FeeBaseMSat: peerNode.baseFee, 206 }, 207 }, 208 ) 209 if err != nil { 210 return err 211 } 212 } 213 return nil 214 } 215 216 // sourceNode returns the source node of the graph. 217 // 218 // NOTE: Part of the Graph interface. 219 func (m *mockGraph) sourceNode() route.Vertex { 220 return m.source.pubkey 221 } 222 223 // fetchNodeFeatures returns the features of the given node. 224 // 225 // NOTE: Part of the Graph interface. 226 func (m *mockGraph) FetchNodeFeatures(_ context.Context, 227 _ route.Vertex) (*lnwire.FeatureVector, error) { 228 229 return lnwire.EmptyFeatureVector(), nil 230 } 231 232 // GraphSession will provide the call-back with access to a 233 // graphdb.NodeTraverser instance which can be used to perform queries against 234 // the channel graph. 235 // 236 // NOTE: Part of the GraphSessionFactory interface. 237 func (m *mockGraph) GraphSession(_ context.Context, 238 cb func(graph graphdb.NodeTraverser) error, _ func()) error { 239 240 return cb(m) 241 } 242 243 // htlcResult describes the resolution of an htlc. If failure is nil, the htlc 244 // was settled. 245 type htlcResult struct { 246 failureSource route.Vertex 247 failure lnwire.FailureMessage 248 } 249 250 // hop describes one hop of a route. 251 type hop struct { 252 node *mockNode 253 amtToFwd lnwire.MilliSatoshi 254 next *hop 255 } 256 257 // sendHtlc sends out an htlc on the mock network and synchronously returns the 258 // final resolution of the htlc. 259 func (m *mockGraph) sendHtlc(route *route.Route) (htlcResult, error) { 260 var next *hop 261 262 // Convert the route into a structure that is suitable for recursive 263 // processing. 264 for i := len(route.Hops) - 1; i >= 0; i-- { 265 routeHop := route.Hops[i] 266 node := m.nodes[routeHop.PubKeyBytes] 267 next = &hop{ 268 node: node, 269 next: next, 270 amtToFwd: routeHop.AmtToForward, 271 } 272 } 273 274 // Create the starting hop instance. 275 source := m.nodes[route.SourcePubKey] 276 next = &hop{ 277 node: source, 278 next: next, 279 amtToFwd: route.TotalAmount, 280 } 281 282 // Recursively walk the path and obtain the htlc resolution. 283 return source.fwd(nil, next) 284 } 285 286 // Compile-time check for the Graph interface. 287 var _ Graph = &mockGraph{}