[thirdparty/pdns.git] / pdns / dnsdist-ecs.cc

/*
 * This file is part of PowerDNS or dnsdist.
 * Copyright -- PowerDNS.COM B.V. and its contributors
 *
 * This program is free software; you can redistribute it and/or modify
 * it under the terms of version 2 of the GNU General Public License as
 * published by the Free Software Foundation.
 *
 * In addition, for the avoidance of any doubt, permission is granted to
 * link this program with OpenSSL and to (re)distribute the binaries
 * produced as the result of such linking.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program; if not, write to the Free Software
 * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.
 */
#include "dolog.hh"
#include "dnsdist.hh"
#include "dnsdist-ecs.hh"
#include "dnsparser.hh"
#include "dnswriter.hh"
#include "ednsoptions.hh"
#include "ednssubnet.hh"

/* when we add EDNS to a query, we don't want to advertise
   a large buffer size */
size_t g_EdnsUDPPayloadSize = 512;
/* draft-ietf-dnsop-edns-client-subnet-04 "11.1.  Privacy" */
uint16_t g_ECSSourcePrefixV4 = 24;
uint16_t g_ECSSourcePrefixV6 = 56;

bool g_ECSOverride{false};

int rewriteResponseWithoutEDNS(const char * packet, const size_t len, vector<uint8_t>& newContent)
{
  assert(packet != NULL);
  assert(len >= sizeof(dnsheader));
  const struct dnsheader* dh = (const struct dnsheader*) packet;

  if (ntohs(dh->arcount) == 0)
    return ENOENT;

  if (ntohs(dh->qdcount) == 0)
    return ENOENT;
    
  vector<uint8_t> content(len - sizeof(dnsheader));
  copy(packet + sizeof(dnsheader), packet + len, content.begin());
  PacketReader pr(content);
  
  size_t idx = 0;
  DNSName rrname;
  uint16_t qdcount = ntohs(dh->qdcount);
  uint16_t ancount = ntohs(dh->ancount);
  uint16_t nscount = ntohs(dh->nscount);
  uint16_t arcount = ntohs(dh->arcount);
  uint16_t rrtype;
  uint16_t rrclass;
  string blob;
  struct dnsrecordheader ah;

  rrname = pr.getName();
  rrtype = pr.get16BitInt();
  rrclass = pr.get16BitInt();
  
  DNSPacketWriter pw(newContent, rrname, rrtype, rrclass, dh->opcode);
  pw.getHeader()->id=dh->id;
  pw.getHeader()->qr=dh->qr;
  pw.getHeader()->aa=dh->aa;
  pw.getHeader()->tc=dh->tc;
  pw.getHeader()->rd=dh->rd;
  pw.getHeader()->ra=dh->ra;
  pw.getHeader()->ad=dh->ad;
  pw.getHeader()->cd=dh->cd;
  pw.getHeader()->rcode=dh->rcode;
  
  /* consume remaining qd if any */
  if (qdcount > 1) {
    for(idx = 1; idx < qdcount; idx++) {
      rrname = pr.getName();
      rrtype = pr.get16BitInt();
      rrclass = pr.get16BitInt();
      (void) rrtype;
      (void) rrclass;
    }
  }

  /* copy AN and NS */
  for (idx = 0; idx < ancount; idx++) {
    rrname = pr.getName();
    pr.getDnsrecordheader(ah);

    pw.startRecord(rrname, ah.d_type, ah.d_ttl, ah.d_class, DNSResourceRecord::ANSWER, true);
    pr.xfrBlob(blob);
    pw.xfrBlob(blob);
  }

  for (idx = 0; idx < nscount; idx++) {
    rrname = pr.getName();
    pr.getDnsrecordheader(ah);

    pw.startRecord(rrname, ah.d_type, ah.d_ttl, ah.d_class, DNSResourceRecord::AUTHORITY, true);
    pr.xfrBlob(blob);
    pw.xfrBlob(blob);
  }
  /* consume AR, looking for OPT */
  for (idx = 0; idx < arcount; idx++) {
    rrname = pr.getName();
    pr.getDnsrecordheader(ah);

    if (ah.d_type != QType::OPT) {
      pw.startRecord(rrname, ah.d_type, ah.d_ttl, ah.d_class, DNSResourceRecord::ADDITIONAL, true);
      pr.xfrBlob(blob);
      pw.xfrBlob(blob);
    } else {
      pr.d_pos += ah.d_clen;
    }
  }
  pw.commit();

  return 0;
}

int locateEDNSOptRR(char * packet, const size_t len, char ** optStart, size_t * optLen, bool * last)
{
  assert(packet != NULL);
  assert(optStart != NULL);
  assert(optLen != NULL);
  assert(last != NULL);
  const struct dnsheader* dh = (const struct dnsheader*) packet;

  if (ntohs(dh->arcount) == 0)
    return ENOENT;

  vector<uint8_t> content(len - sizeof(dnsheader));
  copy(packet + sizeof(dnsheader), packet + len, content.begin());
  PacketReader pr(content);
  size_t idx = 0;
  DNSName rrname;
  uint16_t qdcount = ntohs(dh->qdcount);
  uint16_t ancount = ntohs(dh->ancount);
  uint16_t nscount = ntohs(dh->nscount);
  uint16_t arcount = ntohs(dh->arcount);
  uint16_t rrtype;
  uint16_t rrclass;
  struct dnsrecordheader ah;

  /* consume qd */
  for(idx = 0; idx < qdcount; idx++) {
    rrname = pr.getName();
    rrtype = pr.get16BitInt();
    rrclass = pr.get16BitInt();
    (void) rrtype;
    (void) rrclass;
  }

  /* consume AN and NS */
  for (idx = 0; idx < ancount + nscount; idx++) {
    rrname = pr.getName();
    pr.getDnsrecordheader(ah);
    pr.d_pos += ah.d_clen;
  }

  /* consume AR, looking for OPT */
  for (idx = 0; idx < arcount; idx++) {
    uint16_t start = pr.d_pos;
    rrname = pr.getName();
    pr.getDnsrecordheader(ah);

    if (ah.d_type == QType::OPT) {
      *optStart = packet + sizeof(dnsheader) + start;
      *optLen = (pr.d_pos - start) + ah.d_clen;

      if ((packet + len) < (*optStart + *optLen)) {
        throw std::range_error("Opt record overflow");
      }

      if (idx == ((size_t) arcount - 1)) {
        *last = true;
      }
      else {
        *last = false;
      }
      return 0;
    }
    pr.d_pos += ah.d_clen;
  }

  return ENOENT;
}

