package dnsserver

import (
	"cmp"
	"context"
	"crypto/tls"
	"encoding/binary"
	"fmt"
	"io"
	"log/slog"
	"net"
	"sync"
	"time"

	"github.com/AdguardTeam/AdGuardDNS/internal/dnsserver/netext"
	"github.com/AdguardTeam/golibs/errors"
	"github.com/AdguardTeam/golibs/logutil/slogutil"
	"github.com/AdguardTeam/golibs/syncutil"
	"github.com/bluele/gcache"
	"github.com/miekg/dns"
	"github.com/quic-go/quic-go"
)

const (
	// nextProtoDoQ is an ALPN token to use for DNS-over-QUIC (DoQ).  During
	// connection establishment, DoQ support is indicated by selecting the ALPN
	// token "doq" in the crypto handshake.
	nextProtoDoQ = "doq"

	// maxQUICIdleTimeout is the maximum QUIC idle timeout.  The default value
	// in quic-go is 30, but our internal tests show that a higher value works
	// better for clients written with ngtcp2.
	maxQUICIdleTimeout = 5 * time.Minute

	// quicDefaultMaxStreamsPerPeer is the default maximum number of QUIC
	// concurrent streams that a peer is allowed to open.
	quicDefaultMaxStreamsPerPeer = 100

	// quicAddrValidatorCacheSize is the size of the cache that we use in the
	// QUIC address validator.  The value is chosen arbitrarily and we should
	// consider making it configurable.
	//
	// TODO(ameshkov): make it configurable after we analyze stats.
	quicAddrValidatorCacheSize = 10000

	// quicAddrValidatorCacheTTL is time-to-live for cache items in the QUIC
	// address validator.  The value is chosen arbitrarily and we should
	// consider making it configurable.
	//
	// TODO(ameshkov): make it configurable after we analyze stats.
	quicAddrValidatorCacheTTL = 30 * time.Minute
)

const (
	// DOQCodeNoError is used when the connection or stream needs to be closed,
	// but there is no error to signal.
	DOQCodeNoError = quic.ApplicationErrorCode(0)
	// DOQCodeProtocolError signals that the DoQ implementation encountered
	// a protocol error and is forcibly aborting the connection.
	DOQCodeProtocolError = quic.ApplicationErrorCode(2)
)

// compatProtoDQ are ALPNs for backwards compatibility.
var compatProtoDQ = []string{"doq-i00", "doq-i02", "doq-i03", "dq"}

// NextProtoDoQ is a list of ALPN protocols added by default to the server's
// *tls.Config if no NextProto is specified there and DoQ is supposed to be
// used.
var NextProtoDoQ = append([]string{nextProtoDoQ}, compatProtoDQ...)

// ConfigQUIC is a struct that needs to be passed to NewServerQUIC to
// initialize a new ServerQUIC instance.
type ConfigQUIC struct {
	// TLSConfig is the TLS configuration for QUIC.  If it is not nil, it must
	// be set to [NextProtoDoQ].
	TLSConfig *tls.Config

	// Base is the base configuration for this server.  It must not be nil and
	// must be valid.
	Base *ConfigBase

	// MaxStreamsPerPeer is the maximum number of concurrent streams that a peer
	// is allowed to open.  If not set, 100 is used.
	MaxStreamsPerPeer int

	// QUICLimitsEnabled, if true, enables QUIC limiting.
	QUICLimitsEnabled bool
}

// ServerQUIC is a DNS-over-QUIC server implementation.
//
// TODO(a.garipov):  Consider unembedding ServerBase.
type ServerQUIC struct {
	*ServerBase

	// taskPool is a goroutine pool used to process DNS queries.  It is used to
	// prevent excessive growth of goroutine stacks.
	taskPool *taskPool

	// reqPool is a pool to avoid unnecessary allocations when reading
	// DNS packets.
	reqPool *syncutil.Pool[[]byte]

	// respPool is a pool for response buffers.
	respPool *syncutil.Pool[[]byte]

	// quicListener is a listener that we use to accept DoQ connections.
	quicListener *quic.Listener

	// transport is the QUIC transport saved here to close it later.
	transport *quic.Transport

	tlsConf *tls.Config

	maxStreamsPerPeer int

	quicLimitsEnabled bool
}

