1
//go:build !no_mocks
2
 
3
package mocks
4
 
5
import (
6
	"context"
7
	"crypto/tls"
8
	"crypto/x509"
9
	"errors"
10
	"fmt"
11
	"net"
12
	"sync"
13
 
14
	"go.miragespace.co/specter/spec/protocol"
15
	"go.miragespace.co/specter/spec/transport"
16
	"go.miragespace.co/specter/util/bufconn"
17
 
18
	"github.com/stretchr/testify/mock"
19
)
20
 
21
type MemoryTransport struct {
22
	Identify    *protocol.Node
23
	Other       chan *transport.StreamDelegate
24
	Self        chan *transport.StreamDelegate
25
	Physical    transport.PhysicalConn
26
	Certificate *x509.Certificate
27
	mock.Mock
28
}
29
 
30
var _ transport.Transport = (*MemoryTransport)(nil)
31
 
32
var _ transport.ClientTransport = (*MemoryTransport)(nil)
33
 
34
// SelfTransport returns a transport.Transport that when .DialStream() is invoked, .AcceptStream()
35
// on the same Transport will receive the net.Conn
36
func SelfTransport() *MemoryTransport {
37
	s := make(chan *transport.StreamDelegate, 1)
38
	t := &MemoryTransport{
39
		Other: s,
40
		Self:  s,
41
	}
42
	return t
43
}
44
 
45
// PipeTransport returns two transport.Transport that when one side's .DialStream() is invoked,
46
// .AcceptStream() on the other Transport will receive the net.Conn
47
func PipeTransport() (*MemoryTransport, *MemoryTransport) {
48
	s1 := make(chan *transport.StreamDelegate, 1)
49
	s2 := make(chan *transport.StreamDelegate, 1)
50
	t1 := &MemoryTransport{
51
		Other: s2,
52
		Self:  s1,
53
	}
54
	t2 := &MemoryTransport{
55
		Other: s1,
56
		Self:  s2,
57
	}
58
	return t1, t2
59
}
60
 
61
func (t *MemoryTransport) WithCertificate(cert *x509.Certificate) {
62
	t.Certificate = cert
63
}
64
 
65
func (t *MemoryTransport) WithClientCertificate(cert tls.Certificate) error {
66
	args := t.Called(cert)
67
	return args.Error(0)
68
}
69
 
70
func (t *MemoryTransport) Identity() *protocol.Node {
71
	return t.Identify
72
}
73
 
74
func (t *MemoryTransport) DialStream(ctx context.Context, peer *protocol.Node, kind protocol.Stream_Type) (net.Conn, error) {
75
	c1, c2 := bufconn.BufferedPipe(8192)
76
	if t.Physical != nil {
77
		c1 = &physicalPipe{c1, t.Physical}
78
		c2 = &physicalPipe{c2, t.Physical}
79
	}
80
	select {
81
	case t.Other <- &transport.StreamDelegate{
82
		Conn:        c1,
83
		Identity:    peer,
84
		Kind:        kind,
85
		Certificate: t.Certificate,
86
	}:
87
	default:
88
		panic(fmt.Sprintf("blocked on dialing %s", peer.String()))
89
	}
90
	return c2, nil
91
}
92
 
93
func (t *MemoryTransport) AcceptStream() <-chan *transport.StreamDelegate {
94
	return t.Self
95
}
96
 
97
func (t *MemoryTransport) ListConnected() []transport.ConnectedPeer {
98
	panic("not implemented") // TODO: Implement
99
}
100
 
101
func (t *MemoryTransport) SupportDatagram() bool {
102
	panic("not implemented") // TODO: Implement
103
}
104
 
105
func (t *MemoryTransport) ReceiveDatagram() <-chan *transport.DatagramDelegate {
106
	panic("not implemented") // TODO: Implement
107
}
108
 
109
func (t *MemoryTransport) SendDatagram(_ *protocol.Node, _ []byte) error {
110
	panic("not implemented") // TODO: Implement
111
}
112
 
113
// PhysicalConn is a controllable connection lifetime for session tests.
114
type PhysicalConn struct {
115
	deliver func(*transport.StreamDelegate)
116
	done    chan struct{}
117
	mu      sync.Mutex
118
	err     error
119
	streams []net.Conn
120
}
121
 
122
func NewPhysicalConn(deliver func(*transport.StreamDelegate)) *PhysicalConn {
123
	return &PhysicalConn{
124
		deliver: deliver,
125
		done:    make(chan struct{}),
126
	}
127
}
128
 
129
func (p *PhysicalConn) Done() <-chan struct{} { return p.done }
130
 
131
func (p *PhysicalConn) Err() error { p.mu.Lock(); defer p.mu.Unlock(); return p.err }
132
 
133
func (p *PhysicalConn) Close(reason string) error {
134
	p.mu.Lock()
135
	defer p.mu.Unlock()
136
	if p.err != nil {
137
		return nil
138
	}
139
	p.err = errors.New(reason)
140
	close(p.done)
141
	for _, c := range p.streams {
142
		c.Close()
143
	}
144
	p.streams = nil
145
	return nil
146
}
147
 
148
func (p *PhysicalConn) OpenStream(kind protocol.Stream_Type) (net.Conn, error) {
149
	p.mu.Lock()
150
	if p.err != nil {
151
		p.mu.Unlock()
152
		return nil, transport.ErrClosed
153
	}
154
	c1, c2 := bufconn.BufferedPipe(8192)
155
	p.streams = append(p.streams, c1, c2)
156
	p.mu.Unlock()
157
	p.deliver(&transport.StreamDelegate{
158
		Conn: &physicalPipe{c1, p},
159
		Kind: kind,
160
	})
161
	return &physicalPipe{c2, p}, nil
162
}
163
 
164
type physicalPipe struct {
165
	net.Conn
166
	pc transport.PhysicalConn
167
}
168
 
169
func (p *physicalPipe) PhysicalConn() transport.PhysicalConn { return p.pc }