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.iterative;
012
013import java.net.InetAddress;
014import java.util.HashMap;
015import java.util.HashSet;
016import java.util.Set;
017
018import org.minidns.dnsmessage.DnsMessage;
019import org.minidns.dnsmessage.Question;
020import org.minidns.iterative.IterativeClientException.LoopDetected;
021import org.minidns.iterative.IterativeClientException.MaxIterativeStepsReached;
022
023public class ResolutionState {
024
025    private final IterativeDnsClient recursiveDnsClient;
026    private final HashMap<InetAddress, Set<Question>> map = new HashMap<>();
027    private int steps;
028
029    ResolutionState(IterativeDnsClient recursiveDnsClient) {
030        this.recursiveDnsClient = recursiveDnsClient;
031    }
032
033    void recurse(InetAddress address, DnsMessage query) throws LoopDetected, MaxIterativeStepsReached {
034        Question question = query.getQuestion();
035        if (!map.containsKey(address)) {
036            map.put(address, new HashSet<Question>());
037        } else if (map.get(address).contains(question)) {
038            throw new IterativeClientException.LoopDetected(address, question);
039        }
040
041        if (++steps > recursiveDnsClient.maxSteps) {
042            throw new IterativeClientException.MaxIterativeStepsReached(); 
043        }
044
045        boolean isNew = map.get(address).add(question);
046        assert isNew;
047    }
048
049    void decrementSteps() {
050        steps--;
051    }
052
053}