Skip to content

Commit

Permalink
feat(netsim): start sketching out scenario (#22)
Browse files Browse the repository at this point in the history
The scenario allows us to express integration test more compactly and
less redundantly.
  • Loading branch information
bassosimone authored Nov 23, 2024
1 parent 6ef2b24 commit 11ef9d9
Show file tree
Hide file tree
Showing 8 changed files with 406 additions and 69 deletions.
177 changes: 177 additions & 0 deletions netsim/dns.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,177 @@
// SPDX-License-Identifier: GPL-3.0-or-later

package netsim

import (
"net"

"github.com/miekg/dns"
"github.com/rbmk-project/common/runtimex"
"github.com/rbmk-project/dnscore/dnscoretest"
)

// DNSHandler is an alias for dnscoretest.DNSHandler.
type DNSHandler = dnscoretest.Handler

// dnsDatabase is the global DNS database.
type dnsDatabase struct {
names map[string][]dns.RR
}

// newDNSDatabase creates a new DNS database.
func newDNSDatabase() *dnsDatabase {
return &dnsDatabase{
names: make(map[string][]dns.RR),
}
}

// AddCNAME adds a CNAME alias.
//
// This method IS NOT goroutine safe.
func (dd *dnsDatabase) AddCNAME(name, alias string) {
header := dns.RR_Header{
Name: dns.CanonicalName(name),
Rrtype: dns.TypeCNAME,
Class: dns.ClassINET,
Ttl: 3600,
Rdlength: 0,
}

rr := &dns.CNAME{
Hdr: header,
Target: dns.CanonicalName(alias),
}

dd.names[name] = append(dd.names[name], rr)
}

// AddFromConfig adds DNS records from a [*StackConfig].
//
// This method IS NOT goroutine safe.
func (dd *dnsDatabase) AddFromConfig(cfg *StackConfig) {
for _, name := range cfg.DomainNames {
name = dns.CanonicalName(name)
for _, addr := range cfg.Addresses {
// Make sure the string is a valid IP address
ipAddr := net.ParseIP(addr)
runtimex.Assert(ipAddr != nil, "invalid IP address")

// Create the common DNS header
header := dns.RR_Header{
Name: dns.CanonicalName(name),
Rrtype: 0,
Class: dns.ClassINET,
Ttl: 3600,
Rdlength: 0,
}

// Create the DNS record to add
var rr dns.RR
switch ipAddr.To4() {
case nil:
header.Rrtype = dns.TypeAAAA
rr = &dns.AAAA{Hdr: header, AAAA: ipAddr}
default:
header.Rrtype = dns.TypeA
rr = &dns.A{Hdr: header, A: ipAddr}
}

dd.names[name] = append(dd.names[name], rr)
}
}
}

// Ensure [*dnsDatabase] implements [dnsHandler].
var _ DNSHandler = (*dnsDatabase)(nil)

// Handler implements [dnsHandler] using [*dnsDatabase].
//
// This method is goroutine safe as long as one does not
// modify the database while handling queries.
func (dd *dnsDatabase) Handle(rw dnscoretest.ResponseWriter, rawQuery []byte) {
// Parse the incoming query and make sure it's a
// query containing just one question.
var (
response = &dns.Msg{}
query = &dns.Msg{}
)
if err := query.Unpack(rawQuery); err != nil {
return
}
if query.Response || query.Opcode != dns.OpcodeQuery || len(query.Question) != 1 {
return
}
response.SetReply(query)

// Get the RRs if possible
var (
q0 = query.Question[0]
name = dns.CanonicalName(q0.Name)
)
switch {
case q0.Qclass != dns.ClassINET:
response.Rcode = dns.RcodeRefused
case q0.Qtype == dns.TypeA ||
q0.Qtype == dns.TypeAAAA ||
q0.Qtype == dns.TypeCNAME:
var found bool
response.Answer, found = dd.lookup(q0.Qtype, name)
if !found {
response.Rcode = dns.RcodeNameError
}
default:
response.Rcode = dns.RcodeNameError
}

// Write the response
rawResp, err := response.Pack()
if err != nil {
return
}
rw.Write(rawResp)
}

// lookup returns the DNS records for a domain name.
//
// This method is goroutine safe as long as one does not
// modify the database while handling queries.
func (dd *dnsDatabase) lookup(qtype uint16, name string) ([]dns.RR, bool) {
const maxloops = 10
var rrs []dns.RR
for idx := 0; idx < maxloops; idx++ {

// Search whether the current name is in the database.
var interim []dns.RR
interim, found := dd.names[name]
if !found {
return nil, false
}

// We have definitely found something related.
rrs = append(rrs, interim...)

// Check whether we have found the desired record.
for _, rr := range interim {
if qtype == rr.Header().Rrtype {
return rrs, true
}
}

// Otherwise, follow CNAME redirects.
var cname string
for _, rr := range interim {
if rr, ok := rr.(*dns.CNAME); ok {
cname = rr.Target
break
}
}
if cname == "" {
return nil, false
}

// Continue searching from the CNAME target.
name = cname
}

return nil, false
}
84 changes: 27 additions & 57 deletions netsim/example_dns_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -6,78 +6,51 @@ import (
"context"
"fmt"
"log"
"net"
"net/netip"
"time"

"github.com/miekg/dns"
"github.com/rbmk-project/x/connpool"
"github.com/rbmk-project/x/netsim"
)

// This example shows how to use [netsim] to simulate a DNS
// server that listens for incoming requests over UDP.
func Example_dnsOverUDP() {
// Create a pool to close resources when done.
cpool := connpool.New()
defer cpool.Close()

// Create the server stack.
serverAddr := netip.MustParseAddr("8.8.8.8")
serverStack := netsim.NewStack(serverAddr)
cpool.Add(serverStack)

// Create the client stack.
clientAddr := netip.MustParseAddr("130.192.91.211")
clientStack := netsim.NewStack(clientAddr)
cpool.Add(clientStack)

// Link the client and the server stacks.
link := netsim.NewLink(clientStack, serverStack)
cpool.Add(link)
// Create a new scenario using the given directory to cache
// the certificates used by the simulated PKI
scenario := netsim.NewScenario("testdata")
defer scenario.Close()

// Create server stack running a DNS-over-UDP server.
//
// This includes:
//
// 1. creating, attaching, and enabling routing for a server stack
//
// 2. registering the proper domain names and addresses
//
// 3. updating the PKI database to include the server's certificate
scenario.Attach(scenario.MustNewStack(&netsim.StackConfig{
DomainNames: []string{"dns.google"},
Addresses: []string{"8.8.8.8"},
DNSOverUDPHandler: scenario.DNSHandler(),
}))

// Create and attach the client stack.
clientStack := scenario.MustNewStack(&netsim.StackConfig{
Addresses: []string{"130.192.91.211"},
})
scenario.Attach(clientStack)

// Create a context with a watchdog timeout.
ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second)
defer cancel()

// Create the server UDP listener.
serverEndpoint := netip.AddrPortFrom(serverAddr, 53)
serverConn, err := serverStack.ListenPacket(ctx, "udp", serverEndpoint.String())
if err != nil {
log.Fatal(err)
}
cpool.Add(serverConn)

// Start the server in the background.
serverDNS := &dns.Server{
PacketConn: serverConn,
Handler: dns.HandlerFunc(func(rw dns.ResponseWriter, query *dns.Msg) {
resp := &dns.Msg{}
resp.SetReply(query)
resp.Answer = append(resp.Answer, &dns.A{
Hdr: dns.RR_Header{
Name: "dns.google.",
Rrtype: dns.TypeA,
Class: dns.ClassINET,
Ttl: 3600,
Rdlength: 0,
},
A: net.IPv4(8, 8, 8, 8),
})
if err := rw.WriteMsg(resp); err != nil {
log.Fatal(err)
}
}),
}
go serverDNS.ActivateAndServe()
defer serverDNS.Shutdown()

// Create the client connection with the DNS server.
conn, err := clientStack.DialContext(ctx, "udp", serverEndpoint.String())
conn, err := clientStack.DialContext(ctx, "udp", "8.8.8.8:53")
if err != nil {
log.Fatal(err)
}
cpool.Add(conn)
defer conn.Close()

// Create the query to send
query := new(dns.Msg)
Expand All @@ -103,9 +76,6 @@ func Example_dnsOverUDP() {
}
}