/* extract the start of the OPT RR in a QUERY packet if any */
static int getEDNSOptionsStart(char* packet, const size_t offset, const size_t len, char ** optRDLen, size_t * remaining)
{
  assert(packet != NULL);
  assert(optRDLen != NULL);
  assert(remaining != NULL);
  const struct dnsheader* dh = (const struct dnsheader*) packet;
  
  if (offset >= len)
    return ENOENT;

  if (ntohs(dh->qdcount) != 1 || dh->ancount != 0 || ntohs(dh->arcount) != 1 || dh->nscount != 0)
    return ENOENT;

  size_t pos = sizeof(dnsheader) + offset;
  pos += DNS_TYPE_SIZE + DNS_CLASS_SIZE;

  if (pos >= len)
    return ENOENT;

  uint16_t qtype, qclass;
  unsigned int consumed;
  DNSName aname(packet, len, pos, true, &qtype, &qclass, &consumed);

  if ((len - pos) < (consumed + DNS_TYPE_SIZE + DNS_CLASS_SIZE))
    return ENOENT;

  pos += consumed + DNS_TYPE_SIZE + DNS_CLASS_SIZE;
  if(qtype != QType::OPT || (len - pos) < (DNS_TTL_SIZE + DNS_RDLENGTH_SIZE))
    return ENOENT;

  pos += DNS_TTL_SIZE;
  *optRDLen = packet + pos;
  *remaining = len - pos;

  return 0;
}

static void generateECSOption(const ComboAddress& source, string& res, uint16_t ECSPrefixLength)
{
  Netmask sourceNetmask(source, ECSPrefixLength);
  EDNSSubnetOpts ecsOpts;
  ecsOpts.source = sourceNetmask;
  string payload = makeEDNSSubnetOptsString(ecsOpts);
  generateEDNSOption(EDNSOptionCode::ECS, payload, res);
}

void generateOptRR(const std::string& optRData, string& res)
{
  const uint8_t name = 0;
  dnsrecordheader dh;
  EDNS0Record edns0;
  edns0.extRCode = 0;
  edns0.version = 0;
  edns0.Z = 0;
  
  dh.d_type = htons(QType::OPT);
  dh.d_class = htons(g_EdnsUDPPayloadSize);
  static_assert(sizeof(EDNS0Record) == sizeof(dh.d_ttl), "sizeof(EDNS0Record) must match sizeof(dnsrecordheader.d_ttl)");
  memcpy(&dh.d_ttl, &edns0, sizeof edns0);
  dh.d_clen = htons((uint16_t) optRData.length());
  res.assign((const char *) &name, sizeof name);
  res.append((const char *) &dh, sizeof dh);
  res.append(optRData.c_str(), optRData.length());
}

static bool replaceEDNSClientSubnetOption(char * const packet, const size_t packetSize, uint16_t * const len, const ComboAddress& remote, char * const oldEcsOptionStart, size_t const oldEcsOptionSize, unsigned char * const optRDLen, uint16_t ECSPrefixLength)
{
  assert(packet != NULL);
  assert(len != NULL);
  assert(oldEcsOptionStart != NULL);
  assert(optRDLen != NULL);
  string ECSOption;
  generateECSOption(remote, ECSOption, ECSPrefixLength);

  if (ECSOption.size() == oldEcsOptionSize) {
    /* same size as the existing option */
    memcpy(oldEcsOptionStart, ECSOption.c_str(), oldEcsOptionSize);
  }
  else {
    /* different size than the existing option */
    const unsigned int newPacketLen = *len + (ECSOption.length() - oldEcsOptionSize);
    const size_t beforeOptionLen = oldEcsOptionStart - packet;
    const size_t dataBehindSize = *len - beforeOptionLen - oldEcsOptionSize;

    /* check that it fits in the existing buffer */
    if (newPacketLen > packetSize) {
      return false;
    }

    /* fix the size of ECS Option RDLen */
    uint16_t newRDLen = (optRDLen[0] * 256) + optRDLen[1];
    newRDLen += (ECSOption.size() - oldEcsOptionSize);
    optRDLen[0] = newRDLen / 256;
    optRDLen[1] = newRDLen % 256;

    if (dataBehindSize > 0) {
      memmove(oldEcsOptionStart, oldEcsOptionStart + oldEcsOptionSize, dataBehindSize);
    }
    memcpy(oldEcsOptionStart + dataBehindSize, ECSOption.c_str(), ECSOption.size());
    *len = newPacketLen;
  }

  return true;
}

