feat(DNSConfig): allow to remove entries

Part of ooni/probe#2029

dnsserver.go

@@ -67,8 +67,8 @@ func (ds *DNSServer) Close() error {
 	return nil
 }
 
-// dnsRecord is a DNS record in the [DNSConfig].
-type dnsRecord struct {
+// DNSRecord is a DNS record in the [DNSConfig].
+type DNSRecord struct {
 	// A is the A resource record.
 	A []net.IP
 
@@ -80,14 +80,14 @@ type dnsRecord struct {
 // value is invalid; please use [NewDNSConfig].
 type DNSConfig struct {
 	mu sync.Mutex
-	r  map[string]*dnsRecord
+	r  map[string]*DNSRecord
 }
 
 // NewDNSConfig constructs a [DNSConfig] instance.
 func NewDNSConfig() *DNSConfig {
 	return &DNSConfig{
 		mu: sync.Mutex{},
-		r:  map[string]*dnsRecord{},
+		r:  map[string]*DNSRecord{},
 	}
 }
 
@@ -108,19 +108,27 @@ func (dc *DNSConfig) AddRecord(domain string, cname string, addrs ...string) err
 		cname = dns.CanonicalName(cname)
 	}
 	dc.mu.Lock()
-	dc.r[dns.CanonicalName(domain)] = &dnsRecord{
+	dc.r[dns.CanonicalName(domain)] = &DNSRecord{
 		A:     a,
 		CNAME: cname,
 	}
 	dc.mu.Unlock()
 	return nil
 }
 
+// RemoveRecord removes a record from the DNS server's database. If the record
+// does not exist, this method does nothing.
+func (dc *DNSConfig) RemoveRecord(domain string) {
+	dc.mu.Lock()
+	delete(dc.r, dns.CanonicalName(domain))
+	dc.mu.Unlock()
+}
+
 // Lookup searches a name inside the [DNSConfig].
-func (dc *DNSConfig) Lookup(name string) (*dnsRecord, bool) {
+func (dc *DNSConfig) Lookup(name string) (*DNSRecord, bool) {
 	defer dc.mu.Unlock()
 	dc.mu.Lock()
-	record, found := dc.r[name]
+	record, found := dc.r[dns.CanonicalName(name)]
 	return record, found
 }
 
@@ -193,7 +201,7 @@ func DNSServerRoundTrip(config *DNSConfig, rawQuery []byte) ([]byte, error) {
 }
 
 // dnsServerNewSuccessfulResponse constructs a successful response.
-func dnsServerNewSuccessfulResponse(query *dns.Msg, q0 dns.Question, rr *dnsRecord) ([]byte, error) {
+func dnsServerNewSuccessfulResponse(query *dns.Msg, q0 dns.Question, rr *DNSRecord) ([]byte, error) {
 	// fill the response
 	resp := &dns.Msg{}
 	resp.SetReply(query)

dnsserver_test.go

@@ -0,0 +1,50 @@
+package netem
+
+import (
+	"net"
+	"testing"
+
+	"github.com/google/go-cmp/cmp"
+)
+
+func TestDNSConfig(t *testing.T) {
+	t.Run("removing a nonexisting record does not cause any issue", func(t *testing.T) {
+		dc := NewDNSConfig()
+		dc.RemoveRecord("www.example.com")
+	})
+
+	t.Run("we can remove a previously added record", func(t *testing.T) {
+		dc := NewDNSConfig()
+
+		t.Run("the record should be there once we have added it", func(t *testing.T) {
+			if err := dc.AddRecord("www.example.com", "www1.example.com", "1.2.3.4", "4.5.6.7"); err != nil {
+				t.Fatal(err)
+			}
+			rec, good := dc.Lookup("www.example.com")
+			if !good {
+				t.Fatal("the record is not there")
+			}
+			expect := &DNSRecord{
+				A: []net.IP{
+					net.IPv4(1, 2, 3, 4),
+					net.IPv4(4, 5, 6, 7),
+				},
+				CNAME: "www1.example.com.",
+			}
+			if diff := cmp.Diff(expect, rec); diff != "" {
+				t.Fatal(diff)
+			}
+
+			t.Run("the record should disappear once we have removed it", func(t *testing.T) {
+				dc.RemoveRecord("www.example.com")
+				rec, good := dc.Lookup("www.example.com")
+				if good {
+					t.Fatal("expected the record to be nonexistent")
+				}
+				if rec != nil {
+					t.Fatal("expected a nil record")
+				}
+			})
+		})
+	})
+}

dpiblock.go

@@ -236,7 +236,7 @@ func (r *DPISpoofDNSResponse) Filter(
 	}
 
 	// create a DNS record for preparing a response
-	dnsRecord := &dnsRecord{
+	dnsRecord := &DNSRecord{
 		A:     []net.IP{},
 		CNAME: "",
 	}