<commit_before>#include "dnsmsg.h"
#include "wrappers.h"
#include "dnstypes.h"

#include <stdio.h>
#include <string.h>

#define MAX_NAME_LEN 256

#define DNS_OPCODE_MASK 0x7800
#define DNS_RCODE_MASK 0xf
#define DNS_Z_MASK 0x70

void append_to_buf8(uint8_t **buf, uint8_t val);
void append_to_buf16(uint8_t **buf, uint16_t val);
uint16_t dnsmsg_header_get_opt(struct dnsmsg_header *);
uint16_t get_uint16(uint8_t **buf);
uint8_t get_uint8(uint8_t **buf);
struct dnsmsg_header *dnsmsg_parse_header(uint8_t **buf);
struct dnsmsg_query *dnsmsg_parse_query(uint8_t **buf, uint8_t *original);
char *dnsmsg_parse_name(uint8_t **buf, uint8_t *original);
dnsrecord_t *dnsmsg_parse_record(uint8_t **buf, uint8_t *original);

struct dnsmsg_header *new_dnsmsg_header(uint16_t options) {
    struct dnsmsg_header *result = _malloc(sizeof(struct dnsmsg_header));
    result->id = rand();
    result->qr = (options & DNSMSG_OPT_QR) > 0;
    result->tc = (options & DNSMSG_OPT_TC) > 0;
    result->rd = (options & DNSMSG_OPT_RD) > 0;
    result->ra = (options & DNSMSG_OPT_RA) > 0;
    return result;
}

struct dnsmsg_query *new_dnsmsg_query(void) {
    struct dnsmsg_query *result = _malloc(sizeof(struct dnsmsg_query));
    return result;
}

dnsmsg_t *dnsmsg_new(uint16_t options) {
    dnsmsg_t *result = _malloc(sizeof(dnsmsg_t));
    result->header = new_dnsmsg_header(options);
    result->query = new_dnsmsg_query();
    result->answer = NULL;
    result->auth = NULL;
    result->additional = NULL;
    return result;
}

void dnsmsg_print(dnsmsg_t *msg) {
    printf("-----------------------------------------------\n");
    printf(";; Header\n\n");
    printf("  ID:             %-#4x\n", msg->header->id);
    printf("  Response Code:  %-#4x\n", msg->header->rcode);
    printf("  Flags:          %-#4x\n", dnsmsg_header_get_opt(msg->header));
    printf("  Questions:      %-4d\n", msg->header->qdcount);
    printf("  Answers:        %-4d\n", msg->header->ancount);
    printf("  Authority:      %-4d\n", msg->header->nscount);
    printf("  Additional:     %-4d\n", msg->header->arcount);
    printf("\n");
    printf(";; Question Section\n");
    printf("  QNAME: %s\n", msg->query->qname);
    printf("  QTYPE: %x\n", msg->query->qtype);
    printf("  QCLASS: %x\n", msg->query->qclass);
    printf("\n");
    printf(";; Answer Section\n");
    int i = 0;
    for (; i < msg->header->ancount; i++) {
        dnsrecord_print(msg->answer[i]);
    }
    printf("-----------------------------------------------\n");
}

void dnsmsg_free(dnsmsg_t *msg) {
    if (msg->header) free(msg->header); 
    if (msg->query) {
        if (msg->query->qname) free(msg->query->qname);
        free(msg->query);
    }
    free(msg);
}

dnsmsg_t *dnsmsg_parse(uint8_t **buf, ssize_t buf_len) {
    uint8_t *original = *buf;
    
    dnsmsg_t *result = malloc(sizeof(dnsmsg_t));
    int i = 0;
    
    result->header = dnsmsg_parse_header(buf);
    result->query = dnsmsg_parse_query(buf, original);
    
    result->answer = malloc(sizeof(dnsrecord_t *) * result->header->ancount);
    for (i = 0; i < result->header->ancount; i++) {
        result->answer[i] = dnsmsg_parse_record(buf, original);
    }
    
    result->auth = malloc(sizeof(dnsrecord_t) * result->header->nscount);
    for (i = 0; i < result->header->nscount; i++) {
        result->auth[i] = dnsmsg_parse_record(buf, original);
    }
    
    result->additional = malloc(sizeof(dnsrecord_t) * result->header->arcount);
    for (i = 0; i < result->header->arcount; i++) {
        result->additional[i] = dnsmsg_parse_record(buf, original);
    }
    
    return result;
}

