1// Copyright (c) 2005 Brian Wellington (bwelling@xbill.org)
2
3package org.xbill.DNS;
4
5import java.io.*;
6import java.net.*;
7import java.nio.*;
8import java.nio.channels.*;
9
10final class TCPClient extends Client {
11
12public
13TCPClient(long endTime) throws IOException {
14	super(SocketChannel.open(), endTime);
15}
16
17void
18bind(SocketAddress addr) throws IOException {
19	SocketChannel channel = (SocketChannel) key.channel();
20	channel.socket().bind(addr);
21}
22
23void
24connect(SocketAddress addr) throws IOException {
25	SocketChannel channel = (SocketChannel) key.channel();
26	if (channel.connect(addr))
27		return;
28	key.interestOps(SelectionKey.OP_CONNECT);
29	try {
30		while (!channel.finishConnect()) {
31			if (!key.isConnectable())
32				blockUntil(key, endTime);
33		}
34	}
35	finally {
36		if (key.isValid())
37			key.interestOps(0);
38	}
39}
40
41void
42send(byte [] data) throws IOException {
43	SocketChannel channel = (SocketChannel) key.channel();
44	verboseLog("TCP write", data);
45	byte [] lengthArray = new byte[2];
46	lengthArray[0] = (byte)(data.length >>> 8);
47	lengthArray[1] = (byte)(data.length & 0xFF);
48	ByteBuffer [] buffers = new ByteBuffer[2];
49	buffers[0] = ByteBuffer.wrap(lengthArray);
50	buffers[1] = ByteBuffer.wrap(data);
51	int nsent = 0;
52	key.interestOps(SelectionKey.OP_WRITE);
53	try {
54		while (nsent < data.length + 2) {
55			if (key.isWritable()) {
56				long n = channel.write(buffers);
57				if (n < 0)
58					throw new EOFException();
59				nsent += (int) n;
60				if (nsent < data.length + 2 &&
61				    System.currentTimeMillis() > endTime)
62					throw new SocketTimeoutException();
63			} else
64				blockUntil(key, endTime);
65		}
66	}
67	finally {
68		if (key.isValid())
69			key.interestOps(0);
70	}
71}
72
73private byte []
74_recv(int length) throws IOException {
75	SocketChannel channel = (SocketChannel) key.channel();
76	int nrecvd = 0;
77	byte [] data = new byte[length];
78	ByteBuffer buffer = ByteBuffer.wrap(data);
79	key.interestOps(SelectionKey.OP_READ);
80	try {
81		while (nrecvd < length) {
82			if (key.isReadable()) {
83				long n = channel.read(buffer);
84				if (n < 0)
85					throw new EOFException();
86				nrecvd += (int) n;
87				if (nrecvd < length &&
88				    System.currentTimeMillis() > endTime)
89					throw new SocketTimeoutException();
90			} else
91				blockUntil(key, endTime);
92		}
93	}
94	finally {
95		if (key.isValid())
96			key.interestOps(0);
97	}
98	return data;
99}
100
101byte []
102recv() throws IOException {
103	byte [] buf = _recv(2);
104	int length = ((buf[0] & 0xFF) << 8) + (buf[1] & 0xFF);
105	byte [] data = _recv(length);
106	verboseLog("TCP read", data);
107	return data;
108}
109
110static byte []
111sendrecv(SocketAddress local, SocketAddress remote, byte [] data, long endTime)
112throws IOException
113{
114	TCPClient client = new TCPClient(endTime);
115	try {
116		if (local != null)
117			client.bind(local);
118		client.connect(remote);
119		client.send(data);
120		return client.recv();
121	}
122	finally {
123		client.cleanup();
124	}
125}
126
127static byte []
128sendrecv(SocketAddress addr, byte [] data, long endTime) throws IOException {
129	return sendrecv(null, addr, data, endTime);
130}
131
132}
133