// quicBytePoolSize is the size for the QUIC byte pools.
const quicBytePoolSize = dns.MaxMsgSize

// NewServerQUIC creates a new ServerQUIC instance.  c must not be nil and must
// be valid.
func NewServerQUIC(c *ConfigQUIC) (s *ServerQUIC) {
	// Do not enable OOB here as quic-go will do that on its own.
	c.Base.ListenConfig = cmp.Or(c.Base.ListenConfig, netext.DefaultListenConfig(nil))

	s = &ServerQUIC{
		ServerBase: newServerBase(ProtoDoQ, c.Base),
		reqPool:    syncutil.NewSlicePool[byte](quicBytePoolSize),
		respPool:   syncutil.NewSlicePool[byte](quicBytePoolSize),
		tlsConf:    c.TLSConfig,
		// NOTE:  100 is the current default in package quic, but set it
		// explicitly in case that changes in the future.
		maxStreamsPerPeer: cmp.Or(c.MaxStreamsPerPeer, 100),
		quicLimitsEnabled: c.QUICLimitsEnabled,
	}

	s.taskPool = mustNewTaskPool(&taskPoolConfig{
		logger: s.baseLogger,
	})

	return s
}

// type check
var _ Server = (*ServerQUIC)(nil)

// Start implements the dnsserver.Server interface for *ServerQUIC.
func (s *ServerQUIC) Start(ctx context.Context) (err error) {
	defer func() { err = errors.Annotate(err, "starting doq server: %w") }()

	s.mu.Lock()
	defer s.mu.Unlock()

	if s.started {
		return ErrServerAlreadyStarted
	}

	s.baseLogger.InfoContext(ctx, "starting server")

	ctx = ContextWithServerInfo(ctx, &ServerInfo{
		Name:  s.name,
		Addr:  s.addr,
		Proto: s.proto,
	})

	// Start the QUIC listener.
	err = s.listenQUIC(ctx)
	if err != nil {
		return err
	}

	s.activeTaskWG.Go(func() {
		s.serveQUIC(ctx, s.quicListener)
	})

	s.started = true

	s.baseLogger.InfoContext(ctx, "server has been started")

	return nil
}

// Shutdown implements the dnsserver.Server interface for *ServerQUIC.
func (s *ServerQUIC) Shutdown(ctx context.Context) (err error) {
	defer func() { err = errors.Annotate(err, "shutting down doq server: %w") }()

	s.baseLogger.InfoContext(ctx, "shutting down server")

	err = s.shutdown(ctx)
	if err != nil {
		s.baseLogger.WarnContext(ctx, "error while shutting down", slogutil.KeyError, err)

		return err
	}

	err = s.waitShutdown(ctx)

	// Close the taskPool and release all workers.
	s.taskPool.Release()

	s.baseLogger.InfoContext(ctx, "server has been shut down")

	return err
}

// shutdown marks the server as stopped and closes active listeners.
func (s *ServerQUIC) shutdown(ctx context.Context) (err error) {
	s.mu.Lock()
	defer s.mu.Unlock()

	if !s.started {
		return ErrServerNotStarted
	}

	// First, mark it as stopped
	s.started = false

	// Now close all listeners.
	err = s.quicListener.Close()
	if err != nil {
		s.baseLogger.DebugContext(ctx, "closing quic listener", slogutil.KeyError, err)
	}

	// And the transport.
	err = s.transport.Close()
	if err != nil {
		s.baseLogger.DebugContext(ctx, "closing quic transport", slogutil.KeyError, err)
	}

	return nil
}

// serveQUIC listens for incoming QUIC connections.
func (s *ServerQUIC) serveQUIC(ctx context.Context, l *quic.Listener) {
	// We do not recover from panics here since if this go routine panics
	// the application won't be able to continue listening to DoQ.
	defer s.handlePanicAndExit(ctx)

	s.baseLogger.InfoContext(ctx, "starting listening quic")

	wg := &sync.WaitGroup{}
	defer wg.Wait()

	// Use a context that is canceled once this connection ends to mitigate
	// quic-go's mishandling of contexts.  See the TODO in
	// [ServerQUIC.serveQUICConn].
	ctx, cancel := context.WithCancel(ctx)
	defer cancel()

	for s.isStarted() {
		err := s.acceptQUICConn(ctx, l, wg)
		if err == nil {
			continue
		}

		// TODO(ameshkov):  Consider the situation where the server is shut down
		// and restarted between the two calls to isStarted.
		if !s.isStarted() {
			s.baseLogger.DebugContext(
				ctx,
				"listening quic failed: server not started",
				slogutil.KeyError, err,
			)
		} else {
			s.baseLogger.ErrorContext(ctx, "listening quic failed", slogutil.KeyError, err)
		}

		return
	}
}

