org.apache.hbase.thirdparty.io.netty.resolver.dns.DnsQueryContextManager Maven / Gradle / Ivy

/*
 * Copyright 2015 The Netty Project
 *
 * The Netty Project licenses this file to you under the Apache License,
 * version 2.0 (the "License"); you may not use this file except in compliance
 * with the License. You may obtain a copy of the License at:
 *
 *   https://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
 * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
 * License for the specific language governing permissions and limitations
 * under the License.
 */

package org.apache.hbase.thirdparty.io.netty.resolver.dns;

import org.apache.hbase.thirdparty.io.netty.util.NetUtil;
import org.apache.hbase.thirdparty.io.netty.util.collection.IntObjectHashMap;
import org.apache.hbase.thirdparty.io.netty.util.collection.IntObjectMap;

import java.net.Inet4Address;
import java.net.Inet6Address;
import java.net.InetAddress;
import java.net.InetSocketAddress;
import java.net.UnknownHostException;
import java.util.HashMap;
import java.util.Map;

final class DnsQueryContextManager {

    /**
     * A map whose key is the DNS server address and value is the map of the DNS query ID and its corresponding
     * {@link DnsQueryContext}.
     */
    private final Map map =
            new HashMap();

    /**
     * Add {@link DnsQueryContext} to the context manager and return the ID that should be used for the query.
     * This method will return {@code -1} if an ID could not be generated and the context was not stored.
     *
     * @param nameServerAddr    The {@link InetSocketAddress} of the nameserver to query.
     * @param qCtx              The {@link {@link DnsQueryContext} to store.
     * @return                  the ID that should be used or {@code -1} if none could be generated.
     */
    int add(InetSocketAddress nameServerAddr, DnsQueryContext qCtx) {
        final DnsQueryContextMap contexts = getOrCreateContextMap(nameServerAddr);
        return contexts.add(qCtx);
    }

    /**
     * Return the {@link DnsQueryContext} for the given {@link InetSocketAddress} and id or {@code null} if
     * none could be found.
     *
     * @param nameServerAddr    The {@link InetSocketAddress} of the nameserver.
     * @param id                The id that identifies the {@link DnsQueryContext} and was used for the query.
     * @return                  The context or {@code null} if none could be found.
     */
    DnsQueryContext get(InetSocketAddress nameServerAddr, int id) {
        final DnsQueryContextMap contexts = getContextMap(nameServerAddr);
        if (contexts == null) {
            return null;
        }
        return contexts.get(id);
    }

    /**
     * Remove the {@link DnsQueryContext} for the given {@link InetSocketAddress} and id or {@code null} if
     * none could be found.
     *
     * @param nameServerAddr    The {@link InetSocketAddress} of the nameserver.
     * @param id                The id that identifies the {@link DnsQueryContext} and was used for the query.
     * @return                  The context or {@code null} if none could be removed.
     */
    DnsQueryContext remove(InetSocketAddress nameServerAddr, int id) {
        final DnsQueryContextMap contexts = getContextMap(nameServerAddr);
        if (contexts == null) {
            return null;
        }
        return contexts.remove(id);
    }

    private DnsQueryContextMap getContextMap(InetSocketAddress nameServerAddr) {
        synchronized (map) {
            return map.get(nameServerAddr);
        }
    }

    private DnsQueryContextMap getOrCreateContextMap(InetSocketAddress nameServerAddr) {
        synchronized (map) {
            final DnsQueryContextMap contexts = map.get(nameServerAddr);
            if (contexts != null) {
                return contexts;
            }

            final DnsQueryContextMap newContexts = new DnsQueryContextMap();
            final InetAddress a = nameServerAddr.getAddress();
            final int port = nameServerAddr.getPort();
            DnsQueryContextMap old = map.put(nameServerAddr, newContexts);
            // Assert that we didn't replace an existing mapping.
            assert old == null : "DnsQueryContextMap already exists for " + nameServerAddr;

            InetSocketAddress extraAddress = null;
            if (a instanceof Inet4Address) {
                // Also add the mapping for the IPv4-compatible IPv6 address.
                final Inet4Address a4 = (Inet4Address) a;
                if (a4.isLoopbackAddress()) {
                    extraAddress = new InetSocketAddress(NetUtil.LOCALHOST6, port);
                } else {
                    extraAddress = new InetSocketAddress(toCompactAddress(a4), port);
                }
            } else if (a instanceof Inet6Address) {
                // Also add the mapping for the IPv4 address if this IPv6 address is compatible.
                final Inet6Address a6 = (Inet6Address) a;
                if (a6.isLoopbackAddress()) {
                    extraAddress = new InetSocketAddress(NetUtil.LOCALHOST4, port);
                } else if (a6.isIPv4CompatibleAddress()) {
                    extraAddress = new InetSocketAddress(toIPv4Address(a6), port);
                }
            }
            if (extraAddress != null) {
                old = map.put(extraAddress, newContexts);
                // Assert that we didn't replace an existing mapping.
                assert old == null : "DnsQueryContextMap already exists for " + extraAddress;
            }

            return newContexts;
        }
    }

    private static Inet6Address toCompactAddress(Inet4Address a4) {
        byte[] b4 = a4.getAddress();
        byte[] b6 = { 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, b4[0], b4[1], b4[2], b4[3] };
        try {
            return (Inet6Address) InetAddress.getByAddress(b6);
        } catch (UnknownHostException e) {
            throw new Error(e);
        }
    }

    private static Inet4Address toIPv4Address(Inet6Address a6) {
        assert a6.isIPv4CompatibleAddress();

        byte[] b6 = a6.getAddress();
        byte[] b4 = { b6[12], b6[13], b6[14], b6[15] };
        try {
            return (Inet4Address) InetAddress.getByAddress(b4);
        } catch (UnknownHostException e) {
            throw new Error(e);
        }
    }

    private static final class DnsQueryContextMap {

        private final DnsQueryIdSpace idSpace = new DnsQueryIdSpace();

        // We increment on every usage so start with -1, this will ensure we start with 0 as first id.
        private final IntObjectMap map = new IntObjectHashMap();

        synchronized int add(DnsQueryContext ctx) {
            int id = idSpace.nextId();
            DnsQueryContext oldCtx = map.put(id, ctx);
            assert oldCtx == null;
            return id;
        }

        synchronized DnsQueryContext get(int id) {
            return map.get(id);
        }

        synchronized DnsQueryContext remove(int id) {
            idSpace.pushId(id);
            return map.remove(id);
        }
    }
}