struct dnsmsg_header *dnsmsg_parse_header(uint8_t **buf) {
    struct dnsmsg_header *result = malloc(sizeof(struct dnsmsg_header));
    result->id = get_uint16(buf);
    uint16_t row_2 = get_uint16(buf);
    result->qr = (row_2 & DNSMSG_OPT_QR) >> 15;
    result->opcode = (row_2 & DNS_OPCODE_MASK) >> 11;
    result->aa = (row_2 & DNSMSG_OPT_AA) >> 10;
    result->tc = (row_2 & DNSMSG_OPT_TC) >> 9;
    result->rd = (row_2 & DNSMSG_OPT_RD) >> 8;
    result->ra = (row_2 & DNSMSG_OPT_RA) >> 7;
    result->z = (row_2 & DNS_Z_MASK) >> 4;
    result->rcode = (row_2 & DNS_RCODE_MASK);
    result->qdcount = get_uint16(buf);
    result->ancount = get_uint16(buf);
    result->nscount = get_uint16(buf);
    result->arcount = get_uint16(buf);
    return result;
}

struct dnsmsg_query *dnsmsg_parse_query(uint8_t **buf, uint8_t *original_buf) {
    struct dnsmsg_query *result = malloc(sizeof(struct dnsmsg_query));
    result->qname = dnsmsg_parse_name(buf, original_buf);
    result->qtype = get_uint16(buf);
    result->qclass = get_uint16(buf);
    return result;
}

dnsrecord_t *dnsmsg_parse_record(uint8_t **buf, uint8_t *original_buf) {
    int i;
    dnsrecord_t *result = dnsrecord_new();
    result->name = dnsmsg_parse_name(buf, original_buf);
    result->type = get_uint16(buf);
    result->rclass = get_uint16(buf);
    result->ttl = (get_uint16(buf) << 16) | get_uint16(buf);
    result->rdlength = get_uint16(buf);
    result->rdata = malloc(sizeof(uint8_t *) * result->rdlength);
    for (i = 0; i < result->rdlength; i++) {
        result->rdata[i] = get_uint8(buf);
    }
    return result;
}

char *dnsmsg_parse_name(uint8_t **buf, uint8_t *original_buf) {
    int i = 0;
    uint8_t label_len = get_uint8(buf);
    char *name = calloc(sizeof(char *), MAX_NAME_LEN);
    char *name_ptr = name;
    do {
        if ((label_len & 0xc0) == 0xc0) {
             // label_len indicates a pointer
            uint16_t offset = ((label_len & 0x3f) << 8) | get_uint8(buf);
            // Rewind buf
            uint8_t *new_buf_pos = original_buf + offset;
            int new_label_len = *(new_buf_pos++);
            for (i = 0; i < new_label_len; i++) {
                *(name_ptr++) = (char) *(new_buf_pos++);
            }
            *(name_ptr++) = '.';
        } else {
            for (i = 0; i < label_len; i++) {
                *(name_ptr++) = (char) get_uint8(buf);
            }
            *(name_ptr++) = '.';
        }
        label_len = get_uint8(buf);
    } while (label_len != 0);
    return name;
}

dnsmsg_t *dnsmsg_create_query(const char *name, int qtype, int qclass) {
    dnsmsg_t *result = dnsmsg_new(DNSMSG_OPT_RD);
    result->header->opcode = DNS_OPCODE_QUERY;
    result->header->qdcount = 1;
    
    result->query->qname = calloc(sizeof(char), strlen(name) + 1);
    strcpy(result->query->qname, name);
    result->query->qtype = qtype;
    result->query->qclass = qclass;
    return result;
}

