shater/dns/transport/https.go

package transport

import (
	"bytes"
	"context"
	"errors"
	"io"
	"net"
	"net/http"
	"net/url"
	"strconv"
	"sync"
	"sync/atomic"
	"time"

	"github.com/sagernet/sing-box/adapter"
	"github.com/sagernet/sing-box/common/dialer"
	"github.com/sagernet/sing-box/common/tls"
	C "github.com/sagernet/sing-box/constant"
	"github.com/sagernet/sing-box/dns"
	"github.com/sagernet/sing-box/log"
	"github.com/sagernet/sing-box/option"
	"github.com/sagernet/sing/common"
	"github.com/sagernet/sing/common/buf"
	E "github.com/sagernet/sing/common/exceptions"
	"github.com/sagernet/sing/common/logger"
	M "github.com/sagernet/sing/common/metadata"
	N "github.com/sagernet/sing/common/network"
	sHTTP "github.com/sagernet/sing/protocol/http"

	mDNS "github.com/miekg/dns"
	"golang.org/x/net/http2"
)

const MimeType = "application/dns-message"

var _ adapter.DNSTransport = (*HTTPSTransport)(nil)

func RegisterHTTPS(registry *dns.TransportRegistry) {
	dns.RegisterTransport[option.RemoteHTTPSDNSServerOptions](registry, C.DNSTypeHTTPS, NewHTTPS)
}

type HTTPSTransport struct {
	dns.TransportAdapter
	logger           logger.ContextLogger
	dialer           N.Dialer
	destination      *url.URL
	headers          http.Header
	transportAccess  sync.Mutex
	transport        *HTTPSTransportWrapper
	transportResetAt time.Time
}

func NewHTTPS(ctx context.Context, logger log.ContextLogger, tag string, options option.RemoteHTTPSDNSServerOptions) (adapter.DNSTransport, error) {
	transportDialer, err := dns.NewRemoteDialer(ctx, options.RemoteDNSServerOptions)
	if err != nil {
		return nil, err
	}
	tlsOptions := common.PtrValueOrDefault(options.TLS)
	tlsOptions.Enabled = true
	tlsConfig, err := tls.NewClient(ctx, logger, options.Server, tlsOptions)
	if err != nil {
		return nil, err
	}
	if len(tlsConfig.NextProtos()) == 0 {
		tlsConfig.SetNextProtos([]string{http2.NextProtoTLS, "http/1.1"})
	}
	headers := options.Headers.Build()
	host := headers.Get("Host")
	if host != "" {
		headers.Del("Host")
	} else {
		if tlsConfig.ServerName() != "" {
			host = tlsConfig.ServerName()
		} else {
			host = options.Server
		}
	}
	destinationURL := url.URL{
		Scheme: "https",
		Host:   host,
	}
	if destinationURL.Host == "" {
		destinationURL.Host = options.Server
	}
	if options.ServerPort != 0 && options.ServerPort != 443 {
		destinationURL.Host = net.JoinHostPort(destinationURL.Host, strconv.Itoa(int(options.ServerPort)))
	}
	path := options.Path
	if path == "" {
		path = "/dns-query"
	}
	err = sHTTP.URLSetPath(&destinationURL, path)
	if err != nil {
		return nil, err
	}
	serverAddr := options.DNSServerAddressOptions.Build()
	if serverAddr.Port == 0 {
		serverAddr.Port = 443
	}
	if !serverAddr.IsValid() {
		return nil, E.New("invalid server address: ", serverAddr)
	}
	return NewHTTPSRaw(
		dns.NewTransportAdapterWithRemoteOptions(C.DNSTypeHTTPS, tag, options.RemoteDNSServerOptions),
		logger,
		transportDialer,
		&destinationURL,
		headers,
		serverAddr,
		tlsConfig,
	), nil
}

func NewHTTPSRaw(
	adapter dns.TransportAdapter,
	logger log.ContextLogger,
	dialer N.Dialer,
	destination *url.URL,
	headers http.Header,
	serverAddr M.Socksaddr,
	tlsConfig tls.Config,
) *HTTPSTransport {
	if tlsConfig != nil {
		dialer = tls.NewDialer(dialer, tlsConfig)
	}
	return &HTTPSTransport{
		TransportAdapter: adapter,
		logger:           logger,
		dialer:           dialer,
		destination:      destination,
		headers:          headers,
		transport:        NewHTTPSTransportWrapper(dialer, serverAddr, destination),
	}
}

func (t *HTTPSTransport) Start(stage adapter.StartStage) error {
	if stage != adapter.StartStateStart {
		return nil
	}
	return dialer.InitializeDetour(t.dialer)
}

func (t *HTTPSTransport) Close() error {
	t.Reset()
	return nil
}

func (t *HTTPSTransport) Reset() {
	t.transportAccess.Lock()
	defer t.transportAccess.Unlock()
	t.transport.CloseIdleConnections()
	t.transport = t.transport.Clone()
}

func (t *HTTPSTransport) Exchange(ctx context.Context, message *mDNS.Msg) (*mDNS.Msg, error) {
	startAt := time.Now()
	response, err := t.exchange(ctx, message)
	if err != nil {
		if errors.Is(err, context.DeadlineExceeded) {
			t.transportAccess.Lock()
			defer t.transportAccess.Unlock()
			if t.transportResetAt.After(startAt) {
				return nil, err
			}
			t.transport.CloseIdleConnections()
			t.transport = t.transport.Clone()
			t.transportResetAt = time.Now()
		}
		return nil, err
	}
	return response, nil
}

