@@ -2,14 +2,15 @@ package zeroconf
22
33import (
44 "context"
5+ "fmt"
56 "log"
67 "testing"
78 "time"
89
910 "github.com/pkg/errors"
1011)
1112
12- var (
13+ const (
1314 mdnsName = "test--xxxxxxxxxxxx"
1415 mdnsService = "_test--xxxx._tcp"
1516 mdnsSubtype = "_test--xxxx._tcp,_fancy"
@@ -163,4 +164,67 @@ func TestSubtype(t *testing.T) {
163164 t .Fatalf ("Expected port is %d, but got %d" , mdnsPort , result .Port )
164165 }
165166 })
167+
168+ t .Run ("DoS protection" , func (t * testing.T ) {
169+ origMaxSentEntries := maxSentEntries
170+ maxSentEntries = 10
171+ defer func () { maxSentEntries = origMaxSentEntries }()
172+
173+ ctx , cancel := context .WithTimeout (context .Background (), 10 * time .Second )
174+ defer cancel ()
175+
176+ const firstName = mdnsName
177+
178+ go startMDNS (ctx , mdnsPort , firstName , mdnsSubtype , mdnsDomain )
179+ time .Sleep (time .Second )
180+
181+ resolver , err := NewResolver (nil )
182+ if err != nil {
183+ t .Fatalf ("Expected create resolver success, but got %v" , err )
184+ }
185+ entries := make (chan * ServiceEntry , maxSentEntries + 1 )
186+ received := make (chan * ServiceEntry , 10 )
187+ go func () {
188+ for {
189+ select {
190+ case entry := <- entries :
191+ if entry .Instance == firstName {
192+ received <- entry
193+ }
194+ case <- ctx .Done ():
195+ return
196+ }
197+ }
198+ }()
199+ if err := resolver .Browse (ctx , mdnsService , mdnsDomain , entries ); err != nil {
200+ t .Fatalf ("Expected browse success, but got %v" , err )
201+ }
202+ select {
203+ case <- received :
204+ case <- time .NewTimer (time .Second ).C :
205+ t .Fatal ("expected to discover service" )
206+ }
207+
208+ for i := 1 ; i < maxSentEntries ; i ++ {
209+ go startMDNS (ctx , mdnsPort , fmt .Sprintf ("%s-%d" , mdnsName , i ), mdnsSubtype , mdnsDomain )
210+ }
211+ time .Sleep (time .Second )
212+
213+ select {
214+ case entry := <- entries :
215+ t .Fatalf ("didn't expect to receive an entry, got %v" , entry )
216+ default :
217+ }
218+
219+ // Announcing this service will cause the map to overflow.
220+ go startMDNS (ctx , mdnsPort , fmt .Sprintf ("%s-%d" , mdnsName , maxSentEntries ), mdnsSubtype , mdnsDomain )
221+
222+ // wait for a re-announcement of the firstName service
223+ select {
224+ case <- received :
225+ cancel ()
226+ case <- ctx .Done ():
227+ t .Fatal ("expected to discover service" )
228+ }
229+ })
166230}
0 commit comments