bool handleEDNSClientSubnet(char* const packet, const size_t packetSize, const unsigned int consumed, uint16_t* const len, bool* const ednsAdded, bool* const ecsAdded, const ComboAddress& remote, bool overrideExisting, uint16_t ecsPrefixLength)
{
  assert(packet != NULL);
  assert(len != NULL);
  assert(consumed <= (size_t) *len);
  assert(ednsAdded != NULL);
  assert(ecsAdded != NULL);
  unsigned char * optRDLen = NULL;
  size_t remaining = 0;

  int res = getEDNSOptionsStart(packet, consumed, *len, (char**) &optRDLen, &remaining);
        
  if (res == 0) {
    char * ecsOptionStart = NULL;
    size_t ecsOptionSize = 0;
    
    res = getEDNSOption((char*)optRDLen, remaining, EDNSOptionCode::ECS, &ecsOptionStart, &ecsOptionSize);
    
    if (res == 0) {
      /* there is already an ECS value */
      if (overrideExisting) {
        return replaceEDNSClientSubnetOption(packet, packetSize, len, remote, ecsOptionStart, ecsOptionSize, optRDLen, ecsPrefixLength);
      }
    } else {
      /* we need to add one EDNS0 ECS option, fixing the size of EDNS0 RDLENGTH */
      /* getEDNSOptionsStart has already checked that there is exactly one AR,
         no NS and no AN */
      string ECSOption;
      generateECSOption(remote, ECSOption, ecsPrefixLength);
      const size_t ECSOptionSize = ECSOption.size();
      
      /* check if the existing buffer is large enough */
      if (packetSize - *len <= ECSOptionSize) {
        return false;
      }

      uint16_t newRDLen = (optRDLen[0] * 256) + optRDLen[1];
      newRDLen += ECSOptionSize;
      optRDLen[0] = newRDLen / 256;
      optRDLen[1] = newRDLen % 256;

      memcpy(packet + *len, ECSOption.c_str(), ECSOptionSize);
      *len += ECSOptionSize;
      *ecsAdded = true;
    }
  }
  else {
    /* we need to add a EDNS0 RR with one EDNS0 ECS option, fixing the AR count */
    string EDNSRR;
    struct dnsheader* dh = (struct dnsheader*) packet;
    string optRData;
    generateECSOption(remote, optRData, ecsPrefixLength);
    generateOptRR(optRData, EDNSRR);

    /* does it fit in the existing buffer? */
    if (packetSize - *len <= EDNSRR.size()) {
      return false;
    }

    uint16_t arcount = ntohs(dh->arcount);
    arcount++;
    dh->arcount = htons(arcount);
    *ednsAdded = true;

    memcpy(packet + *len, EDNSRR.c_str(), EDNSRR.size());
    *len += EDNSRR.size();
  }

  return true;
}

static int removeEDNSOptionFromOptions(unsigned char* optionsStart, const uint16_t optionsLen, const uint16_t optionCodeToRemove, uint16_t* newOptionsLen)
{
  unsigned char* p = optionsStart;
  size_t pos = 0;
  while ((pos + 4) <= optionsLen) {
    unsigned char* optionBegin = p;
    const uint16_t optionCode = 0x100*p[0] + p[1];
    p += sizeof(optionCode);
    pos += sizeof(optionCode);
    const uint16_t optionLen = 0x100*p[0] + p[1];
    p += sizeof(optionLen);
    pos += sizeof(optionLen);
    if ((pos + optionLen) > optionsLen) {
      return EINVAL;
    }
    if (optionCode == optionCodeToRemove) {
      if (pos + optionLen < optionsLen) {
        /* move remaining options over the removed one,
           if any */
        memmove(optionBegin, p + optionLen, optionsLen - (pos + optionLen));
      }
      *newOptionsLen = optionsLen - (sizeof(optionCode) + sizeof(optionLen) + optionLen);
      return 0;
    }
    p += optionLen;
    pos += optionLen;
  }
  return ENOENT;
}

int removeEDNSOptionFromOPT(char* optStart, size_t* optLen, const uint16_t optionCodeToRemove)
{
  /* we need at least:
     root label (1), type (2), class (2), ttl (4) + rdlen (2)*/
  if (*optLen < 11) {
    return EINVAL;
  }
  const unsigned char* end = (const unsigned char*) optStart + *optLen;
  unsigned char* p = (unsigned char*) optStart + 9;
  unsigned char* rdLenPtr = p;
  uint16_t rdLen = (0x100*p[0] + p[1]);
  p += sizeof(rdLen);
  if (p + rdLen != end) {
    return EINVAL;
  }
  uint16_t newRdLen = 0;
  int res = removeEDNSOptionFromOptions(p, rdLen, optionCodeToRemove, &newRdLen);
  if (res != 0) {
    return res;
  }
  *optLen -= (rdLen - newRdLen);
  rdLenPtr[0] = newRdLen / 0x100;
  rdLenPtr[1] = newRdLen % 0x100;
  return 0;
}

