/* This code is part of Freenet. It is distributed under the GNU General
 * Public License, version 2 (or at your option any later version). See
 * http://www.gnu.org/ for further details of the GPL. */
package freenet.node;

import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertNotNull;
import static org.junit.Assert.assertTrue;
import static org.junit.Assert.fail;

import java.net.UnknownHostException;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Random;

import freenet.crypt.BlockCipher;
import freenet.crypt.DummyRandomSource;
import freenet.crypt.ciphers.Rijndael;
import freenet.io.comm.FreenetInetAddress;
import freenet.io.comm.Peer;
import org.junit.Before;
import org.junit.Test;

public class NewPacketFormatTest {
	@Before
	public void setUp() {
		// Because we don't call maybeSendPacket, the packet sent times are not updated,
		// so lets turn off the keepalives.
		NewPacketFormat.DO_KEEPALIVES = false;
	}
	
	@Test
	public void testEmptyCreation() throws BlockedTooLongException {
		NewPacketFormat npf = new NewPacketFormat(null, 0, 0);
		PeerMessageQueue pmq = new PeerMessageQueue(new DummyRandomSource(1234));
		SessionKey s = new SessionKey(null, null, null, null, null, null, null, null, new NewPacketFormatKeyContext(0, 0), 1);

		NPFPacket p = npf.createPacket(1400, pmq, s, false);
		if(p != null) fail("Created packet from nothing");
	}

	@Test
	public void testAckOnlyCreation() throws BlockedTooLongException, InterruptedException {
		BasePeerNode pn = new NullBasePeerNode();
		NewPacketFormat npf = new NewPacketFormat(pn, 0, 0);
		PeerMessageQueue pmq = new PeerMessageQueue(new DummyRandomSource(1234));
		SessionKey s = new SessionKey(null, null, null, null, null, null, null, null, new NewPacketFormatKeyContext(0, 0), 1);

		NPFPacket p = null;

		//Packet that should be acked
		p = new NPFPacket();
		p.addMessageFragment(new MessageFragment(true, false, true, 0, 8, 8, 0, new byte[] {(byte) 0x01,
		                (byte) 0x23, (byte) 0x45, (byte) 0x67, (byte) 0x89, (byte) 0xAB, (byte) 0xCD,
		                (byte) 0xEF }, null));
		assertEquals(1, npf.handleDecryptedPacket(p, s).size());

		Thread.sleep(NewPacketFormatKeyContext.MAX_ACK_DELAY*2);
		p = npf.createPacket(1400, pmq, s, false);
		assertEquals(1, p.getAcks().size());
	}

	@Test
	public void testLostLastAck() throws BlockedTooLongException, InterruptedException {
		NullBasePeerNode senderNode = new NullBasePeerNode();
		NewPacketFormat sender = new NewPacketFormat(senderNode, 0, 0);
		PeerMessageQueue senderQueue = new PeerMessageQueue(new DummyRandomSource(1234));
		NullBasePeerNode receiverNode = new NullBasePeerNode();
		NewPacketFormat receiver = new NewPacketFormat(receiverNode, 0, 0);
		PeerMessageQueue receiverQueue = new PeerMessageQueue(new DummyRandomSource(1234));
		SessionKey senderKey = new SessionKey(null, null, null, null, null, null, null, null, new NewPacketFormatKeyContext(0, 0), 1);
		senderNode.currentKey = senderKey;
		SessionKey receiverKey = new SessionKey(null, null, null, null, null, null, null, null, new NewPacketFormatKeyContext(0, 0), 1);

		senderQueue.queueAndEstimateSize(new MessageItem(new byte[1024], null, false, null, (short) 0), 1024);

		NPFPacket fragment1 = sender.createPacket(512, senderQueue, senderKey, false);
		assertEquals(1, fragment1.getFragments().size());
		receiver.handleDecryptedPacket(fragment1, receiverKey);

		NPFPacket fragment2 = sender.createPacket(512, senderQueue, senderKey, false);
		assertEquals(1, fragment2.getFragments().size());
		receiver.handleDecryptedPacket(fragment2, receiverKey);

		Thread.sleep(NewPacketFormatKeyContext.MAX_ACK_DELAY*2);
		NPFPacket ack1 = receiver.createPacket(512, receiverQueue, receiverKey, false);
		assertEquals(2, ack1.getAcks().size());
		assertEquals(0, (int)ack1.getAcks().first());
		assertEquals(1, (int)ack1.getAcks().last());
		sender.handleDecryptedPacket(ack1, senderKey);

		NPFPacket fragment3 = sender.createPacket(512, senderQueue, senderKey, false);
		assertEquals(1, fragment3.getFragments().size());
		receiver.handleDecryptedPacket(fragment3, receiverKey);
		Thread.sleep(NewPacketFormatKeyContext.MAX_ACK_DELAY*2);
		receiver.createPacket(512, senderQueue, receiverKey, false); //Sent, but lost

		try {
			Thread.sleep(2000); //RTT is 250ms by default since there is no PeerNode to track it
		} catch (InterruptedException e) { fail(); }

		NPFPacket resend1 = sender.createPacket(512, senderQueue, senderKey, false);
		if(resend1 == null) fail("No packet to resend");
		assertEquals(0, receiver.handleDecryptedPacket(resend1, receiverKey).size());

		//Make sure an ack is sent
		Thread.sleep(NewPacketFormatKeyContext.MAX_ACK_DELAY*2);
		NPFPacket ack2 = receiver.createPacket(512, receiverQueue, receiverKey, false);
		assertNotNull(ack2);
		assertEquals(1, ack2.getAcks().size());
		assertEquals(0, ack2.getFragments().size());
	}

