Files
godns/pkg/doh/client.go
2026-01-06 02:25:24 +08:00

174 lines
3.5 KiB
Go

package doh
import (
"context"
"encoding/base64"
"io"
"net"
"net/http"
"net/http/httptrace"
"strings"
"time"
"github.com/miekg/dns"
"github.com/pkg/errors"
"golang.org/x/net/proxy"
)
const (
dohMediaType = "application/dns-message"
)
// Logger 定义可选的日志接口
type Logger interface {
Printf(format string, v ...interface{})
}
type clientOptions struct {
timeout time.Duration
server string
bootstrap func(domain string) (net.IP, error)
getDialer func(d *net.Dialer) (proxy.Dialer, proxy.ContextDialer, error)
logger Logger
}
type ClientOption func(*clientOptions) error
func WithTimeout(t time.Duration) ClientOption {
return func(o *clientOptions) error {
o.timeout = t
return nil
}
}
func WithSocksProxy(getDialer func(d *net.Dialer) (proxy.Dialer, proxy.ContextDialer, error)) ClientOption {
return func(o *clientOptions) error {
o.getDialer = getDialer
return nil
}
}
func WithServer(server string) ClientOption {
return func(o *clientOptions) error {
o.server = server
return nil
}
}
func WithBootstrap(resolver func(domain string) (net.IP, error)) ClientOption {
return func(o *clientOptions) error {
o.bootstrap = resolver
return nil
}
}
func WithLogger(logger Logger) ClientOption {
return func(o *clientOptions) error {
o.logger = logger
return nil
}
}
type Client struct {
opt *clientOptions
cli *http.Client
traceCtx context.Context
}
func NewClient(opts ...ClientOption) *Client {
o := new(clientOptions)
for _, f := range opts {
f(o)
}
clientTrace := &httptrace.ClientTrace{
GotConn: func(info httptrace.GotConnInfo) {
if o.logger != nil {
o.logger.Printf("http conn was reused: %t", info.Reused)
}
},
}
var transport *http.Transport
if o.bootstrap != nil {
transport = &http.Transport{
DialContext: func(ctx context.Context, network, address string) (net.Conn, error) {
urls := strings.Split(address, ":")
ipv4, err := o.bootstrap(urls[0])
if err != nil {
return nil, errors.Wrap(err, "bootstrap")
}
urls[0] = ipv4.String()
if o.getDialer != nil {
dialer, _, err := o.getDialer(&net.Dialer{
Timeout: o.timeout,
})
if err != nil {
return nil, err
}
return dialer.Dial("tcp", strings.Join(urls, ":"))
}
return (&net.Dialer{
Timeout: o.timeout,
}).DialContext(ctx, network, strings.Join(urls, ":"))
},
}
}
return &Client{
opt: o,
traceCtx: httptrace.WithClientTrace(context.Background(), clientTrace),
cli: &http.Client{
Transport: transport,
Timeout: o.timeout,
},
}
}
func (c *Client) Exchange(req *dns.Msg) (r *dns.Msg, rtt time.Duration, err error) {
var (
buf []byte
begin = time.Now()
origID = req.Id
hreq *http.Request
)
// Set DNS ID as zero accoreding to RFC8484 (cache friendly)
req.Id = 0
buf, err = req.Pack()
if err != nil {
return
}
hreq, err = http.NewRequestWithContext(c.traceCtx, http.MethodGet, c.opt.server+"?dns="+base64.RawURLEncoding.EncodeToString(buf), nil)
if err != nil {
return
}
hreq.Header.Add("Accept", dohMediaType)
hreq.Header.Add("User-Agent", "godns-doh-client/0.1")
resp, err := c.cli.Do(hreq)
if err != nil {
return
}
defer resp.Body.Close()
content, err := io.ReadAll(resp.Body)
if err != nil {
return
}
if resp.StatusCode != http.StatusOK {
err = errors.New("DoH query failed: " + string(content))
return
}
r = new(dns.Msg)
err = r.Unpack(content)
r.Id = origID
rtt = time.Since(begin)
return
}