// Copyright (c) 1999-2004 Brian Wellington (bwelling@xbill.org)

package org.xbill.DNS;

import java.util.*;
import java.io.*;
import java.net.*;

/**
 * An implementation of Resolver that sends one query to one server.
 * SimpleResolver handles TCP retries, transaction security (TSIG), and
 * EDNS 0.
 * @see Resolver
 * @see TSIG
 * @see OPTRecord
 *
 * @author Brian Wellington
 */


public class SimpleResolver implements Resolver {

/** The default port to send queries to */
public static final int DEFAULT_PORT = 53;

/** The default EDNS payload size */
public static final int DEFAULT_EDNS_PAYLOADSIZE = 1280;

private InetSocketAddress address;
private InetSocketAddress localAddress;
private boolean useTCP, ignoreTruncation;
private OPTRecord queryOPT;
private TSIG tsig;
private long timeoutValue = 10 * 1000;

private static final short DEFAULT_UDPSIZE = 512;

private static String defaultResolver = "localhost";
private static int uniqueID = 0;

/**
 * Creates a SimpleResolver that will query the specified host 
 * @exception UnknownHostException Failure occurred while finding the host
 */
public
SimpleResolver(String hostname) throws UnknownHostException {
	if (hostname == null) {
		hostname = ResolverConfig.getCurrentConfig().server();
		if (hostname == null)
			hostname = defaultResolver;
	}
	InetAddress addr;
	if (hostname.equals("0"))
		addr = InetAddress.getLocalHost();
	else
		addr = InetAddress.getByName(hostname);
	address = new InetSocketAddress(addr, DEFAULT_PORT);
}

/**
 * Creates a SimpleResolver.  The host to query is either found by using
 * ResolverConfig, or the default host is used.
 * @see ResolverConfig
 * @exception UnknownHostException Failure occurred while finding the host
 */
public
SimpleResolver() throws UnknownHostException {
	this(null);
}

/**
 * Gets the destination address associated with this SimpleResolver.
 * Messages sent using this SimpleResolver will be sent to this address.
 * @return The destination address associated with this SimpleResolver.
 */
InetSocketAddress
getAddress() {
	return address;
}

/** Sets the default host (initially localhost) to query */
public static void
setDefaultResolver(String hostname) {
	defaultResolver = hostname;
}

public void
setPort(int port) {
	address = new InetSocketAddress(address.getAddress(), port);
}

/**
 * Sets the address of the server to communicate with.
 * @param addr The address of the DNS server
 */
public void
setAddress(InetSocketAddress addr) {
	address = addr;
}

/**
 * Sets the address of the server to communicate with (on the default
 * DNS port)
 * @param addr The address of the DNS server
 */
public void
setAddress(InetAddress addr) {
	address = new InetSocketAddress(addr, address.getPort());
}

/**
 * Sets the local address to bind to when sending messages.
 * @param addr The local address to send messages from.
 */
public void
setLocalAddress(InetSocketAddress addr) {
	localAddress = addr;
}

/**
 * Sets the local address to bind to when sending messages.  A random port
 * will be used.
 * @param addr The local address to send messages from.
 */
public void
setLocalAddress(InetAddress addr) {
	localAddress = new InetSocketAddress(addr, 0);
}

public void
setTCP(boolean flag) {
	this.useTCP = flag;
}

public void
setIgnoreTruncation(boolean flag) {
	this.ignoreTruncation = flag;
}

public void
setEDNS(int level, int payloadSize, int flags, List options) {
	if (level != 0 && level != -1)
		throw new IllegalArgumentException("invalid EDNS level - " +
						   "must be 0 or -1");
	if (payloadSize == 0)
		payloadSize = DEFAULT_EDNS_PAYLOADSIZE;
	queryOPT = new OPTRecord(payloadSize, 0, level, flags, options);
}

public void
setEDNS(int level) {
	setEDNS(level, 0, 0, null);
}

public void
setTSIGKey(TSIG key) {
	tsig = key;
}

TSIG
getTSIGKey() {
	return tsig;
}

public void
setTimeout(int secs, int msecs) {
	timeoutValue = (long)secs * 1000 + msecs;
}

public void
setTimeout(int secs) {
	setTimeout(secs, 0);
}

long
getTimeout() {
	return timeoutValue;
}

private Message
parseMessage(byte [] b) throws WireParseException {
	try {
		return (new Message(b));
	}
	catch (IOException e) {
		if (Options.check("verbose"))
			e.printStackTrace();
		if (!(e instanceof WireParseException))
			e = new WireParseException("Error parsing message");
		throw (WireParseException) e;
	}
}

private void
verifyTSIG(Message query, Message response, byte [] b, TSIG tsig) {
	if (tsig == null)
		return;
	int error = tsig.verify(response, b, query.getTSIG());
	if (Options.check("verbose"))
		System.err.println("TSIG verify: " + Rcode.TSIGstring(error));
}

private void
applyEDNS(Message query) {
	if (queryOPT == null || query.getOPT() != null)
		return;
	query.addRecord(queryOPT, Section.ADDITIONAL);
}

private int
maxUDPSize(Message query) {
	OPTRecord opt = query.getOPT();
	if (opt == null)
		return DEFAULT_UDPSIZE;
	else
		return opt.getPayloadSize();
}

/**
 * Sends a message to a single server and waits for a response.  No checking
 * is done to ensure that the response is associated with the query.
 * @param query The query to send.
 * @return The response.
 * @throws IOException An error occurred while sending or receiving.
 */
public Message
send(Message query) throws IOException {
	if (Options.check("verbose"))
		System.err.println("Sending to " +
				   address.getAddress().getHostAddress() +
				   ":" + address.getPort());

	if (query.getHeader().getOpcode() == Opcode.QUERY) {
		Record question = query.getQuestion();
		if (question != null && question.getType() == Type.AXFR)
			return sendAXFR(query);
	}

	query = (Message) query.clone();
	applyEDNS(query);
	if (tsig != null)
		tsig.apply(query, null);

	byte [] out = query.toWire(Message.MAXLENGTH);
	int udpSize = maxUDPSize(query);
	boolean tcp = false;
	long endTime = System.currentTimeMillis() + timeoutValue;
	do {
		byte [] in;

		if (useTCP || out.length > udpSize)
			tcp = true;
		if (tcp)
			in = TCPClient.sendrecv(localAddress, address, out,
						endTime);
		else
			in = UDPClient.sendrecv(localAddress, address, out,
						udpSize, endTime);

		/*
		 * Check that the response is long enough.
		 */
		if (in.length < Header.LENGTH) {
			throw new WireParseException("invalid DNS header - " +
						     "too short");
		}
		/*
		 * Check that the response ID matches the query ID.  We want
		 * to check this before actually parsing the message, so that
		 * if there's a malformed response that's not ours, it
		 * doesn't confuse us.
		 */
		int id = ((in[0] & 0xFF) << 8) + (in[1] & 0xFF);
		int qid = query.getHeader().getID();
		if (id != qid) {
			String error = "invalid message id: expected " + qid +
				       "; got id " + id;
			if (tcp) {
				throw new WireParseException(error);
			} else {
				if (Options.check("verbose")) {
					System.err.println(error);
				}
				continue;
			}
		}
		Message response = parseMessage(in);
		verifyTSIG(query, response, in, tsig);
		if (!tcp && !ignoreTruncation &&
		    response.getHeader().getFlag(Flags.TC))
		{
			tcp = true;
			continue;
		}
		return response;
	} while (true);
}

/**
 * Asynchronously sends a message to a single server, registering a listener
 * to receive a callback on success or exception.  Multiple asynchronous
 * lookups can be performed in parallel.  Since the callback may be invoked
 * before the function returns, external synchronization is necessary.
 * @param query The query to send
 * @param listener The object containing the callbacks.
 * @return An identifier, which is also a parameter in the callback
 */
public Object
sendAsync(final Message query, final ResolverListener listener) {
	final Object id;
	synchronized (this) {
		id = new Integer(uniqueID++);
	}
	Record question = query.getQuestion();
	String qname;
	if (question != null)
		qname = question.getName().toString();
	else
		qname = "(none)";
	String name = this.getClass() + ": " + qname;
	Thread thread = new ResolveThread(this, query, id, listener);
	thread.setName(name);
	thread.setDaemon(true);
	thread.start();
	return id;
}

private Message
sendAXFR(Message query) throws IOException {
	Name qname = query.getQuestion().getName();
	ZoneTransferIn xfrin = ZoneTransferIn.newAXFR(qname, address, tsig);
	xfrin.setTimeout((int)(getTimeout() / 1000));
	xfrin.setLocalAddress(localAddress);
	try {
		xfrin.run();
	}
	catch (ZoneTransferException e) {
		throw new WireParseException(e.getMessage());
	}
	List records = xfrin.getAXFR();
	Message response = new Message(query.getHeader().getID());
	response.getHeader().setFlag(Flags.AA);
	response.getHeader().setFlag(Flags.QR);
	response.addRecord(query.getQuestion(), Section.QUESTION);
	Iterator it = records.iterator();
	while (it.hasNext())
		response.addRecord((Record)it.next(), Section.ANSWER);
	return response;
}

}