<commit_before>package dns

import (
	"bytes"
	"encoding/binary"
	"errors"
	"io"
)

const (
	MaxNameLen  = 255
	MaxLabelLen = 63

	headerLen    = 12
	octetPairLen = 2
)

var (
	ErrLabelTooLong = errors.New("a label must be 63 octets or less")
	ErrNameTooLong  = errors.New("a name must be 255 octets or less")
)

type Unpacker struct {
	buffer []byte
	msg    *Msg
	i      int

	nameBuffer bytes.Buffer
}

func NewUnpacker() *Unpacker {
	return &Unpacker{}
}

func (r *Unpacker) Unpack() (msg *Msg, n int, err error) {
	err = r.unpackHeader()
	if err != nil {
		return nil, 0, err
	}

	err = r.unpackQueries()
	if err != nil {
		return nil, 0, err
	}

	return r.msg, r.i, nil
}

func (r *Unpacker) Reset(b []byte) {
	r.buffer = b
	r.msg = &Msg{Header: Header{}}
	r.i = 0
}

func (r *Unpacker) readQType() QType {
	return QType(r.unpackOctetPair())
}

func (r *Unpacker) readQClass() QClass {
	return QClass(r.unpackOctetPair())
}

func (r *Unpacker) unpackHeader() (err error) {
	if len(r.buffer) < headerLen {
		return io.ErrShortBuffer
	}
	h := Header{}
	h.ID = r.unpackOctetPair()

	flags := r.unpackOctetPair()
	h.QR = (flags & maskQR >> maskQROffset) == 1
	h.OpCode = OpCode(flags & maskOpCode >> maskOpCodeOffset)
	h.AA = (flags & maskAA >> maskAAOffset) == 1
	h.TC = (flags & maskTC >> maskTCOffset) == 1
	h.RD = (flags & maskRD >> maskRDOffset) == 1
	h.RA = (flags & maskRA >> maskRAOffset) == 1
	h.Z = byte(flags & maskZ >> maskZOffset)
	h.RCode = byte(flags & maskRCode)

	h.QDCount = r.unpackOctetPair()
	h.ANCount = r.unpackOctetPair()
	h.NSCount = r.unpackOctetPair()
	h.ARCount = r.unpackOctetPair()

	r.msg.Header = h
	return nil
}

func (r *Unpacker) unpackQueries() (err error) {
	queries := make([]Query, r.msg.Header.QDCount)
	r.msg.Queries = queries
	for i := range queries {
		qName, err := r.readName()
		if err != nil {
			return err
		}
		if r.i+octetPairLen >= len(r.buffer) {
			return io.ErrShortBuffer
		}
		queries[i].QName = qName
		queries[i].QType = r.readQType()
		queries[i].QClass = r.readQClass()
	}
	return nil
}

func (r *Unpacker) readName() (string, error) {
	r.nameBuffer.Reset()
	for {
		label, err := unpackLabel(r.buffer, r.i)
		if err == io.EOF {
			r.i++
			return r.nameBuffer.String(), nil
		} else if err != nil {
			return "", err
		}

		r.i += len(label) + 1
		r.nameBuffer.WriteString(label + ".")
		if r.nameBuffer.Len() > MaxNameLen {
			return "", ErrNameTooLong
		}
	}
}

func (r *Unpacker) unpackOctetPair() uint16 {
	ini := r.i
	r.i += octetPairLen
	return binary.BigEndian.Uint16(r.buffer[ini:r.i])
}

func (r *Unpacker) unpackByte() (byte, error) {
	if r.i >= len(r.buffer) {
		return 0, io.ErrShortBuffer
	}
	b := r.buffer[r.i]
	r.i++
	return b, nil
}

func unpackLabel(b []byte, offset int) (string, error) {
	if !checkBounds(b, offset, offset+1) {
		return "", io.ErrShortBuffer
	}
	currentByte := b[offset]
	if currentByte == 0 {
		return "", io.EOF
	}
	offset++
	switch {
	case currentByte>>6 == 3:
		// Compute offsets for the pointer.
		offset = int((currentByte&64)<<8 + b[offset])
		currentByte = b[offset]
		offset++
	}
	end := offset + int(currentByte)
	if end-offset > MaxLabelLen {
		return "", ErrLabelTooLong
	}
	if !checkBounds(b, offset, end) {
		return "", io.ErrShortBuffer
	}
	return string(b[offset:end]), nil
}

func checkBounds(b []byte, begin, end int) bool {
	return len(b[begin:]) >= end-begin
}<commit_msg>(wip) add comments on the unpackLabel<commit_after>package dns

import (
	"bytes"
	"encoding/binary"
	"errors"
	"io"
)

const (
	MaxNameLen  = 255
	MaxLabelLen = 63

	headerLen    = 12
	octetPairLen = 2
)