	@Test
	public void testOutOfOrderDelivery() throws BlockedTooLongException {
		NullBasePeerNode senderNode = new NullBasePeerNode();
		NewPacketFormat sender = new NewPacketFormat(senderNode, 0, 0);
		PeerMessageQueue senderQueue = new PeerMessageQueue(new DummyRandomSource(1234));
		NullBasePeerNode receiverNode = new NullBasePeerNode();
		NewPacketFormat receiver = new NewPacketFormat(receiverNode, 0, 0);
		SessionKey senderKey = new SessionKey(null, null, null, null, null, null, null, null, new NewPacketFormatKeyContext(0, 0), 1);
		SessionKey receiverKey = new SessionKey(null, null, null, null, null, null, null, null, new NewPacketFormatKeyContext(0, 0), 1);

		senderQueue.queueAndEstimateSize(new MessageItem(new byte[1024], null, false, null, (short) 0), 1024);

		NPFPacket fragment1 = sender.createPacket(512, senderQueue, senderKey, false);
		assertEquals(1, fragment1.getFragments().size());

		NPFPacket fragment2 = sender.createPacket(512, senderQueue, senderKey, false);
		assertEquals(1, fragment2.getFragments().size());

		NPFPacket fragment3 = sender.createPacket(512, senderQueue, senderKey, false);
		assertEquals(1, fragment3.getFragments().size());

		receiver.handleDecryptedPacket(fragment1, receiverKey);
		receiver.handleDecryptedPacket(fragment3, receiverKey);
		assertEquals(1, receiver.handleDecryptedPacket(fragment2, receiverKey).size());
	}

	@Test
	public void testReceiveUnknownMessageLength() throws BlockedTooLongException {
		NullBasePeerNode senderNode = new NullBasePeerNode();
		NewPacketFormat sender = new NewPacketFormat(senderNode, 0, 0);
		PeerMessageQueue senderQueue = new PeerMessageQueue(new DummyRandomSource(1234));
		NullBasePeerNode receiverNode = new NullBasePeerNode();
		NewPacketFormat receiver = new NewPacketFormat(receiverNode, 0, 0);
		SessionKey senderKey = new SessionKey(null, null, null, null, null, null, null, null, new NewPacketFormatKeyContext(0, 0), 1);
		SessionKey receiverKey = new SessionKey(null, null, null, null, null, null, null, null, new NewPacketFormatKeyContext(0, 0), 1);

		senderQueue.queueAndEstimateSize(new MessageItem(new byte[1024], null, false, null, (short) 0), 1024);

		NPFPacket fragment1 = sender.createPacket(512, senderQueue, senderKey, false);
		assertEquals(1, fragment1.getFragments().size());
		NPFPacket fragment2 = sender.createPacket(512, senderQueue, senderKey, false);
		assertEquals(1, fragment2.getFragments().size());
		NPFPacket fragment3 = sender.createPacket(512, senderQueue, senderKey, false);
		assertEquals(1, fragment3.getFragments().size());

		receiver.handleDecryptedPacket(fragment3, receiverKey);
		receiver.handleDecryptedPacket(fragment2, receiverKey);
		assertEquals(1, receiver.handleDecryptedPacket(fragment1, receiverKey).size());
	}

	@Test
	public void testResendAlreadyCompleted() throws BlockedTooLongException, InterruptedException {
		NullBasePeerNode senderNode = new NullBasePeerNode();
		NewPacketFormat sender = new NewPacketFormat(senderNode, 0, 0);
		PeerMessageQueue senderQueue = new PeerMessageQueue(new DummyRandomSource(1234));
		NullBasePeerNode receiverNode = new NullBasePeerNode();
		NewPacketFormat receiver = new NewPacketFormat(receiverNode, 0, 0);
		SessionKey senderKey = new SessionKey(null, null, null, null, null, null, null, null, new NewPacketFormatKeyContext(0, 0), 1);
		SessionKey receiverKey = new SessionKey(null, null, null, null, null, null, null, null, new NewPacketFormatKeyContext(0, 0), 1);

		senderQueue.queueAndEstimateSize(new MessageItem(new byte[128], null, false, null, (short) 0), 1024);

		Thread.sleep(PacketSender.MAX_COALESCING_DELAY*2);
		NPFPacket packet1 = sender.createPacket(512, senderQueue, senderKey, false);
		assertEquals(1, receiver.handleDecryptedPacket(packet1, receiverKey).size());

		//Receiving the same packet twice should work
		assertEquals(0, receiver.handleDecryptedPacket(packet1, receiverKey).size());

		//Same message, new sequence number ie. resend
		assertEquals(0, receiver.handleDecryptedPacket(packet1, receiverKey).size());
	}