// requestBuffer owns the pooled buffer that backs one DoH query.
//
// Both transports behind HTTPSTransportWrapper write the request body on a
// goroutine of their own and return from RoundTrip as soon as the response
// HEADERS arrive: net/http's write loop is still copying out of the body a
// bufferful at a time (4 KiB of write buffer, or io.Copy's 32 KiB once it hands
// the body to the connection), and http2's writeRequestBody has read only the
// first max-frame-size bytes of it. Returning the buffer to the pool at that
// point handed live memory to the next caller while the query was still going
// out — everything past that first copy left the router as whatever that caller
// had written there. A data race, and a memory-disclosure primitive aimed at the
// resolver. Measured, not reasoned: with the write parked mid-query the bytes on
// the wire diverge from the bytes we packed at exactly one copy buffer in.
//
// Ownership is counted rather than handed over once, because a retry holds two
// bodies at a time and the two transports order that differently:
// http.Transport.rewindBody CLOSES the old body before asking GetBody for a
// new one, while http2's shouldRetryRequest asks GetBody first and closes the
// old body on a goroutine. exchange keeps a count of its own until RoundTrip
// returns — the only window in which either can call GetBody — so neither
// ordering can free the buffer under the other. If a transport ever fails to
// close a body, the count never reaches zero and the buffer is simply not
// reused: garbage, not corruption.
type requestBuffer struct {
	buffer *buf.Buffer
	raw    []byte
	refs   atomic.Int32
}

func newRequestBuffer(buffer *buf.Buffer, raw []byte) *requestBuffer {
	holder := &requestBuffer{buffer: buffer, raw: raw}
	holder.refs.Store(1)
	return holder
}

// body hands out a reader over the packed query as one more owner. It refuses
// once the buffer is back in the pool, so a late caller gets an error instead
// of a reader over memory that now belongs to somebody else.
func (b *requestBuffer) body() (*pooledRequestBody, bool) {
	for {
		refs := b.refs.Load()
		if refs < 1 {
			return nil, false
		}
		if b.refs.CompareAndSwap(refs, refs+1) {
			return &pooledRequestBody{Reader: bytes.NewReader(b.raw), owner: b}, true
		}
	}
}

func (b *requestBuffer) release() {
	if b.refs.Add(-1) == 0 {
		b.buffer.Release()
	}
}

type pooledRequestBody struct {
	*bytes.Reader
	owner    *requestBuffer
	closeOne sync.Once
}

func (b *pooledRequestBody) Close() error {
	b.closeOne.Do(b.owner.release)
	return nil
}

func (t *HTTPSTransport) exchange(ctx context.Context, message *mDNS.Msg) (*mDNS.Msg, error) {
	exMessage := *message
	exMessage.Id = 0
	exMessage.Compress = true
	requestBuffer := buf.NewSize(1 + message.Len())
	rawMessage, err := exMessage.PackBuffer(requestBuffer.FreeBytes())
	if err != nil {
		requestBuffer.Release()
		return nil, err
	}
	queryBuffer := newRequestBuffer(requestBuffer, rawMessage)
	// Drops the count exchange holds once RoundTrip is done with the request;
	// the bodies handed to the transport keep their own until it closes them.
	defer queryBuffer.release()
	requestBody, _ := queryBuffer.body() // cannot fail: the count above is ours
	request, err := http.NewRequestWithContext(ctx, http.MethodPost, t.destination.String(), requestBody)
	if err != nil {
		requestBody.Close()
		return nil, err
	}
	// http.NewRequestWithContext infers both only for the body types it knows,
	// and pooledRequestBody is not one of them. Upstream got them for free from
	// *bytes.Reader; GetBody is what lets a POST be replayed when a pooled
	// connection turns out to have been closed under us. Being unknown to
	// net/http also costs one packet on the HTTP/1.1 leg: isKnownInMemoryReader
	// no longer recognises the body, so the request headers are flushed before
	// the query instead of travelling with it.
	request.ContentLength = int64(len(rawMessage))
	request.GetBody = func() (io.ReadCloser, error) {
		retryBody, ok := queryBuffer.body()
		if !ok {
			return nil, E.New("DoH request buffer already released")
		}
		return retryBody, nil
	}
	request.Header = t.headers.Clone()
	request.Header.Set("Content-Type", MimeType)
	request.Header.Set("Accept", MimeType)
	t.transportAccess.Lock()
	currentTransport := t.transport
	t.transportAccess.Unlock()
	response, err := currentTransport.RoundTrip(request)
	if err != nil {
		return nil, err
	}
	defer response.Body.Close()
	if response.StatusCode != http.StatusOK {
		return nil, E.New("unexpected status: ", response.Status)
	}
	var responseMessage mDNS.Msg
	if response.ContentLength > 0 {
		responseBuffer := buf.NewSize(int(response.ContentLength))
		defer responseBuffer.Release()
		_, err = responseBuffer.ReadFullFrom(response.Body, int(response.ContentLength))
		if err != nil {
			return nil, err
		}
		err = responseMessage.Unpack(responseBuffer.Bytes())
	} else {
		rawMessage, err = io.ReadAll(response.Body)
		if err != nil {
			return nil, err
		}
		err = responseMessage.Unpack(rawMessage)
	}
	if err != nil {
		return nil, err
	}
	return &responseMessage, nil
}