var (
	ErrLabelTooLong = errors.New("a label must be 63 octets or less")
	ErrNameTooLong  = errors.New("a name must be 255 octets or less")
)

type Unpacker struct {
	buffer []byte
	msg    *Msg
	i      int

	nameBuffer bytes.Buffer
}

func NewUnpacker() *Unpacker {
	return &Unpacker{}
}

func (r *Unpacker) Unpack() (msg *Msg, n int, err error) {
	err = r.unpackHeader()
	if err != nil {
		return nil, 0, err
	}

	err = r.unpackQueries()
	if err != nil {
		return nil, 0, err
	}

	return r.msg, r.i, nil
}

func (r *Unpacker) Reset(b []byte) {
	r.buffer = b
	r.msg = &Msg{Header: Header{}}
	r.i = 0
}

func (r *Unpacker) readQType() QType {
	return QType(r.unpackOctetPair())
}

func (r *Unpacker) readQClass() QClass {
	return QClass(r.unpackOctetPair())
}

func (r *Unpacker) unpackHeader() (err error) {
	if len(r.buffer) < headerLen {
		return io.ErrShortBuffer
	}
	h := Header{}
	h.ID = r.unpackOctetPair()

	flags := r.unpackOctetPair()
	h.QR = (flags & maskQR >> maskQROffset) == 1
	h.OpCode = OpCode(flags & maskOpCode >> maskOpCodeOffset)
	h.AA = (flags & maskAA >> maskAAOffset) == 1
	h.TC = (flags & maskTC >> maskTCOffset) == 1
	h.RD = (flags & maskRD >> maskRDOffset) == 1
	h.RA = (flags & maskRA >> maskRAOffset) == 1
	h.Z = byte(flags & maskZ >> maskZOffset)
	h.RCode = byte(flags & maskRCode)

	h.QDCount = r.unpackOctetPair()
	h.ANCount = r.unpackOctetPair()
	h.NSCount = r.unpackOctetPair()
	h.ARCount = r.unpackOctetPair()

	r.msg.Header = h
	return nil
}

func (r *Unpacker) unpackQueries() (err error) {
	queries := make([]Query, r.msg.Header.QDCount)
	r.msg.Queries = queries
	for i := range queries {
		qName, err := r.readName()
		if err != nil {
			return err
		}
		if r.i+octetPairLen >= len(r.buffer) {
			return io.ErrShortBuffer
		}
		queries[i].QName = qName
		queries[i].QType = r.readQType()
		queries[i].QClass = r.readQClass()
	}
	return nil
}

func (r *Unpacker) readName() (string, error) {
	r.nameBuffer.Reset()
	for {
		label, err := unpackLabel(r.buffer, r.i)
		if err == io.EOF {
			r.i++
			return r.nameBuffer.String(), nil
		} else if err != nil {
			return "", err
		}

		r.i += len(label) + 1
		r.nameBuffer.WriteString(label + ".")
		if r.nameBuffer.Len() > MaxNameLen {
			return "", ErrNameTooLong
		}
	}
}

func (r *Unpacker) unpackOctetPair() uint16 {
	ini := r.i
	r.i += octetPairLen
	return binary.BigEndian.Uint16(r.buffer[ini:r.i])
}

func (r *Unpacker) unpackByte() (byte, error) {
	if r.i >= len(r.buffer) {
		return 0, io.ErrShortBuffer
	}
	b := r.buffer[r.i]
	r.i++
	return b, nil
}

func unpackLabel(b []byte, offset int) (string, error) {
	if !checkBounds(b, offset, offset+1) {
		return "", io.ErrShortBuffer
	}

	// Current byte indicates the length of the label.
	// If its a null byte, label is over.
	currentByte := b[offset]
	if currentByte == 0 {
		return "", io.EOF
	}

	offset++

	// Check if its a pointer.
	if currentByte>>6 == 3 {

		// Compute the offset to the pointer.
		offset = int((currentByte&64)<<8 + b[offset])
		currentByte = b[offset]
		offset++
	}

	// Check if the label has valid length.
	endOffset := offset + int(currentByte)
	if endOffset-offset > MaxLabelLen {
		return "", ErrLabelTooLong
	}
	if !checkBounds(b, offset, endOffset) {
		return "", io.ErrShortBuffer
	}

	// Return the label.
	return string(b[offset:endOffset]), nil
}

// Check if begin and end are within bounds of a byte slice.
func checkBounds(b []byte, begin, end int) bool {
	return len(b) >= begin && len(b[begin:]) >= end-begin
}
{'fraction_non_alphanumeric': 0.15512621633055987, 'fraction_numerical': 0.009589620645889155, 'mean_word_length': 3.413192283758556, 'pattern_counts': {'":': 0, '<': 9, '<?xml version=': 0, '>': 46, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5764704', 'n_tokens_mistral': 3006, 'n_tokens_neox': 2824, 'n_words': 722}