001/*
002 * Copyright 2015-2020 the original author or authors
003 *
004 * This software is licensed under the Apache License, Version 2.0,
005 * the GNU Lesser General Public License version 2 or later ("LGPL")
006 * and the WTFPL.
007 * You may choose either license to govern your use of this software only
008 * upon the condition that you accept all of the terms of either
009 * the Apache License 2.0, the LGPL 2.1+ or the WTFPL.
010 */
011package org.minidns.edns;
012
013import java.util.ArrayList;
014import java.util.Collections;
015import java.util.HashMap;
016import java.util.Iterator;
017import java.util.List;
018import java.util.Map;
019
020import org.minidns.dnsname.DnsName;
021import org.minidns.record.Data;
022import org.minidns.record.OPT;
023import org.minidns.record.Record;
024import org.minidns.record.Record.TYPE;
025
026/**
027 * EDNS - Extension Mechanism for DNS.
028 *
029 * @see <a href="https://tools.ietf.org/html/rfc6891">RFC 6891 - Extension Mechanisms for DNS (EDNS(0))</a>
030 *
031 */
032public class Edns {
033
034    /**
035     * Inform the dns server that the client supports DNSSEC.
036     */
037    public static final int FLAG_DNSSEC_OK = 0x8000;
038
039    /**
040     * The EDNS option code.
041     *
042     * @see <a href="http://www.iana.org/assignments/dns-parameters/dns-parameters.xhtml#dns-parameters-11">IANA - DNS EDNS0 Option Codes (OPT)</a>
043     */
044    public enum OptionCode {
045        UNKNOWN(-1, UnknownEdnsOption.class),
046        NSID(3, Nsid.class),
047        ;
048
049        private static Map<Integer, OptionCode> INVERSE_LUT = new HashMap<>(OptionCode.values().length);
050
051        static {
052            for (OptionCode optionCode : OptionCode.values()) {
053                INVERSE_LUT.put(optionCode.asInt, optionCode);
054            }
055        }
056
057        public final int asInt;
058        public final Class<? extends EdnsOption> clazz;
059
060        OptionCode(int optionCode, Class<? extends EdnsOption> clazz) {
061            this.asInt = optionCode;
062            this.clazz = clazz;
063        }
064
065        public static OptionCode from(int optionCode) {
066            OptionCode res = INVERSE_LUT.get(optionCode);
067            if (res == null) res = OptionCode.UNKNOWN;
068            return res;
069        }
070    }
071
072    public final int udpPayloadSize;
073
074    /**
075     * 8-bit extended return code.
076     *
077     * RFC 6891 § 6.1.3 EXTENDED-RCODE 
078     */
079    public final int extendedRcode;
080
081    /**
082     * 8-bit version field.
083     *
084     * RFC 6891 § 6.1.3 VERSION
085     */
086    public final int version;
087
088    /**
089     * 16-bit flags.
090     *
091     * RFC 6891 § 6.1.4
092     */
093    public final int flags;
094
095    public final List<EdnsOption> variablePart;
096
097    public final boolean dnssecOk;
098
099    private Record<OPT> optRecord;
100
101    public Edns(Record<OPT> optRecord) {
102        assert optRecord.type == TYPE.OPT;
103        udpPayloadSize = optRecord.clazzValue;
104        extendedRcode = (int) ((optRecord.ttl >> 8) & 0xff);
105        version = (int) ((optRecord.ttl >> 16) & 0xff);
106        flags = (int) optRecord.ttl & 0xffff;
107
108        dnssecOk = (optRecord.ttl & FLAG_DNSSEC_OK) > 0;
109
110        OPT opt = optRecord.payloadData;
111        variablePart = opt.variablePart;
112        this.optRecord = optRecord;
113    }
114
115    public Edns(Builder builder) {
116        udpPayloadSize = builder.udpPayloadSize;
117        extendedRcode = builder.extendedRcode;
118        version = builder.version;
119        int flags = 0;
120        if (builder.dnssecOk) {
121            flags |= FLAG_DNSSEC_OK;
122        }
123        dnssecOk = builder.dnssecOk;
124        this.flags = flags;
125        if (builder.variablePart != null) {
126            variablePart = builder.variablePart;
127        } else {
128            variablePart = Collections.emptyList();
129        }
130    }
131
132    @SuppressWarnings("unchecked")
133    public <O extends EdnsOption> O getEdnsOption(OptionCode optionCode) {
134        for (EdnsOption o : variablePart) {
135            if (o.getOptionCode().equals(optionCode)) {
136                return (O) o;
137            }
138        }
139        return null;
140    }
141
142    public Record<OPT> asRecord() {
143        if (optRecord == null) {
144            long optFlags = flags;
145            optFlags |= extendedRcode << 8;
146            optFlags |= version << 16;
147            optRecord = new Record<OPT>(DnsName.ROOT, Record.TYPE.OPT, udpPayloadSize, optFlags, new OPT(variablePart));
148        }
149        return optRecord;
150    }
151
152    private String terminalOutputCache;
153
154    public String asTerminalOutput() {
155        if (terminalOutputCache == null) {
156            StringBuilder sb = new StringBuilder();
157            sb.append("EDNS: version: ").append(version).append(", flags:");
158            if (dnssecOk)
159                sb.append(" do");
160            sb.append("; udp: ").append(udpPayloadSize);
161            if (!variablePart.isEmpty()) {
162                sb.append('\n');
163                Iterator<EdnsOption> it = variablePart.iterator();
164                while (it.hasNext()) {
165                    EdnsOption edns = it.next();
166                    sb.append(edns.getOptionCode()).append(": ");
167                    sb.append(edns.asTerminalOutput());
168                    if (it.hasNext()) {
169                        sb.append('\n');
170                    }
171                }
172            }
173            terminalOutputCache = sb.toString();
174        }
175        return terminalOutputCache;
176    }
177
178    @Override
179    public String toString() {
180        return asTerminalOutput();
181    }
182
183    public static Edns fromRecord(Record<? extends Data> record) {
184        if (record.type != TYPE.OPT) return null;
185
186        @SuppressWarnings("unchecked")
187        Record<OPT> optRecord = (Record<OPT>) record;
188        return new Edns(optRecord);
189    }
190
191    public static Builder builder() {
192        return new Builder();
193    }
194
195    public static final class Builder {
196        private int udpPayloadSize;
197        private int extendedRcode;
198        private int version;
199        private boolean dnssecOk;
200        private List<EdnsOption> variablePart;
201
202        private Builder() {
203        }
204
205        public Builder setUdpPayloadSize(int udpPayloadSize) {
206            if (udpPayloadSize > 0xffff) {
207                throw new IllegalArgumentException("UDP payload size must not be greater than 65536, was " + udpPayloadSize);
208            }
209            this.udpPayloadSize = udpPayloadSize;
210            return this;
211        }
212
213        public Builder setDnssecOk(boolean dnssecOk) {
214            this.dnssecOk = dnssecOk;
215            return this;
216        }
217
218        public Builder setDnssecOk() {
219            dnssecOk = true;
220            return this;
221        }
222
223        public Builder addEdnsOption(EdnsOption ednsOption) {
224            if (variablePart == null) {
225                variablePart = new ArrayList<>(4);
226            }
227            variablePart.add(ednsOption);
228            return this;
229        }
230
231        public Edns build() {
232            return new Edns(this);
233        }
234    }
235}