Add fragmentation to DNS registrar

pkg/registrars/dns-registrar/examples/client/main.go

@@ -0,0 +1,61 @@
+package main
+
+import (
+	"bufio"
+	"encoding/hex"
+	"flag"
+	"fmt"
+	"net"
+	"os"
+	"strings"
+
+	"github.com/pion/dtls/v2/examples/util"
+	"github.com/refraction-networking/conjure/pkg/registrars/dns-registrar/requester"
+	"github.com/refraction-networking/conjure/pkg/registrars/dns-registrar/tworeqresp"
+)
+
+const key = "0b63baad7f2f4bb5b547c53adc0fbb179852910607935e6f4b5639fd989b1156"
+
+func main() {
+	var remoteAddr = flag.String("saddr", "127.0.0.1:6666", "remote address")
+	var baseDomain = flag.String("domain", "test.xyz", "base domain to use")
+	flag.Parse()
+
+	addr, err := net.ResolveUDPAddr("udp", *remoteAddr)
+	util.Check(err)
+
+	pubKey, err := hex.DecodeString(key)
+	util.Check(err)
+
+	req, err := requester.NewRequester(&requester.Config{
+		TransportMethod: requester.UDP,
+		Target:          addr.String(),
+		BaseDomain:      *baseDomain,
+		Pubkey:          pubKey,
+	})
+
+	util.Check(err)
+
+	tworeq, err := tworeqresp.NewRequester(req, 80)
+	util.Check(err)
+
+	reader := bufio.NewReader(os.Stdin)
+
+	fmt.Println("type 'exit' to shutdown gracefully")
+
+	for {
+		text, err := reader.ReadString('\n')
+		util.Check(err)
+
+		if strings.TrimSpace(text) == "exit" {
+			return
+		}
+
+		resp, err := tworeq.RequestAndRecv([]byte(text))
+		util.Check(err)
+
+		fmt.Printf("Got message: %s\n", string(resp))
+
+	}
+
+}

pkg/registrars/dns-registrar/examples/server/main.go

@@ -0,0 +1,42 @@
+package main
+
+import (
+	"encoding/hex"
+	"flag"
+	"fmt"
+	"net"
+
+	"github.com/pion/dtls/v2/examples/util"
+	"github.com/refraction-networking/conjure/pkg/registrars/dns-registrar/responder"
+	"github.com/refraction-networking/conjure/pkg/registrars/dns-registrar/tworeqresp"
+)
+
+const key = "203963feed62ddda89b98857940f09866ae840f42e8c90160e411a0029b87e60"
+
+func main() {
+	var localAddr = flag.String("laddr", "[::]:6666", "source address")
+	var domain = flag.String("domain", "test.xyz", "domain to use")
+	var msg = flag.String("msg", "hey", "message to respond")
+	flag.Parse()
+
+	privKey, err := hex.DecodeString(key)
+	util.Check(err)
+
+	// Prepare the IP to connect to
+	laddr, err := net.ResolveUDPAddr("udp", *localAddr)
+	util.Check(err)
+
+	responder, err := responder.NewDnsResponder(*domain, laddr.String(), privKey)
+	util.Check(err)
+
+	tworesponder, err := tworeqresp.NewResponder(responder)
+	util.Check(err)
+
+	fmt.Println("Listening")
+
+	tworesponder.RecvAndRespond(func(b []byte) ([]byte, error) {
+		fmt.Println(string(b))
+		return []byte(*msg), nil
+	})
+
+}

pkg/registrars/dns-registrar/tworeqresp/requester.go

