github.com/xmplusdev/xmcore@v1.8.11-0.20240412132628-5518b55526af/transport/internet/headers/dns/dns.go (about)

     1  package dns
     2  
     3  import (
     4  	"context"
     5  	"encoding/binary"
     6  	"errors"
     7  
     8  	"github.com/xmplusdev/xmcore/common"
     9  	"github.com/xmplusdev/xmcore/common/dice"
    10  )
    11  
    12  type DNS struct {
    13  	header []byte
    14  }
    15  
    16  func (d DNS) Size() int32 {
    17  	return int32(len(d.header))
    18  }
    19  
    20  // Serialize implements PacketHeader.
    21  func (d DNS) Serialize(b []byte) {
    22  	copy(b, d.header)
    23  	binary.BigEndian.PutUint16(b[0:], dice.RollUint16()) // random transaction ID
    24  }
    25  
    26  // NewDNS returns a new DNS instance based on given config.
    27  func NewDNS(ctx context.Context, config interface{}) (interface{}, error) {
    28  	var header []byte
    29  
    30  	header = binary.BigEndian.AppendUint16(header, 0x0000) // Transaction ID
    31  	header = binary.BigEndian.AppendUint16(header, 0x0100) // Flags: Standard query
    32  	header = binary.BigEndian.AppendUint16(header, 0x0001) // Questions
    33  	header = binary.BigEndian.AppendUint16(header, 0x0000) // Answer RRs
    34  	header = binary.BigEndian.AppendUint16(header, 0x0000) // Authority RRs
    35  	header = binary.BigEndian.AppendUint16(header, 0x0000) // Additional RRs
    36  
    37  	buf := make([]byte, 0x100)
    38  
    39  	off1, err := packDomainName(config.(*Config).Domain+".", buf)
    40  	if err != nil {
    41  		return nil, err
    42  	}
    43  
    44  	header = append(header, buf[:off1]...)
    45  
    46  	header = binary.BigEndian.AppendUint16(header, 0x0001) // Type: A
    47  	header = binary.BigEndian.AppendUint16(header, 0x0001) // Class: IN
    48  
    49  	return DNS{
    50  		header: header,
    51  	}, nil
    52  }
    53  
    54  // copied from github.com/miekg/dns
    55  func packDomainName(s string, msg []byte) (off1 int, err error) {
    56  	off := 0
    57  	ls := len(s)
    58  	// Each dot ends a segment of the name.
    59  	// We trade each dot byte for a length byte.
    60  	// Except for escaped dots (\.), which are normal dots.
    61  	// There is also a trailing zero.
    62  
    63  	// Emit sequence of counted strings, chopping at dots.
    64  	var (
    65  		begin int
    66  		bs    []byte
    67  	)
    68  	for i := 0; i < ls; i++ {
    69  		var c byte
    70  		if bs == nil {
    71  			c = s[i]
    72  		} else {
    73  			c = bs[i]
    74  		}
    75  
    76  		switch c {
    77  		case '\\':
    78  			if off+1 > len(msg) {
    79  				return len(msg), errors.New("buffer size too small")
    80  			}
    81  
    82  			if bs == nil {
    83  				bs = []byte(s)
    84  			}
    85  
    86  			copy(bs[i:ls-1], bs[i+1:])
    87  			ls--
    88  		case '.':
    89  			labelLen := i - begin
    90  			if labelLen >= 1<<6 { // top two bits of length must be clear
    91  				return len(msg), errors.New("bad rdata")
    92  			}
    93  
    94  			// off can already (we're in a loop) be bigger than len(msg)
    95  			// this happens when a name isn't fully qualified
    96  			if off+1+labelLen > len(msg) {
    97  				return len(msg), errors.New("buffer size too small")
    98  			}
    99  
   100  			// The following is covered by the length check above.
   101  			msg[off] = byte(labelLen)
   102  
   103  			if bs == nil {
   104  				copy(msg[off+1:], s[begin:i])
   105  			} else {
   106  				copy(msg[off+1:], bs[begin:i])
   107  			}
   108  			off += 1 + labelLen
   109  			begin = i + 1
   110  		default:
   111  		}
   112  	}
   113  
   114  	if off < len(msg) {
   115  		msg[off] = 0
   116  	}
   117  
   118  	return off + 1, nil
   119  }
   120  
   121  func init() {
   122  	common.Must(common.RegisterConfig((*Config)(nil), NewDNS))
   123  }