int rewriteResponseWithoutEDNSOption(const char * packet, const size_t len, const uint16_t optionCodeToSkip, vector<uint8_t>& newContent)
{
  assert(packet != NULL);
  assert(len >= sizeof(dnsheader));
  const struct dnsheader* dh = (const struct dnsheader*) packet;

  if (ntohs(dh->arcount) == 0)
    return ENOENT;

  if (ntohs(dh->qdcount) == 0)
    return ENOENT;

  vector<uint8_t> content(len - sizeof(dnsheader));
  copy(packet + sizeof(dnsheader), packet + len, content.begin());
  PacketReader pr(content);

  size_t idx = 0;
  DNSName rrname;
  uint16_t qdcount = ntohs(dh->qdcount);
  uint16_t ancount = ntohs(dh->ancount);
  uint16_t nscount = ntohs(dh->nscount);
  uint16_t arcount = ntohs(dh->arcount);
  uint16_t rrtype;
  uint16_t rrclass;
  string blob;
  struct dnsrecordheader ah;

  rrname = pr.getName();
  rrtype = pr.get16BitInt();
  rrclass = pr.get16BitInt();

  DNSPacketWriter pw(newContent, rrname, rrtype, rrclass, dh->opcode);
  pw.getHeader()->id=dh->id;
  pw.getHeader()->qr=dh->qr;
  pw.getHeader()->aa=dh->aa;
  pw.getHeader()->tc=dh->tc;
  pw.getHeader()->rd=dh->rd;
  pw.getHeader()->ra=dh->ra;
  pw.getHeader()->ad=dh->ad;
  pw.getHeader()->cd=dh->cd;
  pw.getHeader()->rcode=dh->rcode;

  /* consume remaining qd if any */
  if (qdcount > 1) {
    for(idx = 1; idx < qdcount; idx++) {
      rrname = pr.getName();
      rrtype = pr.get16BitInt();
      rrclass = pr.get16BitInt();
      (void) rrtype;
      (void) rrclass;
    }
  }

  /* copy AN and NS */
  for (idx = 0; idx < ancount; idx++) {
    rrname = pr.getName();
    pr.getDnsrecordheader(ah);

    pw.startRecord(rrname, ah.d_type, ah.d_ttl, ah.d_class, DNSResourceRecord::ANSWER, true);
    pr.xfrBlob(blob);
    pw.xfrBlob(blob);
  }

  for (idx = 0; idx < nscount; idx++) {
    rrname = pr.getName();
    pr.getDnsrecordheader(ah);

    pw.startRecord(rrname, ah.d_type, ah.d_ttl, ah.d_class, DNSResourceRecord::AUTHORITY, true);
    pr.xfrBlob(blob);
    pw.xfrBlob(blob);
  }

  /* consume AR, looking for OPT */
  for (idx = 0; idx < arcount; idx++) {
    rrname = pr.getName();
    pr.getDnsrecordheader(ah);

    if (ah.d_type != QType::OPT) {
      pw.startRecord(rrname, ah.d_type, ah.d_ttl, ah.d_class, DNSResourceRecord::ADDITIONAL, true);
      pr.xfrBlob(blob);
      pw.xfrBlob(blob);
    } else {
      pw.startRecord(rrname, ah.d_type, ah.d_ttl, ah.d_class, DNSResourceRecord::ADDITIONAL, false);
      pr.xfrBlob(blob);
      uint16_t rdLen = blob.length();
      removeEDNSOptionFromOptions((unsigned char*)blob.c_str(), rdLen, optionCodeToSkip, &rdLen);
      /* xfrBlob(string, size) completely ignores size.. */
      if (rdLen > 0) {
        blob.resize((size_t)rdLen);
        pw.xfrBlob(blob);
      } else {
        pw.commit();
      }
    }
  }
  pw.commit();

  return 0;
}
Commit	Line	Data
12471842 PL	1	/*
	2	* This file is part of PowerDNS or dnsdist.
	3	* Copyright -- PowerDNS.COM B.V. and its contributors
	4	*
	5	* This program is free software; you can redistribute it and/or modify
	6	* it under the terms of version 2 of the GNU General Public License as
	7	* published by the Free Software Foundation.
	8	*
	9	* In addition, for the avoidance of any doubt, permission is granted to
	10	* link this program with OpenSSL and to (re)distribute the binaries
	11	* produced as the result of such linking.
	12	*
	13	* This program is distributed in the hope that it will be useful,
	14	* but WITHOUT ANY WARRANTY; without even the implied warranty of
	15	* MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
	16	* GNU General Public License for more details.
	17	*
	18	* You should have received a copy of the GNU General Public License
	19	* along with this program; if not, write to the Free Software
	20	* Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.
	21	*/
ca404e94 RG	22	#include "dolog.hh"
	23	#include "dnsdist.hh"
	24	#include "dnsdist-ecs.hh"
	25	#include "dnsparser.hh"
	26	#include "dnswriter.hh"
5c3b5e7f	27	#include "ednsoptions.hh"
ca404e94 RG	28	#include "ednssubnet.hh"
	29
	30	/* when we add EDNS to a query, we don't want to advertise
	31	a large buffer size */
