memrelay.go raw
1 // Package memrelay is an in-process NIP-01 relay used by tests. It implements
2 // just REQ/EVENT/CLOSE with local filter matching, and can deterministically
3 // drop every Nth delivery to exercise a client's gap recovery.
4 package memrelay
5
6 import (
7 "errors"
8 "io"
9 "sync"
10
11 "git.smesh.lol/nostrino/pkg/nostr"
12 "git.smesh.lol/nostrino/pkg/relay"
13 )
14
15 type Relay struct {
16 mu sync.Mutex
17 conns map[*memTransport]map[string]nostr.Filter
18 dropEvery int
19 drop map[*memTransport]int
20 closed bool
21 }
22
23 func New() *Relay {
24 return &Relay{
25 conns: make(map[*memTransport]map[string]nostr.Filter),
26 drop: make(map[*memTransport]int),
27 }
28 }
29
30 func (r *Relay) Dial(string) (relay.Transport, error) {
31 r.mu.Lock()
32 defer r.mu.Unlock()
33 if r.closed {
34 return nil, errors.New("memrelay: closed")
35 }
36 t := &memTransport{
37 relay: r,
38 in: make(chan []byte, 512),
39 done: make(chan struct{}),
40 }
41 r.conns[t] = make(map[string]nostr.Filter)
42 return t, nil
43 }
44
45 func (r *Relay) SetDrop(n int) {
46 r.mu.Lock()
47 r.dropEvery = n
48 r.drop = make(map[*memTransport]int)
49 r.mu.Unlock()
50 }
51
52 func (r *Relay) Close() {
53 r.mu.Lock()
54 r.closed = true
55 conns := make([]*memTransport, 0, len(r.conns))
56 for c := range r.conns {
57 conns = append(conns, c)
58 }
59 r.mu.Unlock()
60 for _, c := range conns {
61 c.Close()
62 }
63 }
64
65 type memTransport struct {
66 relay *Relay
67 in chan []byte
68 done chan struct{}
69 once sync.Once
70 }
71
72 func (t *memTransport) WriteText(raw []byte) error {
73 parts, err := nostr.SplitArray(raw)
74 if err != nil || len(parts) == 0 {
75 return errors.New("memrelay: bad message")
76 }
77 typ, err := nostr.ParseJSONString(parts[0])
78 if err != nil {
79 return err
80 }
81 switch typ {
82 case "REQ":
83 if len(parts) < 3 {
84 return errors.New("memrelay: REQ")
85 }
86 id, err := nostr.ParseJSONString(parts[1])
87 if err != nil {
88 return err
89 }
90 f, err := nostr.ParseFilter(parts[2])
91 if err != nil {
92 return err
93 }
94 t.relay.mu.Lock()
95 if subs := t.relay.conns[t]; subs != nil {
96 subs[id] = *f
97 }
98 t.relay.mu.Unlock()
99 t.deliver([]byte(`["EOSE","` + id + `"]`))
100 return nil
101 case "CLOSE":
102 if len(parts) < 2 {
103 return errors.New("memrelay: CLOSE")
104 }
105 id, _ := nostr.ParseJSONString(parts[1])
106 t.relay.mu.Lock()
107 if subs := t.relay.conns[t]; subs != nil {
108 delete(subs, id)
109 }
110 t.relay.mu.Unlock()
111 return nil
112 case "EVENT":
113 if len(parts) < 2 {
114 return errors.New("memrelay: EVENT")
115 }
116 ev, err := nostr.ParseEvent(parts[1])
117 if err != nil {
118 return err
119 }
120 if err := ev.Verify(); err != nil {
121 return err
122 }
123 id := nostr.Hex(ev.ID[:])
124 t.deliver([]byte(`["OK","` + id + `",true,""]`))
125 t.relay.route(t, ev)
126 return nil
127 default:
128 return errors.New("memrelay: unsupported " + typ)
129 }
130 }
131
132 func (r *Relay) route(from *memTransport, ev *nostr.Event) {
133 type target struct {
134 t *memTransport
135 id string
136 }
137 r.mu.Lock()
138 var targets []target
139 drop := make(map[target]bool)
140 for t, subs := range r.conns {
141 for id, f := range subs {
142 if !f.Matches(ev) {
143 continue
144 }
145 tg := target{t: t, id: id}
146 if r.dropEvery > 0 {
147 r.drop[t]++
148 if r.drop[t]%r.dropEvery == 0 {
149 drop[tg] = true
150 }
151 }
152 targets = append(targets, tg)
153 }
154 }
155 r.mu.Unlock()
156 raw := ev.AppendJSON(nil)
157 for _, tg := range targets {
158 if drop[tg] {
159 continue
160 }
161 msg := make([]byte, 0, len(raw)+len(tg.id)+16)
162 msg = append(msg, `["EVENT","`...)
163 msg = append(msg, tg.id...)
164 msg = append(msg, `",`...)
165 msg = append(msg, raw...)
166 msg = append(msg, ']')
167 tg.t.deliver(msg)
168 }
169 _ = from
170 }
171
172 func (t *memTransport) deliver(msg []byte) {
173 select {
174 case t.in <- msg:
175 case <-t.done:
176 default:
177 // A real relay has a bounded send buffer; the client re-syncs or the
178 // tunnel retransmits, so drop rather than block the whole relay.
179 }
180 }
181
182 func (t *memTransport) ReadMessage() ([]byte, error) {
183 select {
184 case msg := <-t.in:
185 return msg, nil
186 case <-t.done:
187 return nil, io.EOF
188 }
189 }
190
191 func (t *memTransport) Close() error {
192 t.once.Do(func() {
193 close(t.done)
194 t.relay.mu.Lock()
195 delete(t.relay.conns, t)
196 t.relay.mu.Unlock()
197 })
198 return nil
199 }
200
201 var _ relay.Transport = (*memTransport)(nil)
202