@@ -2,6 +2,8 @@ package mail
22
33import (
44 "net"
5+ "reflect"
6+ "strings"
57 "testing"
68
79 mdns "github.com/miekg/dns"
@@ -87,6 +89,164 @@ func TestQueryTCPFallbackOnTruncation(t *testing.T) {
8789 }
8890}
8991
92+ func TestMergeDKIMSelectors (t * testing.T ) {
93+ tests := []struct {
94+ name string
95+ user []string
96+ defaults []string
97+ want []string
98+ }{
99+ {
100+ name : "no user selectors keeps defaults in order" ,
101+ defaults : []string {"default" , "google" , "selector1" },
102+ want : []string {"default" , "google" , "selector1" },
103+ },
104+ {
105+ name : "user selectors come first" ,
106+ user : []string {"foo" , "bar" },
107+ defaults : []string {"default" , "google" },
108+ want : []string {"foo" , "bar" , "default" , "google" },
109+ },
110+ {
111+ name : "dedupes case-insensitively, preserving user spelling" ,
112+ user : []string {"Google" , "custom" },
113+ defaults : []string {"default" , "google" },
114+ want : []string {"Google" , "custom" , "default" },
115+ },
116+ {
117+ name : "trims whitespace and drops empty entries" ,
118+ user : []string {" foo " , "" , " " },
119+ defaults : []string {"default" },
120+ want : []string {"foo" , "default" },
121+ },
122+ }
123+
124+ for _ , tc := range tests {
125+ t .Run (tc .name , func (t * testing.T ) {
126+ got := mergeDKIMSelectors (tc .user , tc .defaults )
127+ if ! reflect .DeepEqual (got , tc .want ) {
128+ t .Errorf ("mergeDKIMSelectors(%v, %v) = %v, want %v" , tc .user , tc .defaults , got , tc .want )
129+ }
130+ })
131+ }
132+ }
133+
134+ // startDKIMServer spins up a UDP DNS server that returns a v=DKIM1 TXT record
135+ // for each selector in records, and NXDOMAIN otherwise.
136+ func startDKIMServer (t * testing.T , domain string , records map [string ]string ) string {
137+ t .Helper ()
138+
139+ conn , err := net .ListenPacket ("udp" , "127.0.0.1:0" )
140+ if err != nil {
141+ t .Fatal (err )
142+ }
143+
144+ handler := mdns .HandlerFunc (func (w mdns.ResponseWriter , r * mdns.Msg ) {
145+ m := new (mdns.Msg )
146+ m .SetReply (r )
147+ q := r .Question [0 ]
148+ suffix := "._domainkey." + domain + "."
149+ if q .Qtype == mdns .TypeTXT && strings .HasSuffix (q .Name , suffix ) {
150+ sel := strings .TrimSuffix (q .Name , suffix )
151+ if val , ok := records [sel ]; ok {
152+ m .Answer = append (m .Answer , & mdns.TXT {
153+ Hdr : mdns.RR_Header {Name : q .Name , Rrtype : mdns .TypeTXT , Class : mdns .ClassINET , Ttl : 60 },
154+ Txt : []string {val },
155+ })
156+ } else {
157+ m .Rcode = mdns .RcodeNameError
158+ }
159+ }
160+ _ = w .WriteMsg (m )
161+ })
162+
163+ server := & mdns.Server {PacketConn : conn , Handler : handler , Net : "udp" }
164+ go func () { _ = server .ActivateAndServe () }()
165+ t .Cleanup (func () { _ = server .Shutdown () })
166+
167+ return conn .LocalAddr ().String ()
168+ }
169+
170+ func TestLookupDKIMUserSelectorHit (t * testing.T ) {
171+ domain := "example.com"
172+ resolver := startDKIMServer (t , domain , map [string ]string {
173+ "custom" : "v=DKIM1; k=rsa; p=PUBKEY" ,
174+ })
175+
176+ got := lookupDKIM (domain , []string {"custom" , "default" }, resolver )
177+ if len (got ) != 1 {
178+ t .Fatalf ("expected 1 DKIM record, got %d" , len (got ))
179+ }
180+ if got [0 ].Selector != "custom" {
181+ t .Errorf ("selector: got %q, want %q" , got [0 ].Selector , "custom" )
182+ }
183+ if ! strings .Contains (strings .ToLower (got [0 ].Value ), "v=dkim1" ) {
184+ t .Errorf ("value missing v=dkim1: %q" , got [0 ].Value )
185+ }
186+ }
187+
188+ func TestLookupDKIMPreservesSelectorOrder (t * testing.T ) {
189+ domain := "example.com"
190+ resolver := startDKIMServer (t , domain , map [string ]string {
191+ "google" : "v=DKIM1; k=rsa; p=A" ,
192+ "custom" : "v=DKIM1; k=rsa; p=B" ,
193+ "default" : "v=DKIM1; k=rsa; p=C" ,
194+ })
195+
196+ // Order: user-supplied "custom" first, then built-ins "default", "google".
197+ selectors := mergeDKIMSelectors ([]string {"custom" }, []string {"default" , "google" })
198+ got := lookupDKIM (domain , selectors , resolver )
199+
200+ want := []string {"custom" , "default" , "google" }
201+ if len (got ) != len (want ) {
202+ t .Fatalf ("expected %d records, got %d" , len (want ), len (got ))
203+ }
204+ for i , sel := range want {
205+ if got [i ].Selector != sel {
206+ t .Errorf ("position %d: got %q, want %q" , i , got [i ].Selector , sel )
207+ }
208+ }
209+ }
210+
211+ func TestSelectorsFromEnv (t * testing.T ) {
212+ tests := []struct {
213+ name string
214+ env string
215+ want []string
216+ }{
217+ {name : "unset" , env : "" , want : nil },
218+ {name : "single value" , env : "foo" , want : []string {"foo" }},
219+ {name : "comma separated" , env : "foo,bar,baz" , want : []string {"foo" , "bar" , "baz" }},
220+ {name : "trims whitespace and drops empties" , env : " foo , , bar ," , want : []string {"foo" , "bar" }},
221+ }
222+
223+ for _ , tc := range tests {
224+ t .Run (tc .name , func (t * testing.T ) {
225+ t .Setenv (DKIMSelectorsEnvVar , tc .env )
226+ got := selectorsFromEnv ()
227+ if ! reflect .DeepEqual (got , tc .want ) {
228+ t .Errorf ("selectorsFromEnv() = %v, want %v" , got , tc .want )
229+ }
230+ })
231+ }
232+ }
233+
234+ func TestLookupDKIMSkipsNonDKIMTXT (t * testing.T ) {
235+ domain := "example.com"
236+ resolver := startDKIMServer (t , domain , map [string ]string {
237+ "foo" : "not a dkim record" ,
238+ "bar" : "v=DKIM1; k=rsa; p=PUBKEY" ,
239+ })
240+
241+ got := lookupDKIM (domain , []string {"foo" , "bar" }, resolver )
242+ if len (got ) != 1 {
243+ t .Fatalf ("expected 1 DKIM record, got %d" , len (got ))
244+ }
245+ if got [0 ].Selector != "bar" {
246+ t .Errorf ("selector: got %q, want %q" , got [0 ].Selector , "bar" )
247+ }
248+ }
249+
90250func TestQueryUDPSuccessSkipsTCPFallback (t * testing.T ) {
91251 domain := "small.example."
92252 spf := "v=spf1 -all"
0 commit comments