diff --git a/watchtower/wtmock/peer.go b/watchtower/wtmock/peer.go index 9103c499..16aeb0d8 100644 --- a/watchtower/wtmock/peer.go +++ b/watchtower/wtmock/peer.go @@ -13,6 +13,8 @@ import ( type MockPeer struct { remotePub *btcec.PublicKey remoteAddr net.Addr + localPub *btcec.PublicKey + localAddr net.Addr IncomingMsgs chan []byte OutgoingMsgs chan []byte @@ -20,30 +22,71 @@ type MockPeer struct { writeDeadline <-chan time.Time readDeadline <-chan time.Time - Quit chan struct{} + RemoteQuit chan struct{} + Quit chan struct{} } // NewMockPeer returns a fresh MockPeer. func NewMockPeer(pk *btcec.PublicKey, addr net.Addr, bufferSize int) *MockPeer { return &MockPeer{ - remotePub: pk, - remoteAddr: addr, + remotePub: pk, + remoteAddr: addr, + localAddr: &net.TCPAddr{ + IP: net.IP{0x32, 0x31, 0x30, 0x29}, + Port: 36723, + }, IncomingMsgs: make(chan []byte, bufferSize), OutgoingMsgs: make(chan []byte, bufferSize), Quit: make(chan struct{}), } } +// NewMockConn establishes a bidirectional connection between two MockPeers. +func NewMockConn(localPk, remotePk *btcec.PublicKey, + localAddr, remoteAddr net.Addr, + bufferSize int) (*MockPeer, *MockPeer) { + + localPeer := &MockPeer{ + remotePub: remotePk, + remoteAddr: remoteAddr, + localPub: localPk, + localAddr: localAddr, + IncomingMsgs: make(chan []byte, bufferSize), + OutgoingMsgs: make(chan []byte, bufferSize), + Quit: make(chan struct{}), + } + + remotePeer := &MockPeer{ + remotePub: localPk, + remoteAddr: localAddr, + localPub: remotePk, + localAddr: remoteAddr, + IncomingMsgs: localPeer.OutgoingMsgs, + OutgoingMsgs: localPeer.IncomingMsgs, + Quit: make(chan struct{}), + } + + localPeer.RemoteQuit = remotePeer.Quit + remotePeer.RemoteQuit = localPeer.Quit + + return localPeer, remotePeer +} + // Write sends the raw bytes as the next full message read to the remote peer. // The write will fail if either party closes the connection or the write // deadline expires. The passed bytes slice is copied before sending, thus the // bytes may be reused once the method returns. func (p *MockPeer) Write(b []byte) (n int, err error) { + bb := make([]byte, len(b)) + copy(bb, b) + select { - case p.OutgoingMsgs <- b: + case p.OutgoingMsgs <- bb: return len(b), nil case <-p.writeDeadline: return 0, fmt.Errorf("write timeout expired") + case <-p.RemoteQuit: + return 0, fmt.Errorf("remote closed connected") case <-p.Quit: return 0, fmt.Errorf("connection closed") } @@ -69,6 +112,8 @@ func (p *MockPeer) ReadNextMessage() ([]byte, error) { return b, nil case <-p.readDeadline: return nil, fmt.Errorf("read timeout expired") + case <-p.RemoteQuit: + return nil, fmt.Errorf("remote closed connected") case <-p.Quit: return nil, fmt.Errorf("connection closed") } @@ -112,6 +157,25 @@ func (p *MockPeer) RemoteAddr() net.Addr { return p.remoteAddr } +// LocalAddr returns the local net address of the peer. +func (p *MockPeer) LocalAddr() net.Addr { + return p.localAddr +} + +// Read is not implemented. +func (p *MockPeer) Read(dst []byte) (int, error) { + panic("not implemented") +} + +// SetDeadline is not implemented. +func (p *MockPeer) SetDeadline(t time.Time) error { + panic("not implemented") +} + // Compile-time constraint ensuring the MockPeer implements the wserver.Peer // interface. var _ wtserver.Peer = (*MockPeer)(nil) + +// Compile-time constraint ensuring the MockPeer implements the net.Conn +// interface. +var _ net.Conn = (*MockPeer)(nil)