	/* This checks the output of the sequence number encryption function to
	 * make sure it doesn't change accidentally. */
	@Test
	public void testSequenceNumberEncryption() {
		BlockCipher ivCipher = new Rijndael();
		ivCipher.initialize(new byte[] {
				0x00, 0x00, 0x00, 0x00,
				0x00, 0x00, 0x00, 0x00,
				0x00, 0x00, 0x00, 0x00,
				0x00, 0x00, 0x00, 0x00
		});

		byte[] ivNonce = new byte[16];

		BlockCipher incommingCipher = new Rijndael();
		incommingCipher.initialize(new byte[] {
				0x00, 0x00, 0x00, 0x00,
				0x00, 0x00, 0x00, 0x00,
				0x00, 0x00, 0x00, 0x00,
				0x00, 0x00, 0x00, 0x00
		});

		SessionKey sessionKey = new SessionKey(null, null, null, incommingCipher, null, ivCipher, ivNonce, null, null, -1);

		byte[] encrypted = NewPacketFormat.encryptSequenceNumber(0, sessionKey);

		/* This result has not been checked, but it was the output when
		 * this test was added and we are (in this test) only
		 * interested in making sure the output doesn't change. */
		byte[] correct = new byte[] {(byte) 0xF7, (byte) 0x95, (byte) 0xBD, (byte) 0x4A};

		assertTrue(Arrays.equals(correct, encrypted));
	}

	@Test
	public void testEncryption()
			throws BlockedTooLongException, UnknownHostException, InterruptedException {
		Random random = new Random(120116);
		NullBasePeerNode senderNode = new NullBasePeerNode();
		NullBasePeerNode receiverNode = new NullBasePeerNode();
		byte[] outgoingKey = new byte[32];
		random.nextBytes(outgoingKey);
		BlockCipher outgoingCipher = new Rijndael();
		outgoingCipher.initialize(outgoingKey);
		byte[] incomingKey = new byte[32];
		random.nextBytes(incomingKey);
		BlockCipher incomingCipher = new Rijndael();
		incomingCipher.initialize(incomingKey);
		BlockCipher ivCipher = new Rijndael();
		byte[] ivKey = new byte[32];
		random.nextBytes(ivKey);
		ivCipher.initialize(ivKey);
		byte[] ivNonce = new byte[16];
		random.nextBytes(ivNonce);
		byte[] hmacKey = new byte[32];
		random.nextBytes(hmacKey);
		int senderStartSeq = 1000;
		int receiverStartSeq = 2000;

		NewPacketFormatKeyContext senderContext =
				new NewPacketFormatKeyContext(senderStartSeq, receiverStartSeq);

		NewPacketFormatKeyContext receiverContext =
				new NewPacketFormatKeyContext(receiverStartSeq, senderStartSeq);

		SessionKey senderSessionKey = new SessionKey(null, outgoingCipher, outgoingKey,
							     incomingCipher, incomingKey, ivCipher,
							     ivNonce, hmacKey, senderContext, 0);

		SessionKey receiverSessionKey = new SessionKey(null, incomingCipher, incomingKey,
							       outgoingCipher, outgoingKey,
							       ivCipher, ivNonce, hmacKey,
							       receiverContext, 0);

		senderNode.currentKey = senderSessionKey;
		receiverNode.currentKey = receiverSessionKey;

		NewPacketFormat
				senderNPF =
				new NewPacketFormat(senderNode, senderStartSeq, receiverStartSeq);
		NewPacketFormat
				receiverNPF =
				new NewPacketFormat(receiverNode, receiverStartSeq, senderStartSeq);

		PeerMessageQueue senderQueue = new PeerMessageQueue(new DummyRandomSource(1234));

		byte[] message = new byte[1024];
		random.nextBytes(message);
		byte[] copyOfMessage = Arrays.copyOf(message, message.length);

		senderQueue.queueAndEstimateSize(
				new MessageItem(message, null, false, null, (short) 0), 1024);

		senderNode.messageQueue = senderQueue;
		Thread.sleep(PacketSender.MAX_COALESCING_DELAY * 2);
		senderNPF.maybeSendPacket(false, senderSessionKey);

		FreenetInetAddress LOCALHOST = new FreenetInetAddress("127.0.0.1", true);
		Peer PEER = new Peer(LOCALHOST, 1234);

		byte[] data = senderNode.sentEncryptedPacket;

		receiverNode.decryptedMessages = new ArrayList<byte[]>();
		receiverNPF.handleReceivedPacket(data, 0, data.length, System.currentTimeMillis(),
						 PEER);

		assertEquals(1, receiverNode.decryptedMessages.size());
		assertTrue(Arrays.equals(message, copyOfMessage));
		assertTrue(Arrays.equals(message, receiverNode.decryptedMessages.get(0)));
	}
}