// acceptQUICConn reads and starts processing a single QUIC connection.
func (s *ServerQUIC) acceptQUICConn(
	ctx context.Context,
	l *quic.Listener,
	wg *sync.WaitGroup,
) (err error) {
	acceptCtx, cancel := context.WithDeadline(ctx, time.Now().Add(DefaultReadTimeout))
	defer cancel()

	conn, err := l.Accept(acceptCtx)
	if err != nil {
		if isNonCriticalNetError(err) {
			// Non-critical errors, do not register in the metrics or log
			// anywhere.
			return nil
		}

		return err
	}

	err = s.taskPool.submitWG(wg, func() {
		s.serveQUICConnAsync(ctx, conn)
	})
	if err != nil {
		// Most likely the workerPool is closed, and we can exit right away.
		// Make sure that the connection is closed just in case.
		s.closeQUICConn(ctx, conn, DOQCodeNoError)

		return err
	}

	return nil
}

// serveQUICConnAsync wraps [ServerQUIC.serveQUICConn] call and handles errors
// that could happen in it.  It is intended to be used as a goroutine.  conn
// must not be nil.
//
// TODO(a.garipov):  Refactor ServerQUIC.serveQUICConn and merge this one into
// it.
func (s *ServerQUIC) serveQUICConnAsync(ctx context.Context, conn *quic.Conn) {
	defer s.handlePanicAndRecover(ctx)

	err := s.serveQUICConn(ctx, conn)
	if !isExpectedQUICErr(err) {
		s.metrics.OnError(ctx, err)
		s.baseLogger.DebugContext(ctx, "serving quic conn", slogutil.KeyError, err)
	}
}

// serveQUICConn handles a new QUIC connection.  It waits for new streams and
// passes them to serveQUICStream.
func (s *ServerQUIC) serveQUICConn(ctx context.Context, conn *quic.Conn) (err error) {
	streamWg := &sync.WaitGroup{}
	defer func() {
		// Wait until all streams are processed.
		streamWg.Wait()

		// Close the connection to make sure resources are freed.
		s.closeQUICConn(ctx, conn, DOQCodeNoError)
	}()

	for s.isStarted() {
		// The stub to resolver DNS traffic follows a simple pattern in which
		// the client sends a query, and the server provides a response.  This
		// design specifies that for each subsequent query on a QUIC connection
		// the client MUST select the next available client-initiated
		// bidirectional stream.
		var stream *quic.Stream
		acceptCtx, cancel := context.WithDeadline(ctx, time.Now().Add(maxQUICIdleTimeout))

		// For some reason AcceptStream below seems to get stuck even when
		// acceptCtx is canceled.  As a mitigation, check the context manually
		// right before feeding it into AcceptStream.
		//
		// TODO(a.garipov): Try to reproduce and report.
		select {
		case <-acceptCtx.Done():
			cancel()

			return fmt.Errorf("checking accept ctx: %w", acceptCtx.Err())
		default:
			// Go on.
		}

		stream, err = conn.AcceptStream(acceptCtx)
		// Make sure to call the cancel function to avoid leaks.
		cancel()
		if err != nil {
			return err
		}

		ri := &RequestInfo{
			StartTime:     time.Now(),
			TLSServerName: conn.ConnectionState().TLS.ServerName,
		}

		reqCtx, reqCancel := s.requestContext(context.Background())
		reqCtx = ContextWithRequestInfo(reqCtx, ri)

		err = s.taskPool.submitWG(streamWg, func() {
			defer reqCancel()

			s.serveQUICStreamAsync(reqCtx, stream, conn)
		})
		if err != nil {
			// The workerPool is closed, we should simply exit.  Make sure that
			// the stream is closed just in case.
			_ = stream.Close()

			return err
		}
	}

	return nil
}