ff0902ec	32	size_t g_EdnsUDPPayloadSize = 512;
ca404e94 RG	33	/* draft-ietf-dnsop-edns-client-subnet-04 "11.1. Privacy" */
	34	uint16_t g_ECSSourcePrefixV4 = 24;
	35	uint16_t g_ECSSourcePrefixV6 = 56;
	36
	37	bool g_ECSOverride{false};
	38
	39	int rewriteResponseWithoutEDNS(const char * packet, const size_t len, vector<uint8_t>& newContent)
	40	{
	41	assert(packet != NULL);
	42	assert(len >= sizeof(dnsheader));
	43	const struct dnsheader* dh = (const struct dnsheader*) packet;
	44
	45	if (ntohs(dh->arcount) == 0)
	46	return ENOENT;
	47
	48	if (ntohs(dh->qdcount) == 0)
	49	return ENOENT;
	50
	51	vector<uint8_t> content(len - sizeof(dnsheader));
	52	copy(packet + sizeof(dnsheader), packet + len, content.begin());
	53	PacketReader pr(content);
	54
	55	size_t idx = 0;
	56	DNSName rrname;
	57	uint16_t qdcount = ntohs(dh->qdcount);
	58	uint16_t ancount = ntohs(dh->ancount);
	59	uint16_t nscount = ntohs(dh->nscount);
	60	uint16_t arcount = ntohs(dh->arcount);
	61	uint16_t rrtype;
	62	uint16_t rrclass;
	63	string blob;
	64	struct dnsrecordheader ah;
	65
	66	rrname = pr.getName();
	67	rrtype = pr.get16BitInt();
	68	rrclass = pr.get16BitInt();
	69
	70	DNSPacketWriter pw(newContent, rrname, rrtype, rrclass, dh->opcode);
	71	pw.getHeader()->id=dh->id;
	72	pw.getHeader()->qr=dh->qr;
	73	pw.getHeader()->aa=dh->aa;
	74	pw.getHeader()->tc=dh->tc;
	75	pw.getHeader()->rd=dh->rd;
	76	pw.getHeader()->ra=dh->ra;
	77	pw.getHeader()->ad=dh->ad;
	78	pw.getHeader()->cd=dh->cd;
	79	pw.getHeader()->rcode=dh->rcode;
	80
	81	/* consume remaining qd if any */
	82	if (qdcount > 1) {
	83	for(idx = 1; idx < qdcount; idx++) {
	84	rrname = pr.getName();
	85	rrtype = pr.get16BitInt();
	86	rrclass = pr.get16BitInt();
	87	(void) rrtype;
	88	(void) rrclass;
	89	}
	90	}
	91
	92	/* copy AN and NS */
	93	for (idx = 0; idx < ancount; idx++) {
	94	rrname = pr.getName();
	95	pr.getDnsrecordheader(ah);
	96
97	pw.startRecord(rrname, ah.d_type, ah.d_ttl, ah.d_class, DNSResourceRecord::ANSWER, true);
98	pr.xfrBlob(blob);
99	pw.xfrBlob(blob);
100	}
101
102	for (idx = 0; idx < nscount; idx++) {
103	rrname = pr.getName();
104	pr.getDnsrecordheader(ah);
105
106	pw.startRecord(rrname, ah.d_type, ah.d_ttl, ah.d_class, DNSResourceRecord::AUTHORITY, true);
107	pr.xfrBlob(blob);
108	pw.xfrBlob(blob);
109	}
ca404e94 RG	110	/* consume AR, looking for OPT */
	111	for (idx = 0; idx < arcount; idx++) {
	112	rrname = pr.getName();
	113	pr.getDnsrecordheader(ah);
	114
	115	if (ah.d_type != QType::OPT) {
	116	pw.startRecord(rrname, ah.d_type, ah.d_ttl, ah.d_class, DNSResourceRecord::ADDITIONAL, true);
	117	pr.xfrBlob(blob);
	118	pw.xfrBlob(blob);
	119	} else {
	120	pr.d_pos += ah.d_clen;
	121	}
	122	}
	123	pw.commit();
	124
	125	return 0;
	126	}
	127
ff73f02b	128	int locateEDNSOptRR(char * packet, const size_t len, char ** optStart, size_t * optLen, bool * last)
ca404e94 RG	129	{
	130	assert(packet != NULL);
	131	assert(optStart != NULL);
	132	assert(optLen != NULL);
	133	assert(last != NULL);
	134	const struct dnsheader* dh = (const struct dnsheader*) packet;
	135
	136	if (ntohs(dh->arcount) == 0)
	137	return ENOENT;
	138
	139	vector<uint8_t> content(len - sizeof(dnsheader));
	140	copy(packet + sizeof(dnsheader), packet + len, content.begin());
	141	PacketReader pr(content);
	142	size_t idx = 0;
	143	DNSName rrname;
	144	uint16_t qdcount = ntohs(dh->qdcount);
	145	uint16_t ancount = ntohs(dh->ancount);
	146	uint16_t nscount = ntohs(dh->nscount);
	147	uint16_t arcount = ntohs(dh->arcount);
	148	uint16_t rrtype;
	149	uint16_t rrclass;
	150	struct dnsrecordheader ah;
	151
	152	/* consume qd */
	153	for(idx = 0; idx < qdcount; idx++) {
	154	rrname = pr.getName();
	155	rrtype = pr.get16BitInt();
	156	rrclass = pr.get16BitInt();
	157	(void) rrtype;
	158	(void) rrclass;
	159	}
	160
	161	/* consume AN and NS */
	162	for (idx = 0; idx < ancount + nscount; idx++) {
	163	rrname = pr.getName();
	164	pr.getDnsrecordheader(ah);
	165	pr.d_pos += ah.d_clen;
	166	}
	167
	168	/* consume AR, looking for OPT */
	169	for (idx = 0; idx < arcount; idx++) {
	170	uint16_t start = pr.d_pos;
	171	rrname = pr.getName();
	172	pr.getDnsrecordheader(ah);
	173
	174	if (ah.d_type == QType::OPT) {
	175	*optStart = packet + sizeof(dnsheader) + start;
	176	*optLen = (pr.d_pos - start) + ah.d_clen;
	177
	178	if ((packet + len) < (optStart + optLen)) {
	179	throw std::range_error("Opt record overflow");
	180	}
	181
	182	if (idx == ((size_t) arcount - 1)) {
	183	*last = true;
	184	}
	185	else {
	186	*last = false;
	187	}
	188	return 0;
	189	}
	190	pr.d_pos += ah.d_clen;
	191	}
	192
193	return ENOENT;
194	}
195
196	/* extract the start of the OPT RR in a QUERY packet if any */
5c3b5e7f	197	static int getEDNSOptionsStart(char* packet, const size_t offset, const size_t len, char ** optRDLen, size_t * remaining)
ca404e94 RG	198	{
ca404e94 RG	199	assert(packet != NULL);
5c3b5e7f	200	assert(optRDLen != NULL);
ca404e94 RG	201	assert(remaining != NULL);
	202	const struct dnsheader* dh = (const struct dnsheader*) packet;
	203
	204	if (offset >= len)
	205	return ENOENT;
	206
	207	if (ntohs(dh->qdcount) != 1 \|\| dh->ancount != 0 \|\| ntohs(dh->arcount) != 1 \|\| dh->nscount != 0)
	208	return ENOENT;
	209
	210	size_t pos = sizeof(dnsheader) + offset;
	211	pos += DNS_TYPE_SIZE + DNS_CLASS_SIZE;
	212
	213	if (pos >= len)
	214	return ENOENT;
	215
	216	uint16_t qtype, qclass;
	217	unsigned int consumed;
	218	DNSName aname(packet, len, pos, true, &qtype, &qclass, &consumed);
	219
bd910269 RG	220	if ((len - pos) < (consumed + DNS_TYPE_SIZE + DNS_CLASS_SIZE))
	221	return ENOENT;
	222
ca404e94 RG	223	pos += consumed + DNS_TYPE_SIZE + DNS_CLASS_SIZE;
	224	if(qtype != QType::OPT \|\| (len - pos) < (DNS_TTL_SIZE + DNS_RDLENGTH_SIZE))
	225	return ENOENT;
	226
5c3b5e7f RG	227	pos += DNS_TTL_SIZE;
5c3b5e7f RG	228	*optRDLen = packet + pos;
ca404e94 RG	229	*remaining = len - pos;
ca404e94 RG	230
ca404e94 RG	231	return 0;
	232	}
	233