@@ -0,0 +1,94 @@
+package tworeqresp
+
+import (
+	"context"
+	"crypto/rand"
+	"fmt"
+	"net"
+
+	pb "github.com/refraction-networking/conjure/proto"
+	"google.golang.org/protobuf/proto"
+)
+
+type dialFunc = func(ctx context.Context, network, addr string) (net.Conn, error)
+
+const idLen = 8
+
+type onerequester interface {
+	RequestAndRecv(sendBytes []byte) ([]byte, error)
+	Close() error
+	SetDialer(dialer dialFunc) error
+}
+
+type Requester struct {
+	parent onerequester
+	mtu    uint
+}
+
+func NewRequester(parent onerequester, mtu uint) (*Requester, error) {
+	return &Requester{parent: parent, mtu: mtu}, nil
+}
+
+func (r *Requester) RequestAndRecv(sendBytes []byte) ([]byte, error) {
+
+	id := [idLen]byte{}
+	_, err := rand.Read(id[:])
+	if err != nil {
+		return nil, fmt.Errorf("error generating id: %v", err)
+	}
+
+	parts := splitIntoChunks(sendBytes, int(r.mtu))
+
+	for i, partBytes := range parts {
+		toSend := &pb.DnsPartReq{Id: id[:], PartNum: proto.Uint32(uint32(i)), TotalParts: proto.Uint32(uint32(len(parts))), Data: partBytes}
+		toSendBytes, err := proto.Marshal(toSend)
+		if err != nil {
+			return nil, fmt.Errorf("error marshal part %v: %v", i, err)
+		}
+
+		respBytes, err := r.parent.RequestAndRecv(toSendBytes)
+		if err != nil {
+			return nil, fmt.Errorf("error request part %v: %v", i, err)
+		}
+
+		resp := &pb.DnsPartResp{}
+		err = proto.Unmarshal(respBytes, resp)
+		if err != nil {
+			return nil, fmt.Errorf("error unmarshal response: %v", err)
+		}
+
+		if resp.GetWaiting() {
+			continue
+		}
+
+		return resp.GetData(), nil
+	}
+
+	return nil, fmt.Errorf("no response")
+}
+
+func splitIntoChunks(data []byte, mtu int) [][]byte {
+	var chunks [][]byte
+
+	for i := 0; i < len(data); i += mtu {
+		end := i + mtu
+
+		if end > len(data) {
+			end = len(data)
+		}
+
+		chunks = append(chunks, data[i:end])
+	}
+
+	return chunks
+}
+
+// Close closes the parent transport
+func (r *Requester) Close() error {
+	return r.parent.Close()
+}
+
+// SetDialer sets the parent dialer
+func (r *Requester) SetDialer(dialer dialFunc) error {
+	return r.parent.SetDialer(dialer)
+}

pkg/registrars/dns-registrar/tworeqresp/responder.go

@@ -0,0 +1,99 @@
+package tworeqresp
+
+import (
+	"bytes"
+	"fmt"
+
+	pb "github.com/refraction-networking/conjure/proto"
+	"google.golang.org/protobuf/proto"
+)
+
+type oneresponder interface {
+	RecvAndRespond(getResponse func([]byte) ([]byte, error)) error
+	Close() error
+}
+
+type Responder struct {
+	parent oneresponder
+	parts  map[[idLen]byte][][]byte
+}
+
+func NewResponder(parent oneresponder) (*Responder, error) {
+	return &Responder{
+		parent: parent,
+		parts:  make(map[[idLen]byte][][]byte),
+	}, nil
+}
+
+func (r *Responder) RecvAndRespond(parentGetResponse func([]byte) ([]byte, error)) error {
+	getResponse := func(data []byte) ([]byte, error) {
+		partIn := &pb.DnsPartReq{}
+		err := proto.Unmarshal(data, partIn)
+		if err != nil {
+			return nil, fmt.Errorf("error umarshal part: %v", err)
+		}
+
+		if len(partIn.GetId()) != idLen {
+			return nil, fmt.Errorf("invalid part ID")
+		}
+
+		partId := (*[idLen]byte)(partIn.GetId())
+
+		if _, ok := r.parts[*partId]; !ok {
+			r.parts[*partId] = make([][]byte, partIn.GetTotalParts())
+		}
+
+		if int(partIn.GetTotalParts()) != len(r.parts[*partId]) {
+			return nil, fmt.Errorf("invalid total parts")
+		}
+
+		if int(partIn.GetPartNum()) >= len(r.parts[*partId]) {
+			return nil, fmt.Errorf("part number out of bound")
+		}
+
+		r.parts[*partId][partIn.GetPartNum()] = partIn.GetData()
+
+		waiting := false
+		for _, part := range r.parts[*partId] {
+			if part == nil {
+				waiting = true
+				break
+			}
+		}
+
+		if waiting {
+			resp := &pb.DnsPartResp{Waiting: proto.Bool(true)}
+			respBytes, err := proto.Marshal(resp)
+			if err != nil {
+				return nil, fmt.Errorf("error marshal resp: %v", err)
+			}
+
+			return respBytes, nil
+		}
+
+		var buffer bytes.Buffer
+		for _, part := range r.parts[*partId] {
+			buffer.Write(part)
+		}
+		res, err := parentGetResponse(buffer.Bytes())
+		if err != nil {
+			return nil, fmt.Errorf("error from parent getResponse: %v", err)
+		}
+
+		resp := &pb.DnsPartResp{Waiting: proto.Bool(false), Data: res}
+
+		respBytes, err := proto.Marshal(resp)
+		if err != nil {
+			return nil, fmt.Errorf("error marshal resp: %v", err)
+		}
+
+		return respBytes, nil
+
+	}
+	return r.parent.RecvAndRespond(getResponse)
+}
+
+// Close closes the parent transport
+func (r *Responder) Close() error {
+	return r.parent.Close()
+}