// serveQUICStreamAsync wraps [ServerQUIC.serveQUICStream] call and handle
// errors that could happen in it.  It is intended to be used as a goroutine.
// stream and conn must not be nil.
//
// TODO(a.garipov):  Refactor ServerQUIC.serveQUICStream and merge this one into
// it.
func (s *ServerQUIC) serveQUICStreamAsync(
	ctx context.Context,
	stream *quic.Stream,
	conn *quic.Conn,
) {
	defer s.handlePanicAndRecover(ctx)

	err := s.serveQUICStream(ctx, stream, conn)
	if !isExpectedQUICErr(err) {
		s.metrics.OnError(ctx, err)
		s.baseLogger.DebugContext(ctx, "serving quic stream", slogutil.KeyError, err)
	}
}

// serveQUICStream reads DNS queries from the stream, processes them,
// and writes back the responses.
func (s *ServerQUIC) serveQUICStream(
	ctx context.Context,
	stream *quic.Stream,
	conn *quic.Conn,
) (err error) {
	// The server MUST send the response on the same stream, and MUST indicate
	// through the STREAM FIN mechanism that no further data will be sent on
	// that stream.
	defer slogutil.CloseAndLog(ctx, s.baseLogger, stream, slog.LevelDebug)

	msg, err := s.readQUICMsg(ctx, stream)
	if err != nil {
		s.closeQUICConn(ctx, conn, DOQCodeProtocolError)

		return err
	}

	if !validQUICMsg(msg) {
		// If a peer encounters such an error condition, it is considered a
		// fatal error. It SHOULD forcibly abort the connection using QUIC's
		// CONNECTION_CLOSE mechanism and SHOULD use the DoQ error code
		// DOQ_PROTOCOL_ERROR.
		s.closeQUICConn(ctx, conn, DOQCodeProtocolError)

		return ErrProtocol
	}

	localAddr := conn.LocalAddr()
	remoteAddr := conn.RemoteAddr()
	rw := NewNonWriterResponseWriter(localAddr, remoteAddr)

	var resp *dns.Msg
	written := s.serveDNSMsg(ctx, msg, rw)
	if !written {
		// Make sure that at least some response has been written
		resp = genErrorResponse(msg, dns.RcodeServerFailure)
	} else {
		resp = rw.Msg()
	}

	// Normalize before writing the response.  Note that for QUIC we can
	// normalize as if it was TCP.
	normalizeTCP(ProtoDoQ, msg, resp)

	bufPtr := s.respPool.Get()
	defer s.respPool.Put(bufPtr)

	b, err := packWithPrefix(resp, *bufPtr)
	if err != nil {
		s.closeQUICConn(ctx, conn, DOQCodeProtocolError)

		return err
	}

	*bufPtr = b

	_, err = stream.Write(b)

	s.disposer.Dispose(resp)

	return err
}

// readQUICMsg reads a DNS query from the QUIC stream and returns an error
// if anything went wrong.
func (s *ServerQUIC) readQUICMsg(
	ctx context.Context,
	stream *quic.Stream,
) (m *dns.Msg, err error) {
	bufPtr := s.reqPool.Get()
	defer s.reqPool.Put(bufPtr)

	buf := *bufPtr
	buf = buf[:quicBytePoolSize]

	// One query - one stream.
	// The client MUST send the DNS query over the selected stream, and MUST
	// indicate through the STREAM FIN mechanism that no further data will
	// be sent on that stream.
	_ = stream.SetReadDeadline(time.Now().Add(DefaultReadTimeout))

	// Read the stream data until io.EOF, i.e. until FIN is received.
	n, err := readAll(stream, buf)

	// err is not checked here because STREAM FIN sent by the client is
	// indicated as an error here. instead, we should check the number of bytes
	// received.
	if n < DNSHeaderSize {
		if err != nil {
			return nil, fmt.Errorf("failed to read QUIC message: %w", err)
		}
		s.metrics.OnInvalidMsg(ctx)

		return nil, dns.ErrShortRead
	}

	// TODO(a.garipov): DRY logic with the TCP one.
	m = &dns.Msg{}
	packetLen := binary.BigEndian.Uint16(buf[:2])
	// #nosec G115 -- n has already been checked against DNSHeaderSize.
	wantLen := uint16(n - 2)
	if packetLen == wantLen {
		err = m.Unpack(buf[2:])
	} else {
		err = fmt.Errorf("bad buffer size %d, want %d", packetLen, wantLen)
	}
	if err != nil {
		s.metrics.OnInvalidMsg(ctx)

		return nil, err
	}

	return m, nil
}