// Explicitly close the connections
cpool.Close()

// Output:
// 8.8.8.8
}
4 changes: 2 additions & 2 deletions netsim/example_https_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -44,8 +44,8 @@ func Example_https() {
defer cancel()

// Create a PKI for the server and obtain the certificate.
pki := simpki.MustNewPKI("testdata")
serverCert := pki.MustNewCert(&simpki.PKICertConfig{
pki := simpki.MustNew("testdata")
serverCert := pki.MustNewCert(&simpki.Config{
CommonName: "dns.google",
DNSNames: []string{
"dns.google.com",
Expand Down
4 changes: 2 additions & 2 deletions netsim/example_router_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -49,8 +49,8 @@ func Example_router() {
defer cancel()

// Create a PKI for the server and obtain the certificate.
pki := simpki.MustNewPKI("testdata")
serverCert := pki.MustNewCert(&simpki.PKICertConfig{
pki := simpki.MustNew("testdata")
serverCert := pki.MustNewCert(&simpki.Config{
CommonName: "dns.google",
DNSNames: []string{
"dns.google.com",
Expand Down
4 changes: 2 additions & 2 deletions netsim/example_tls_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -50,8 +50,8 @@ func Example_tls() {
cpool.Add(listener)

// Create a PKI for the server and obtain the certificate.
pki := simpki.MustNewPKI("testdata")
serverCert := pki.MustNewCert(&simpki.PKICertConfig{
pki := simpki.MustNew("testdata")
serverCert := pki.MustNewCert(&simpki.Config{
CommonName: "dns.google",
DNSNames: []string{
"dns.google.com",
Expand Down
Loading

0 comments on commit 11ef9d9

Please sign in to comment.