ff0902ec	234	static void generateECSOption(const ComboAddress& source, string& res, uint16_t ECSPrefixLength)
ca404e94	235	{
ff0902ec	236	Netmask sourceNetmask(source, ECSPrefixLength);
ca404e94 RG	237	EDNSSubnetOpts ecsOpts;
	238	ecsOpts.source = sourceNetmask;
	239	string payload = makeEDNSSubnetOptsString(ecsOpts);
5c3b5e7f	240	generateEDNSOption(EDNSOptionCode::ECS, payload, res);
ca404e94 RG	241	}
ca404e94 RG	242
01520028	243	void generateOptRR(const std::string& optRData, string& res)
ca404e94 RG	244	{
	245	const uint8_t name = 0;
	246	dnsrecordheader dh;
	247	EDNS0Record edns0;
	248	edns0.extRCode = 0;
	249	edns0.version = 0;
	250	edns0.Z = 0;
	251
	252	dh.d_type = htons(QType::OPT);
ff0902ec	253	dh.d_class = htons(g_EdnsUDPPayloadSize);
a683e8bd	254	static_assert(sizeof(EDNS0Record) == sizeof(dh.d_ttl), "sizeof(EDNS0Record) must match sizeof(dnsrecordheader.d_ttl)");
ca404e94	255	memcpy(&dh.d_ttl, &edns0, sizeof edns0);
ca404e94 RG	256	dh.d_clen = htons((uint16_t) optRData.length());
	257	res.assign((const char *) &name, sizeof name);
	258	res.append((const char *) &dh, sizeof dh);
	259	res.append(optRData.c_str(), optRData.length());
	260	}
	261
