github.com/metacubex/mihomo@v1.18.5/component/resolver/relay.go (about)

     1  package resolver
     2  
     3  import (
     4  	"context"
     5  	"encoding/binary"
     6  	"io"
     7  	"net"
     8  	"time"
     9  
    10  	"github.com/metacubex/mihomo/common/pool"
    11  
    12  	D "github.com/miekg/dns"
    13  )
    14  
    15  const DefaultDnsReadTimeout = time.Second * 10
    16  const DefaultDnsRelayTimeout = time.Second * 5
    17  
    18  const SafeDnsPacketSize = 2 * 1024 // safe size which is 1232 from https://dnsflagday.net/2020/, so 2048 is enough
    19  
    20  func RelayDnsConn(ctx context.Context, conn net.Conn, readTimeout time.Duration) error {
    21  	buff := pool.Get(pool.UDPBufferSize)
    22  	defer func() {
    23  		_ = pool.Put(buff)
    24  		_ = conn.Close()
    25  	}()
    26  	for {
    27  		if readTimeout > 0 {
    28  			_ = conn.SetReadDeadline(time.Now().Add(readTimeout))
    29  		}
    30  
    31  		length := uint16(0)
    32  		if err := binary.Read(conn, binary.BigEndian, &length); err != nil {
    33  			break
    34  		}
    35  
    36  		if int(length) > len(buff) {
    37  			break
    38  		}
    39  
    40  		n, err := io.ReadFull(conn, buff[:length])
    41  		if err != nil {
    42  			break
    43  		}
    44  
    45  		err = func() error {
    46  			ctx, cancel := context.WithTimeout(ctx, DefaultDnsRelayTimeout)
    47  			defer cancel()
    48  			inData := buff[:n]
    49  			msg, err := relayDnsPacket(ctx, inData, buff, 0)
    50  			if err != nil {
    51  				return err
    52  			}
    53  
    54  			err = binary.Write(conn, binary.BigEndian, uint16(len(msg)))
    55  			if err != nil {
    56  				return err
    57  			}
    58  
    59  			_, err = conn.Write(msg)
    60  			if err != nil {
    61  				return err
    62  			}
    63  			return nil
    64  		}()
    65  		if err != nil {
    66  			return err
    67  		}
    68  	}
    69  	return nil
    70  }
    71  
    72  func relayDnsPacket(ctx context.Context, payload []byte, target []byte, maxSize int) ([]byte, error) {
    73  	msg := &D.Msg{}
    74  	if err := msg.Unpack(payload); err != nil {
    75  		return nil, err
    76  	}
    77  
    78  	r, err := ServeMsg(ctx, msg)
    79  	if err != nil {
    80  		m := new(D.Msg)
    81  		m.SetRcode(msg, D.RcodeServerFailure)
    82  		return m.PackBuffer(target)
    83  	}
    84  
    85  	r.SetRcode(msg, r.Rcode)
    86  	if maxSize > 0 {
    87  		r.Truncate(maxSize)
    88  	}
    89  	r.Compress = true
    90  	return r.PackBuffer(target)
    91  }
    92  
    93  // RelayDnsPacket will truncate udp message up to SafeDnsPacketSize
    94  func RelayDnsPacket(ctx context.Context, payload []byte, target []byte) ([]byte, error) {
    95  	return relayDnsPacket(ctx, payload, target, SafeDnsPacketSize)
    96  }