refactor: add diale options to internal resolver

internal/resolver/doh.go

@@ -2,80 +2,93 @@ package resolver
 
 import (
 	"bytes"
+	"context"
+	"fmt"
 	"io"
 	"net"
 	"net/http"
 	"net/url"
-	"sync"
 	"time"
 )
 
 var _ net.Conn = &dohConn{}
 
 type dohConn struct {
-	addr string
+	query, resp *bytes.Buffer
 
-	once    sync.Once
-	onceErr error
-
-	in, ret bytes.Buffer
+	do func() error
 }
 
-func (c *dohConn) Close() error                       { return nil }
-func (c *dohConn) LocalAddr() net.Addr                { return nil }
-func (c *dohConn) RemoteAddr() net.Addr               { return nil }
-func (c *dohConn) SetDeadline(t time.Time) error      { return nil }
-func (c *dohConn) SetReadDeadline(t time.Time) error  { return nil }
-func (c *dohConn) SetWriteDeadline(t time.Time) error { return nil }
+func newDoHConn(ctx context.Context, client *http.Client, addr string) (*dohConn, error) {
+	c := &dohConn{
+		query: &bytes.Buffer{},
+		resp:  &bytes.Buffer{},
+	}
 
-func (c *dohConn) Write(b []byte) (int, error) { return c.in.Write(b) }
+	url, err := url.Parse(addr)
+	if err != nil {
+		return nil, err
+	}
+	// RFC 8484
+	url.Path = "/dns-query"
 
-func (c *dohConn) Read(b []byte) (int, error) {
-	c.once.Do(func() {
-		url, err := url.Parse(c.addr)
-		if err != nil {
-			c.onceErr = err
-			return
+	c.do = func() error {
+		if c.query.Len() <= 2 || c.resp.Len() > 0 {
+			return nil
 		}
-		// RFC 8484
-		url.Path = "/dns-query"
 
-		// Skip 2 bytes which are length
-		reqBody := bytes.NewReader(c.in.Bytes()[2:])
-		req, err := http.NewRequest("POST", url.String(), reqBody)
+		// Skip length header
+		c.query.Next(2)
+
+		req, err := http.NewRequestWithContext(ctx, http.MethodPost, url.String(), c.query)
 		if err != nil {
-			c.onceErr = err
-			return
+			return err
 		}
 		req.Header.Set("content-type", "application/dns-message")
 		req.Header.Set("accept", "application/dns-message")
-		resp, err := http.DefaultClient.Do(req)
+
+		resp, err := client.Do(req)
 		if err != nil {
-			c.onceErr = err
-			return
+			return err
 		}
 		defer resp.Body.Close()
+
 		respBody, err := io.ReadAll(resp.Body)
 		if err != nil {
-			c.onceErr = err
-			return
+			return err
 		}
 
+		if resp.StatusCode != http.StatusOK {
+			return fmt.Errorf("server return %d: %s", resp.StatusCode, respBody)
+		}
+
+		// Add length header
 		l := uint16(len(respBody))
-		_, err = c.ret.Write([]byte{uint8(l >> 8), uint8(l & ((1 << 8) - 1))})
+		_, err = c.resp.Write([]byte{uint8(l >> 8), uint8(l & ((1 << 8) - 1))})
 		if err != nil {
-			c.onceErr = err
-			return
+			return err
 		}
 
-		_, err = c.ret.Write(respBody)
-		if err != nil {
-			c.onceErr = err
-			return
-		}
-	})
-	if c.onceErr != nil {
-		return 0, c.onceErr
+		_, err = c.resp.Write(respBody)
+		return err
 	}
-	return c.ret.Read(b)
+
+	return c, nil
+}
+
+func (c *dohConn) Close() error                     { return nil }
+func (c *dohConn) LocalAddr() net.Addr              { return nil }
+func (c *dohConn) RemoteAddr() net.Addr             { return nil }
+func (c *dohConn) SetDeadline(time.Time) error      { return nil }
+func (c *dohConn) SetReadDeadline(time.Time) error  { return nil }
+func (c *dohConn) SetWriteDeadline(time.Time) error { return nil }
+
+func (c *dohConn) Write(b []byte) (int, error) { return c.query.Write(b) }
+
+func (c *dohConn) Read(b []byte) (int, error) {
+	if err := c.do(); err != nil {
+		return 0, err
+	}
+
+	return c.resp.Read(b)
 }