uk.num.numlib.internal.dns.DNSServicesDefaultImpl Maven / Gradle / Ivy

Go to download
package uk.num.numlib.internal.dns;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.xbill.DNS.*;
import uk.num.numlib.exc.NumDNSQueryException;
import uk.num.numlib.exc.NumInvalidDNSQueryException;
import uk.num.numlib.exc.NumNoRecordAvailableException;
import uk.num.numlib.exc.NumNotImplementedException;
import uk.num.numlib.internal.ctx.AppContext;

import java.io.IOException;
import java.util.Arrays;
import java.util.List;
import java.util.function.Predicate;

/**
 * A default implementation of the DNSServices interface.
 *
 * @author tonywalmsley
 */
public class DNSServicesDefaultImpl implements DNSServices {
    private static final Logger LOG = LoggerFactory.getLogger(DNSServicesDefaultImpl.class);

    /**
     * Is a record an SPF or CNAME record?
     */
    private static Predicate isCNAMEOrSPFRecord = (r) -> r.rdataToString()
            .startsWith("v=spf") || r.rdataToString()
            .startsWith("\"v=spf") || r.getType() == Type.CNAME || r.getType() == Type.SPF;

    /**
     * Get the Module Configuration from DNS as an array of Records
     *
     * @param appContext    the AppContext
     * @param moduleId      The module Id String
     * @param timeoutMillis The number of milliseconds to wait for a response.
     * @return An array of Records
     * @throws NumInvalidDNSQueryException on error
     * @throws NumDNSQueryException        on error
     */
    public Record[] getConfigFileTXTRecords(final AppContext appContext, final String moduleId, final int timeoutMillis) throws
                                                                                                                         NumInvalidDNSQueryException,
                                                                                                                         NumDNSQueryException {
        assert moduleId != null && moduleId.trim()
                .length() > 0;
        assert timeoutMillis > 0;

        Record[] records;
        final Resolver resolver = Lookup.getDefaultResolver();
        resolver.setTimeout(timeoutMillis / 1000, timeoutMillis % 1000);
        resolver.setIgnoreTruncation(false);

        final String query = moduleId + appContext.stringConstants.CONFIG_FILE_SUFFIX();
        try {
            LOG.trace("Retrieving configuration file from DNS using query: {}", query);
            final Lookup lookup = new Lookup(query, Type.TXT);
            records = lookup.run();

            if (lookup.getResult() != Lookup.SUCCESSFUL) {
                LOG.error("Error retrieving configuration file from DNS: {} for query {}", lookup.getErrorString(), query);
                throw new NumDNSQueryException("Error retrieving configuration file from DNS: " + lookup.getErrorString() + " for query: " + query);
            }
        } catch (TextParseException e) {
            LOG.error("Error parsing DNS query.", e);
            throw new NumInvalidDNSQueryException(query);
        }

        return records;
    }

    /**
     * Concatenate an array of TXT record values to a single String
     *
     * @param records The array of Records
     * @return The concatenated result.
     */
    public String rebuildTXTRecordContent(final Record[] records) {
        assert records != null && records.length > 0;
        final StringBuilder buffer = new StringBuilder();

        for (Record record : records) {
            final TXTRecord txt = (TXTRecord) record;
            final List strings = txt.getStrings();

            // Append all of the String parts of the TXT record to a single String via the StringBuffer
            for (Object string : strings) {
                buffer.append(string);
            }
        }

        LOG.trace("Concatenated {} records to get: {}", records.length, buffer.toString());
        return buffer.toString();
    }

    /**
     * Get a NUM record from DNS without caching.
     *
     * @param query               The NUM formatted DNS query.
     * @param timeoutMillis       The number of milliseconds to wait for a response.
     * @param checkDnsSecValidity true if the result should be checked for DNSSEC validity.
     * @return An array of Records
     * @throws NumNotImplementedException    on error
     * @throws NumInvalidDNSQueryException   on error
     * @throws NumNoRecordAvailableException if a CNAME or SPF record is received instead of a TXT record
     */
    @Override
    public Record[] getRecordFromDnsNoCache(final String query, final int timeoutMillis, final boolean checkDnsSecValidity) throws
                                                                                                                            NumNotImplementedException,
                                                                                                                            NumInvalidDNSQueryException,
                                                                                                                            NumNoRecordAvailableException {
        assert query != null && query.trim()
                .length() > 0;
        assert timeoutMillis > 0;

        // Return the cached value if we have one.
        Record[] records;

        // No cached value so look it up in DNS
        final Resolver resolver = Lookup.getDefaultResolver();
        resolver.setTimeout(timeoutMillis / 1000, timeoutMillis % 1000);
        resolver.setIgnoreTruncation(false);

        try {
            final Record queryTxtRecord = Record.newRecord(new Name(query), Type.TXT, DClass.IN);
            final Message queryMessage = Message.newQuery(queryTxtRecord);
            LOG.debug("Sending DNS Query: {}", queryMessage);

            final Message response = resolver.send(queryMessage);
            LOG.debug("Received DNS Response: {}", response);

            if (response.getRcode() == Rcode.NOERROR) {
                if (checkDnsSecValidity) {
                    LOG.error("DNSSEC checks not implemented.");
                    throw new NumNotImplementedException("DNSSEC checks not implemented.");
                }
                records = response.getSectionArray(Section.ANSWER);
                if (Arrays.stream(records)
                        .anyMatch(isCNAMEOrSPFRecord)) {
                    throw new NumNoRecordAvailableException("Received CNAME or SPF record instead of TXT record.");
                }
            } else {
                if (response.getHeader()
                        .getFlag(Flags.TC)) {
                    LOG.info("Response truncated, possible Multi-part record.");
                    throw new PossibleMultiPartRecordException();
                }
                return null;
            }
        } catch (IOException e) {
            LOG.error("Error querying for NUM record.", e);
            throw new NumInvalidDNSQueryException("Invalid DNS query: " + query);
        }

        return records;
    }
}