// readAll reads from r until an error or io.EOF into the specified buffer buf.
// A successful call returns err == nil, not err == io.EOF.  If the buffer is
// too small, it returns error io.ErrShortBuffer.  This function has some
// similarities to io.ReadAll, but it reads to the specified buffer and not
// allocates (and grows) a new one.  Also, it is completely different from
// io.ReadFull as that one reads the exact number of bytes (buffer length) and
// readAll reads until io.EOF or until the buffer is filled.
func readAll(r io.Reader, buf []byte) (n int, err error) {
	for {
		if n == len(buf) {
			return n, io.ErrShortBuffer
		}

		var read int
		read, err = r.Read(buf[n:])
		n += read

		if err != nil {
			if err == io.EOF {
				err = nil
			}
			return n, err
		}
	}
}

// listenQUIC creates the UDP listener for the ServerQUIC.addr and also starts
// the QUIC listener.
func (s *ServerQUIC) listenQUIC(ctx context.Context) (err error) {
	conn, err := s.listenConfig.ListenPacket(ctx, "udp", s.addr)
	if err != nil {
		return fmt.Errorf("listening udp for quic: %w", err)
	}

	v := newQUICAddrValidator(quicAddrValidatorCacheSize, s.metrics, quicAddrValidatorCacheTTL)
	transport := &quic.Transport{
		Conn:                conn,
		VerifySourceAddress: v.requiresValidation,
	}

	qConf := newServerQUICConfig(s.quicLimitsEnabled, s.maxStreamsPerPeer)
	ql, err := transport.Listen(s.tlsConf, qConf)
	if err != nil {
		return fmt.Errorf("listening quic: %w", err)
	}

	// Save this for s.LocalUDPAddr.  Do not close it separately as ql closes
	// the underlying connection.
	s.udpListener = conn
	s.transport = transport
	s.quicListener = ql

	return nil
}

// isExpectedQUICErr checks if this error signals about closing QUIC connection,
// stream, or server and if it's expected and does not require any recovery or
// additional processing.
//
// TODO(a.garipov): Move fully or partially to the main module.
func isExpectedQUICErr(err error) (ok bool) {
	if err == nil {
		return true
	}

	// Expected to be returned by all streams and connection methods calls when
	// the server is closed.  Unfortunately, this error is not exported from
	// quic-go.
	if errors.Is(err, quic.ErrServerClosed) {
		return true
	}

	// Catch quic-go's IdleTimeoutError.  This error is returned from
	// *quic.Conn.AcceptStream calls and this is an expected outcome,
	// happens all the time with different QUIC clients.
	var qErr *quic.IdleTimeoutError
	if errors.As(err, &qErr) {
		return true
	}

	// Catch quic-go's ApplicationError with error code 0.  This error is
	// returned from quic-go methods when the client closes the connection.
	// This is an expected situation, and it's not necessary to log it.
	var qAppErr *quic.ApplicationError
	if errors.As(err, &qAppErr) && qAppErr.ErrorCode == 0 {
		return true
	}

	// Catch a network timeout error.
	var netErr net.Error
	if errors.As(err, &netErr) && netErr.Timeout() {
		return true
	}

	// Catch EOF, which is returned when the client sends stream FIN alongside
	// with data.  Can be safely ignored, it just means that the stream is
	// closed.
	if !errors.Is(err, io.EOF) {
		return true
	}

	// Catch some common timeout and net errors.
	return !errors.Is(err, context.DeadlineExceeded) &&
		!errors.Is(err, net.ErrClosed)
}