uint8_t *dnsmsg_create_name(const char *name) {<commit_msg>Fixed name parsing not following pointer to the end of the name<commit_after>#include "dnsmsg.h"
#include "wrappers.h"
#include "dnstypes.h"

#include <stdio.h>
#include <string.h>

#define MAX_NAME_LEN 256

#define DNS_OPCODE_MASK 0x7800
#define DNS_RCODE_MASK 0xf
#define DNS_Z_MASK 0x70

void append_to_buf8(uint8_t **buf, uint8_t val);
void append_to_buf16(uint8_t **buf, uint16_t val);
uint16_t dnsmsg_header_get_opt(struct dnsmsg_header *);
uint16_t get_uint16(uint8_t **buf);
uint8_t get_uint8(uint8_t **buf);
struct dnsmsg_header *dnsmsg_parse_header(uint8_t **buf);
struct dnsmsg_query *dnsmsg_parse_query(uint8_t **buf, uint8_t *original);
char *dnsmsg_parse_name(uint8_t **buf, uint8_t *original);
dnsrecord_t *dnsmsg_parse_record(uint8_t **buf, uint8_t *original);

struct dnsmsg_header *new_dnsmsg_header(uint16_t options) {
    struct dnsmsg_header *result = _malloc(sizeof(struct dnsmsg_header));
    result->id = rand();
    result->qr = (options & DNSMSG_OPT_QR) > 0;
    result->tc = (options & DNSMSG_OPT_TC) > 0;
    result->rd = (options & DNSMSG_OPT_RD) > 0;
    result->ra = (options & DNSMSG_OPT_RA) > 0;
    return result;
}

struct dnsmsg_query *new_dnsmsg_query(void) {
    struct dnsmsg_query *result = _malloc(sizeof(struct dnsmsg_query));
    return result;
}

dnsmsg_t *dnsmsg_new(uint16_t options) {
    dnsmsg_t *result = _malloc(sizeof(dnsmsg_t));
    result->header = new_dnsmsg_header(options);
    result->query = new_dnsmsg_query();
    result->answer = NULL;
    result->auth = NULL;
    result->additional = NULL;
    return result;
}

void dnsmsg_print(dnsmsg_t *msg) {
    printf("-----------------------------------------------\n");
    printf(";; Header\n");
    printf("  ID:             %-#4x\n", msg->header->id);
    printf("  Response Code:  %-#4x\n", msg->header->rcode);
    printf("  Flags:          %-#4x\n", dnsmsg_header_get_opt(msg->header));
    printf("  Questions:      %-4d\n", msg->header->qdcount);
    printf("  Answers:        %-4d\n", msg->header->ancount);
    printf("  Authority:      %-4d\n", msg->header->nscount);
    printf("  Additional:     %-4d\n", msg->header->arcount);
    printf("\n");
    printf(";; Question Section\n");
    printf("  QNAME: %s\n", msg->query->qname);
    printf("  QTYPE: %x\n", msg->query->qtype);
    printf("  QCLASS: %x\n", msg->query->qclass);
    printf("\n");
    printf(";; Answer Section\n");
    int i = 0;
    for (; i < msg->header->ancount; i++) {
        dnsrecord_print(msg->answer[i]);
    }
    printf("-----------------------------------------------\n");
}

void dnsmsg_free(dnsmsg_t *msg) {
    if (msg->header) free(msg->header); 
    if (msg->query) {
        if (msg->query->qname) free(msg->query->qname);
        free(msg->query);
    }
    free(msg);
}

dnsmsg_t *dnsmsg_parse(uint8_t **buf, ssize_t buf_len) {
    uint8_t *original = *buf;
    
    dnsmsg_t *result = malloc(sizeof(dnsmsg_t));
    int i = 0;
    
    result->header = dnsmsg_parse_header(buf);
    result->query = dnsmsg_parse_query(buf, original);
    
    result->answer = malloc(sizeof(dnsrecord_t *) * result->header->ancount);
    for (i = 0; i < result->header->ancount; i++) {
        result->answer[i] = dnsmsg_parse_record(buf, original);
    }
    
    result->auth = malloc(sizeof(dnsrecord_t) * result->header->nscount);
    for (i = 0; i < result->header->nscount; i++) {
        result->auth[i] = dnsmsg_parse_record(buf, original);
    }
    
    result->additional = malloc(sizeof(dnsrecord_t) * result->header->arcount);
    for (i = 0; i < result->header->arcount; i++) {
        result->additional[i] = dnsmsg_parse_record(buf, original);
    }
    
    return result;
}

struct dnsmsg_header *dnsmsg_parse_header(uint8_t **buf) {
    struct dnsmsg_header *result = malloc(sizeof(struct dnsmsg_header));
    result->id = get_uint16(buf);
    uint16_t row_2 = get_uint16(buf);
    result->qr = (row_2 & DNSMSG_OPT_QR) >> 15;
    result->opcode = (row_2 & DNS_OPCODE_MASK) >> 11;
    result->aa = (row_2 & DNSMSG_OPT_AA) >> 10;
    result->tc = (row_2 & DNSMSG_OPT_TC) >> 9;
    result->rd = (row_2 & DNSMSG_OPT_RD) >> 8;
    result->ra = (row_2 & DNSMSG_OPT_RA) >> 7;
    result->z = (row_2 & DNS_Z_MASK) >> 4;
    result->rcode = (row_2 & DNS_RCODE_MASK);
    result->qdcount = get_uint16(buf);
    result->ancount = get_uint16(buf);
    result->nscount = get_uint16(buf);
    result->arcount = get_uint16(buf);
    return result;
}

struct dnsmsg_query *dnsmsg_parse_query(uint8_t **buf, uint8_t *original_buf) {
    struct dnsmsg_query *result = malloc(sizeof(struct dnsmsg_query));
    result->qname = dnsmsg_parse_name(buf, original_buf);
    result->qtype = get_uint16(buf);
    result->qclass = get_uint16(buf);
    return result;
}

dnsrecord_t *dnsmsg_parse_record(uint8_t **buf, uint8_t *original_buf) {
    int i;
    dnsrecord_t *result = dnsrecord_new();
    result->name = dnsmsg_parse_name(buf, original_buf);
    result->type = get_uint16(buf);
    result->rclass = get_uint16(buf);
    result->ttl = (get_uint16(buf) << 16) | get_uint16(buf);
    result->rdlength = get_uint16(buf);
    result->rdata = malloc(sizeof(uint8_t *) * result->rdlength);
    for (i = 0; i < result->rdlength; i++) {
        result->rdata[i] = get_uint8(buf);
    }
    return result;
}

// char *dnsmsg_parse_name(uint8_t **buf, uint8_t *original_buf) {
//     int i;
    
//     // Read name from buffer
//     uint8_t label_len = get_uint8(buf);
//     char *name = calloc(sizeof(char *), MAX_NAME_LEN);
//     char *name_ptr = name; // This variable is used like a string builder
    
//     do {
        
//     } 
    
//     // If encountered pointer, recurse
//     if ((label_len & 0xc0) == 0xc0) {
//     }
// }

char *dnsmsg_parse_name(uint8_t **buf, uint8_t *original_buf) {
    int i = 0;
    uint8_t label_len = get_uint8(buf);
    char *name = calloc(sizeof(char *), MAX_NAME_LEN);
    char *name_ptr = name;
    do {
        if ((label_len & 0xc0) == 0xc0) {
            // Get new buf position
            int offset = ((label_len & 0x3f << 8) | get_uint8(buf));
            uint8_t *new_buf_pos = original_buf + offset;
            char *pointed = dnsmsg_parse_name(&new_buf_pos, original_buf);
            for (i = 0; pointed[i] != '\0'; i++) {
                *(name_ptr++) = pointed[i];
            }
            *(name_ptr++) = '\0';
            free(pointed); // pointed is no longer needed. Goodbye
        } else {
            for (i = 0; i < label_len; i++) {
                *(name_ptr++) = (char) get_uint8(buf);
            }
            *(name_ptr++) = '.';
        }
        label_len = get_uint8(buf);
    } while (label_len != 0);
    return name;
}

dnsmsg_t *dnsmsg_create_query(const char *name, int qtype, int qclass) {
    dnsmsg_t *result = dnsmsg_new(DNSMSG_OPT_RD);
    result->header->opcode = DNS_OPCODE_QUERY;
    result->header->qdcount = 1;
    
    result->query->qname = calloc(sizeof(char), strlen(name) + 1);
    strcpy(result->query->qname, name);
    result->query->qtype = qtype;
    result->query->qclass = qclass;
    return result;
}

uint8_t *dnsmsg_create_name(const char *name) {
{'fraction_non_alphanumeric': 0.1690284136327302, 'fraction_numerical': 0.020928711576193592, 'mean_word_length': 3.002908667830134, 'pattern_counts': {'":': 0, '<': 28, '<?xml version=': 0, '>': 237, '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': '5059465', 'n_tokens_mistral': 5324, 'n_tokens_neox': 5116, 'n_words': 983}