pkg/registrars/registration/dns-registrar.go

@@ -10,6 +10,7 @@ import (
 
 	"github.com/pion/stun"
 	"github.com/refraction-networking/conjure/pkg/registrars/dns-registrar/requester"
+	"github.com/refraction-networking/conjure/pkg/registrars/dns-registrar/tworeqresp"
 	"github.com/refraction-networking/conjure/pkg/registrars/lib"
 	pb "github.com/refraction-networking/conjure/proto"
 	"github.com/refraction-networking/gotapdance/tapdance"
@@ -18,7 +19,7 @@ import (
 )
 
 type DNSRegistrar struct {
-	req             *requester.Requester
+	req             *tworeqresp.Requester
 	maxRetries      int
 	connectionDelay time.Duration
 	bidirectional   bool
@@ -63,13 +64,18 @@ func NewDNSRegistrar(config *Config) (*DNSRegistrar, error) {
 		return nil, fmt.Errorf("error creating requester: %v", err)
 	}
 
+	tworeq, err := tworeqresp.NewRequester(req, 80)
+	if err != nil {
+		return nil, fmt.Errorf("error adding fragmentation layer: %v", err)
+	}
+
 	ip, err := getPublicIp(config.STUNAddr)
 	if err != nil {
 		return nil, fmt.Errorf("failed to get public IP: %v", err)
 	}
 
 	return &DNSRegistrar{
-		req:             req,
+		req:             tworeq,
 		ip:              ip,
 		maxRetries:      config.MaxRetries,
 		bidirectional:   config.Bidirectional,

pkg/regserver/dnsregserver/dnsregserver.go

@@ -8,6 +8,7 @@ import (
 
 	"github.com/refraction-networking/conjure/pkg/metrics"
 	"github.com/refraction-networking/conjure/pkg/registrars/dns-registrar/responder"
+	"github.com/refraction-networking/conjure/pkg/registrars/dns-registrar/tworeqresp"
 	"github.com/refraction-networking/conjure/pkg/regserver/regprocessor"
 	pb "github.com/refraction-networking/conjure/proto"
 	log "github.com/sirupsen/logrus"
@@ -22,7 +23,7 @@ type registrar interface {
 // DNSRegServer provides an interface to forward DNS registration requests. Use a dns responder to receive requests and send responses.
 type DNSRegServer struct {
 	// dns responder to recieve and forward responses with
-	dnsResponder *responder.Responder
+	dnsResponder *tworeqresp.Responder
 	processor    registrar
 	latestCCGen  uint32
 	logger       log.FieldLogger
@@ -45,8 +46,13 @@ func NewDNSRegServer(domain string, udpAddr string, privkey []byte, regprocessor
 		return nil, fmt.Errorf("failed to create DNS responder: %v", err)
 	}
 
+	tworesponder, err := tworeqresp.NewResponder(respder)
+	if err != nil {
+		return nil, fmt.Errorf("error adding fragmentation layer: %v", err)
+	}
+
 	return &DNSRegServer{
-		dnsResponder: respder,
+		dnsResponder: tworesponder,
 		processor:    regprocessor,
 		latestCCGen:  latestClientConfGeneration,
 		logger:       logger,