Not finished stuff

internal/protocols/dot/dot.go

@@ -1,122 +1,161 @@
 package dot
 
 import (
-	"bytes"
+	"context"
 	"crypto/tls"
 	"encoding/binary"
 	"fmt"
 	"io"
+	"log"
 	"net"
 	"os"
+	"sync"
+	"time"
 
-	"github.com/afonsofrancof/sdns-perf/internal/protocols/do53"
 	"github.com/miekg/dns"
 )
 
-type DoTClient struct {
-	tcpConn    *net.TCPConn
-	tlsConn    *tls.Conn
-	keyLogFile *os.File
+type Config struct {
+	Host   string
+	Port   string
+	DNSSEC bool
+	Debug  bool
 }
 
-func New(target string) (*DoTClient, error) {
+type Client struct {
+	config Config
 
-	tcpAddr, err := net.ResolveTCPAddr("tcp", target)
+	serverAddr *net.TCPAddr
+
+	tcpConn    *net.TCPConn
+	tlsConn    *tls.Conn
+	tlsConfig  *tls.Config
+	keyLogFile *os.File
+
+	sendChannel chan *dns.Msg
+
+	responseChannels map[uint16]chan *dns.Msg
+	responseMutex    *sync.Mutex
+}
+
+func New(config Config) (*Client, error) {
+	serverAddr, err := net.ResolveTCPAddr("tcp", net.JoinHostPort(config.Host, config.Port))
 	if err != nil {
-		return nil, fmt.Errorf("failed to resolve TCP address: %v", err)
+		return nil, fmt.Errorf("dot: failed to resolve TCP address %q: %w", config.Host, err)
 	}
 
-	tcpConn, err := net.DialTCP("tcp", nil, tcpAddr)
-	if err != nil {
-		return nil, fmt.Errorf("failed to establish TCP connection: %v", err)
-	}
-
-	keyLogFile, err := os.OpenFile(
-		"tls-key-log.txt",
-		os.O_APPEND|os.O_CREATE|os.O_WRONLY,
-		0600,
-	)
-	if err != nil {
-		return nil, fmt.Errorf("failed opening key log file: %v", err)
+	var keyLogFile *os.File
+	if config.Debug {
+		keyLogFile, err = os.OpenFile(
+			"tls-key-log.txt",
+			os.O_APPEND|os.O_CREATE|os.O_WRONLY,
+			0600,
+		)
+		if err != nil {
+			log.Printf("dot: failed opening TLS key log file: %v", err)
+			keyLogFile = nil
+		}
 	}
 
 	tlsConfig := &tls.Config{
-		InsecureSkipVerify: true,
+		ServerName:         serverAddr.IP.String(),
 		MinVersion:         tls.VersionTLS12,
 		KeyLogWriter:       keyLogFile,
+		ClientSessionCache: tls.NewLRUClientSessionCache(100),
 	}
 
-	tlsConn := tls.Client(tcpConn, tlsConfig)
-	err = tlsConn.Handshake()
-	if err != nil {
-		return nil, fmt.Errorf("failed to execute the TLS handshake: %v", err)
+	client := &Client{
+		config:     config,
+		serverAddr: serverAddr,
+		tlsConfig:  tlsConfig,
+		keyLogFile: keyLogFile,
 	}
 
-	return &DoTClient{tcpConn: tcpConn, tlsConn: tlsConn, keyLogFile: keyLogFile}, nil
+	go client.receiveLoop()
+
+	return client, nil
 }
 
-func (c *DoTClient) Close() {
-	if c.tcpConn != nil {
-		c.tcpConn.Close()
-	}
+func (c *Client) Close() {
 	if c.tlsConn != nil {
 		c.tlsConn.Close()
+		c.tlsConn = nil
 	}
+
+	if c.tcpConn != nil {
+		c.tcpConn.Close()
+		c.tcpConn = nil
+	}
+
 	if c.keyLogFile != nil {
 		c.keyLogFile.Close()
+		c.keyLogFile = nil
 	}
 }
 
-func (c *DoTClient) Query(domain, queryType, target string, dnssec bool) error {
+func (c *Client) receiveLoop() {
 
-	DNSMessage, err := do53.NewDNSMessage(domain, queryType)
+	lengthBuffer := make([]byte, 2)
+	buffer := make([]byte, dns.MaxMsgSize)
+
+	for {
+		msgSize, err := io.ReadFull(c.tlsConn, lengthBuffer)
+		if err != nil {
+			log.Printf("doh: failed to read the DNS message's size: %s", err.Error())
+			// FIX: HANDLE RECONNECTION
+		}
+		n, err := io.ReadFull(c.tlsConn, buffer[:msgSize])
+		if err != nil {
+			log.Printf("doh: failed to read the DNS message: %s", err.Error())
+			// FIX: HANDLE RECONNECTION
+		}
+
+		recvMsg := new(dns.Msg)
+		err = recvMsg.Unpack(buffer[:n])
+		if err != nil {
+			log.Printf("do53: failed to unpack DNS response: %s", err.Error())
+			continue
+		}
+
+		c.responseMutex.Lock()
+		respChan, ok := c.responseChannels[recvMsg.Id]
+		delete(c.responseChannels, recvMsg.Id)
+		c.responseMutex.Unlock()
+
+		if ok {
+			respChan <- recvMsg
+		} else {
+			log.Printf("Receiver: Received DNS response for unknown or already processed msg ID: %v\n", recvMsg.Id)
+		}
+
+	}
+
+}
+
+func (c *Client) connect(ctx context.Context) error {
+	tcpConn, err := net.DialTCP("tcp", nil, c.serverAddr)
 	if err != nil {
-		return err
+		return fmt.Errorf("dot: failed to establish TCP connection: %w", err)
 	}
 
-	var lengthPrefixedMessage bytes.Buffer
-	err = binary.Write(&lengthPrefixedMessage, binary.BigEndian, uint16(len(DNSMessage)))
+	c.tcpConn.SetKeepAlive(true)
+	c.tcpConn.SetKeepAlivePeriod(1 * time.Minute)
+
+	tlsConn := tls.Client(c.tcpConn, c.tlsConfig)
+	err = tlsConn.HandshakeContext(ctx)
 	if err != nil {
-		return fmt.Errorf("failed to write message length: %v", err)
-	}
-	_, err = lengthPrefixedMessage.Write(DNSMessage)
-	if err != nil {
-		return fmt.Errorf("failed to write DNS message: %v", err)
+		c.tcpConn.Close()
+		c.tcpConn = nil
+		return fmt.Errorf("dot: failed to execute the TLS handshake: %w", err)
 	}
 
-	_, err = c.tlsConn.Write(lengthPrefixedMessage.Bytes())
-	if err != nil {
-		return fmt.Errorf("failed writing TLS request: %v", err)
-	}
+	c.tlsConn = tlsConn
 
-	lengthBuf := make([]byte, 2)
-	_, err = io.ReadFull(c.tlsConn, lengthBuf)
-	if err != nil {
-		return fmt.Errorf("failed reading response length: %v", err)
-	}
-
-	messageLength := binary.BigEndian.Uint16(lengthBuf)
-	if messageLength == 0 {
-		return fmt.Errorf("received zero-length message")
-	}
-
-	responseBuf := make([]byte, messageLength)
-	_, err = io.ReadFull(c.tlsConn, responseBuf)
-	if err != nil {
-		return fmt.Errorf("failed reading TLS response: %v", err)
-	}
-
-	recvMsg := new(dns.Msg)
-	err = recvMsg.Unpack(responseBuf)
-	if err != nil {
-		return fmt.Errorf("failed to parse DNS response: %v", err)
-	}
-
-	// TODO: Check if the response had no errors or TD bit set
-
-	for _, answer := range recvMsg.Answer {
-		fmt.Println(answer.String())
-	}
+	log.Println("dot: TCP/TLS connection established successfully.")
 
 	return nil
 }
+
+func (c *Client) Query(domain string, queryType uint16) (*dns.Msg, error) {
+	//TODO
+}