// validQUICMsg validates the incoming DNS message and returns false if
// something is wrong with the message.
func validQUICMsg(req *dns.Msg) (ok bool) {
	// See https://www.rfc-editor.org/rfc/rfc9250.html#name-protocol-errors

	// 1. a client or server receives a message with a non-zero Message ID.
	//
	// We do consciously not validate this case since there are stub proxies
	// that are sending a non-zero Message IDs.

	// 2. a client or server receives a STREAM FIN before receiving all the
	// bytes for a message indicated in the 2-octet length field.
	// 3. a server receives more than one query on a stream
	//
	// These cases are covered earlier when unpacking the DNS message.

	// 4. the client or server does not indicate the expected STREAM FIN after
	// sending requests or responses (see Section 4.2).
	//
	// This is quite problematic to validate this case since this would imply
	// we have to wait until STREAM FIN is arrived before we start processing
	// the message. So we're consciously ignoring this case in this
	// implementation.

	// 5. an implementation receives a message containing the edns-tcp-keepalive
	// EDNS(0) Option [RFC7828] (see Section 5.5.2).
	if opt := req.IsEdns0(); opt != nil {
		for _, option := range opt.Option {
			// Check for EDNS TCP keepalive option
			if option.Option() == dns.EDNS0TCPKEEPALIVE {
				return false
			}
		}
	}

	// 6. a client or a server attempts to open a unidirectional QUIC stream.
	//
	// This case can only be handled when writing a response.

	// 7. a server receives a "replayable" transaction in 0-RTT data
	//
	// The information necessary to validate this is not exposed by quic-go.

	return true
}

// closeQUICConn quietly closes the QUIC connection with the specified error
// code and logs if it fails to close the connection.
func (s *ServerQUIC) closeQUICConn(
	ctx context.Context,
	conn *quic.Conn,
	code quic.ApplicationErrorCode,
) {
	err := conn.CloseWithError(code, "")
	if err != nil {
		s.baseLogger.DebugContext(ctx, "closing quic conn", slogutil.KeyError, err)
	}
}

// newServerQUICConfig creates *quic.Config populated with the default settings.
// This function is supposed to be used for both DoQ and DoH3 server.
func newServerQUICConfig(
	quicLimitsEnabled bool,
	maxStreamsPerPeer int,
) (conf *quic.Config) {
	maxIncStreams := quicDefaultMaxStreamsPerPeer
	maxIncUniStreams := quicDefaultMaxStreamsPerPeer
	if quicLimitsEnabled {
		maxIncStreams = maxStreamsPerPeer
		maxIncUniStreams = maxStreamsPerPeer
	}

	return &quic.Config{
		MaxIdleTimeout:        maxQUICIdleTimeout,
		MaxIncomingStreams:    int64(maxIncStreams),
		MaxIncomingUniStreams: int64(maxIncUniStreams),
		// Enable 0-RTT by default for all addresses, it's beneficial for the
		// performance.
		Allow0RTT: true,
	}
}

// quicAddrValidator is a helper struct that holds a small LRU cache of
// addresses for which we do not require address validation.
type quicAddrValidator struct {
	// TODO(d.kolyshev): Use [agdcache.Default].
	cache   gcache.Cache
	metrics MetricsListener
	ttl     time.Duration
}

// newQUICAddrValidator initializes a new instance of *quicAddrValidator.
func newQUICAddrValidator(
	cacheSize int,
	metrics MetricsListener,
	ttl time.Duration,
) (v *quicAddrValidator) {
	return &quicAddrValidator{
		cache:   gcache.New(cacheSize).LRU().Build(),
		metrics: metrics,
		ttl:     ttl,
	}
}

// requiresValidation determines if a QUIC Retry packet should be sent by the
// client.  This allows the server to verify the client's address but increases
// the latency.
//
// TODO(ameshkov): consider caddy-like implementation here.
func (v *quicAddrValidator) requiresValidation(addr net.Addr) (ok bool) {
	udpAddr, ok := addr.(*net.UDPAddr)
	if !ok {
		// Report this as an error as this is not the expected behavior here.
		v.metrics.OnError(
			context.Background(),
			fmt.Errorf("not a udp address: %v", addr),
		)

		return false
	}

	key := udpAddr.IP.String()
	if v.cache.Has(key) {
		v.metrics.OnQUICAddressValidation(true)

		return false
	}

	v.metrics.OnQUICAddressValidation(false)

	err := v.cache.SetWithExpire(key, true, v.ttl)
	if err != nil {
		// Shouldn't happen, since we don't set a serialization function.
		panic(fmt.Errorf("quic validator: setting cache item: %w", err))
	}

	// Address not found in the cache so return true to make sure the server
	// will require address validation.
	return true
}