0beaa5c8	262	static bool replaceEDNSClientSubnetOption(char * const packet, const size_t packetSize, uint16_t * const len, const ComboAddress& remote, char * const oldEcsOptionStart, size_t const oldEcsOptionSize, unsigned char * const optRDLen, uint16_t ECSPrefixLength)
ca404e94 RG	263	{
	264	assert(packet != NULL);
	265	assert(len != NULL);
	266	assert(oldEcsOptionStart != NULL);
	267	assert(optRDLen != NULL);
	268	string ECSOption;
ff0902ec	269	generateECSOption(remote, ECSOption, ECSPrefixLength);
ca404e94 RG	270
	271	if (ECSOption.size() == oldEcsOptionSize) {
	272	/* same size as the existing option */
	273	memcpy(oldEcsOptionStart, ECSOption.c_str(), oldEcsOptionSize);
	274	}
	275	else {
	276	/* different size than the existing option */
	277	const unsigned int newPacketLen = *len + (ECSOption.length() - oldEcsOptionSize);
	278	const size_t beforeOptionLen = oldEcsOptionStart - packet;
	279	const size_t dataBehindSize = *len - beforeOptionLen - oldEcsOptionSize;
0beaa5c8 RG	280
	281	/* check that it fits in the existing buffer */
	282	if (newPacketLen > packetSize) {
	283	return false;
	284	}
	285
ca404e94	286	/* fix the size of ECS Option RDLen */
df367233	287	uint16_t newRDLen = (optRDLen[0] * 256) + optRDLen[1];
ca404e94	288	newRDLen += (ECSOption.size() - oldEcsOptionSize);
df367233 RG	289	optRDLen[0] = newRDLen / 256;
df367233 RG	290	optRDLen[1] = newRDLen % 256;
0beaa5c8 RG	291
	292	if (dataBehindSize > 0) {
	293	memmove(oldEcsOptionStart, oldEcsOptionStart + oldEcsOptionSize, dataBehindSize);
ca404e94	294	}
0beaa5c8 RG	295	memcpy(oldEcsOptionStart + dataBehindSize, ECSOption.c_str(), ECSOption.size());
0beaa5c8 RG	296	*len = newPacketLen;
ca404e94	297	}
0beaa5c8 RG	298
0beaa5c8 RG	299	return true;
ca404e94 RG	300	}
ca404e94 RG	301
0beaa5c8	302	bool handleEDNSClientSubnet(char* const packet, const size_t packetSize, const unsigned int consumed, uint16_t* const len, bool* const ednsAdded, bool* const ecsAdded, const ComboAddress& remote, bool overrideExisting, uint16_t ecsPrefixLength)
ca404e94 RG	303	{
	304	assert(packet != NULL);
	305	assert(len != NULL);
ca404e94 RG	306	assert(consumed <= (size_t) *len);
ca404e94 RG	307	assert(ednsAdded != NULL);
ff73f02b	308	assert(ecsAdded != NULL);
df367233	309	unsigned char * optRDLen = NULL;
ca404e94	310	size_t remaining = 0;
ff0902ec	311
5c3b5e7f	312	int res = getEDNSOptionsStart(packet, consumed, len, (char*) &optRDLen, &remaining);
ca404e94 RG	313
	314	if (res == 0) {
	315	char * ecsOptionStart = NULL;
	316	size_t ecsOptionSize = 0;
	317
5c3b5e7f	318	res = getEDNSOption((char*)optRDLen, remaining, EDNSOptionCode::ECS, &ecsOptionStart, &ecsOptionSize);
ca404e94 RG	319
ca404e94 RG	320	if (res == 0) {
5c3b5e7f	321	/* there is already an ECS value */
ff0902ec	322	if (overrideExisting) {
0beaa5c8	323	return replaceEDNSClientSubnetOption(packet, packetSize, len, remote, ecsOptionStart, ecsOptionSize, optRDLen, ecsPrefixLength);
ca404e94 RG	324	}
	325	} else {
	326	/* we need to add one EDNS0 ECS option, fixing the size of EDNS0 RDLENGTH */
	327	/* getEDNSOptionsStart has already checked that there is exactly one AR,
	328	no NS and no AN */
	329	string ECSOption;
ff0902ec	330	generateECSOption(remote, ECSOption, ecsPrefixLength);
ca404e94 RG	331	const size_t ECSOptionSize = ECSOption.size();
ca404e94 RG	332
0beaa5c8 RG	333	/* check if the existing buffer is large enough */
	334	if (packetSize - *len <= ECSOptionSize) {
	335	return false;
	336	}
	337
df367233	338	uint16_t newRDLen = (optRDLen[0] * 256) + optRDLen[1];
ca404e94	339	newRDLen += ECSOptionSize;
df367233 RG	340	optRDLen[0] = newRDLen / 256;
df367233 RG	341	optRDLen[1] = newRDLen % 256;
ca404e94	342
0beaa5c8 RG	343	memcpy(packet + *len, ECSOption.c_str(), ECSOptionSize);
0beaa5c8 RG	344	*len += ECSOptionSize;
ff73f02b	345	*ecsAdded = true;
ca404e94 RG	346	}
	347	}
	348	else {
	349	/* we need to add a EDNS0 RR with one EDNS0 ECS option, fixing the AR count */
	350	string EDNSRR;
	351	struct dnsheader* dh = (struct dnsheader*) packet;
01520028	352	string optRData;
ff0902ec	353	generateECSOption(remote, optRData, ecsPrefixLength);
01520028	354	generateOptRR(optRData, EDNSRR);
0beaa5c8 RG	355
	356	/* does it fit in the existing buffer? */
	357	if (packetSize - *len <= EDNSRR.size()) {
	358	return false;
	359	}
	360
ca404e94 RG	361	uint16_t arcount = ntohs(dh->arcount);
	362	arcount++;
	363	dh->arcount = htons(arcount);
	364	*ednsAdded = true;
	365
0beaa5c8 RG	366	memcpy(packet + *len, EDNSRR.c_str(), EDNSRR.size());
0beaa5c8 RG	367	*len += EDNSRR.size();
ca404e94	368	}
0beaa5c8 RG	369
0beaa5c8 RG	370	return true;
ca404e94	371	}
ff73f02b RG	372
	373	static int removeEDNSOptionFromOptions(unsigned char* optionsStart, const uint16_t optionsLen, const uint16_t optionCodeToRemove, uint16_t* newOptionsLen)
	374	{
	375	unsigned char* p = optionsStart;
09b3b357 RG	376	size_t pos = 0;
09b3b357 RG	377	while ((pos + 4) <= optionsLen) {
ff73f02b RG	378	unsigned char* optionBegin = p;
	379	const uint16_t optionCode = 0x100*p[0] + p[1];
	380	p += sizeof(optionCode);
09b3b357	381	pos += sizeof(optionCode);
ff73f02b RG	382	const uint16_t optionLen = 0x100*p[0] + p[1];
ff73f02b RG	383	p += sizeof(optionLen);
09b3b357 RG	384	pos += sizeof(optionLen);
09b3b357 RG	385	if ((pos + optionLen) > optionsLen) {
ff73f02b RG	386	return EINVAL;
	387	}
	388	if (optionCode == optionCodeToRemove) {
09b3b357	389	if (pos + optionLen < optionsLen) {
ff73f02b RG	390	/* move remaining options over the removed one,
ff73f02b RG	391	if any */
09b3b357	392	memmove(optionBegin, p + optionLen, optionsLen - (pos + optionLen));
ff73f02b RG	393	}
	394	*newOptionsLen = optionsLen - (sizeof(optionCode) + sizeof(optionLen) + optionLen);
	395	return 0;
	396	}
	397	p += optionLen;
09b3b357	398	pos += optionLen;
ff73f02b RG	399	}
	400	return ENOENT;
	401	}
	402
	403	int removeEDNSOptionFromOPT(char* optStart, size_t* optLen, const uint16_t optionCodeToRemove)
	404	{
	405	/* we need at least:
	406	root label (1), type (2), class (2), ttl (4) + rdlen (2)*/
	407	if (*optLen < 11) {
	408	return EINVAL;
	409	}
	410	const unsigned char* end = (const unsigned char) optStart + optLen;
	411	unsigned char* p = (unsigned char*) optStart + 9;
	412	unsigned char* rdLenPtr = p;
	413	uint16_t rdLen = (0x100*p[0] + p[1]);
	414	p += sizeof(rdLen);
	415	if (p + rdLen != end) {
	416	return EINVAL;
	417	}
	418	uint16_t newRdLen = 0;
	419	int res = removeEDNSOptionFromOptions(p, rdLen, optionCodeToRemove, &newRdLen);
	420	if (res != 0) {
	421	return res;
	422	}
	423	*optLen -= (rdLen - newRdLen);
	424	rdLenPtr[0] = newRdLen / 0x100;
	425	rdLenPtr[1] = newRdLen % 0x100;
	426	return 0;
	427	}
	428
	429	int rewriteResponseWithoutEDNSOption(const char * packet, const size_t len, const uint16_t optionCodeToSkip, vector<uint8_t>& newContent)
	430	{
	431	assert(packet != NULL);
	432	assert(len >= sizeof(dnsheader));
	433	const struct dnsheader* dh = (const struct dnsheader*) packet;
	434
	435	if (ntohs(dh->arcount) == 0)
	436	return ENOENT;
	437
	438	if (ntohs(dh->qdcount) == 0)
	439	return ENOENT;
	440
	441	vector<uint8_t> content(len - sizeof(dnsheader));
	442	copy(packet + sizeof(dnsheader), packet + len, content.begin());
	443	PacketReader pr(content);
	444
	445	size_t idx = 0;
	446	DNSName rrname;
	447	uint16_t qdcount = ntohs(dh->qdcount);
	448	uint16_t ancount = ntohs(dh->ancount);
	449	uint16_t nscount = ntohs(dh->nscount);
	450	uint16_t arcount = ntohs(dh->arcount);
	451	uint16_t rrtype;
	452	uint16_t rrclass;
	453	string blob;
	454	struct dnsrecordheader ah;
	455
	456	rrname = pr.getName();
	457	rrtype = pr.get16BitInt();
	458	rrclass = pr.get16BitInt();
	459
	460	DNSPacketWriter pw(newContent, rrname, rrtype, rrclass, dh->opcode);
	461	pw.getHeader()->id=dh->id;
	462	pw.getHeader()->qr=dh->qr;
463	pw.getHeader()->aa=dh->aa;
464	pw.getHeader()->tc=dh->tc;
465	pw.getHeader()->rd=dh->rd;
466	pw.getHeader()->ra=dh->ra;
467	pw.getHeader()->ad=dh->ad;
468	pw.getHeader()->cd=dh->cd;
469	pw.getHeader()->rcode=dh->rcode;
470
471	/* consume remaining qd if any */
472	if (qdcount > 1) {
473	for(idx = 1; idx < qdcount; idx++) {
474	rrname = pr.getName();
475	rrtype = pr.get16BitInt();
476	rrclass = pr.get16BitInt();
477	(void) rrtype;
478	(void) rrclass;
479	}
480	}
481
482	/* copy AN and NS */
483	for (idx = 0; idx < ancount; idx++) {
484	rrname = pr.getName();
485	pr.getDnsrecordheader(ah);
486
487	pw.startRecord(rrname, ah.d_type, ah.d_ttl, ah.d_class, DNSResourceRecord::ANSWER, true);
488	pr.xfrBlob(blob);
489	pw.xfrBlob(blob);
490	}
491
492	for (idx = 0; idx < nscount; idx++) {
493	rrname = pr.getName();
494	pr.getDnsrecordheader(ah);
495
496	pw.startRecord(rrname, ah.d_type, ah.d_ttl, ah.d_class, DNSResourceRecord::AUTHORITY, true);
497	pr.xfrBlob(blob);
498	pw.xfrBlob(blob);
499	}
500
501	/* consume AR, looking for OPT */
502	for (idx = 0; idx < arcount; idx++) {
503	rrname = pr.getName();
504	pr.getDnsrecordheader(ah);
505
506	if (ah.d_type != QType::OPT) {
507	pw.startRecord(rrname, ah.d_type, ah.d_ttl, ah.d_class, DNSResourceRecord::ADDITIONAL, true);
508	pr.xfrBlob(blob);
509	pw.xfrBlob(blob);
510	} else {
511	pw.startRecord(rrname, ah.d_type, ah.d_ttl, ah.d_class, DNSResourceRecord::ADDITIONAL, false);
512	pr.xfrBlob(blob);
513	uint16_t rdLen = blob.length();
514	removeEDNSOptionFromOptions((unsigned char*)blob.c_str(), rdLen, optionCodeToSkip, &rdLen);
515	/* xfrBlob(string, size) completely ignores size.. */
516	if (rdLen > 0) {
517	blob.resize((size_t)rdLen);
518	pw.xfrBlob(blob);
519	} else {
520	pw.commit();
521	}
522	}
523	}
524	pw.commit();
525
526	return 0;
527	}