TFTPClient.java

import java.io.*;
import java.net.*;
import java.nio.*;
import java.nio.file.*;
import java.util.*;
import java.util.concurrent.*;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicInteger;

public class TFTPClient {
    private static final int PORT = 1069;
    private static final int BLOCK_SIZE = 512;
    private static final int TIMEOUT = 5000;
    private static final int MAX_RETRIES = 5;
    private static final int MAX_THREADS = 4; // Limit number of parallel threads
    private static final String CLIENT_FILES_DIR = "clientfiles";

    public static void main(String[] args) {
        if (args.length < 3) {
            System.out.println("Usage: java TFTPClient <put|get> <filename> <serverAddress> [username] [password]");
            return;
        }

        String action = args[0];
        String filename = args[1];
        String serverAddress = args[2];
        String username = null;
        String password = null;

        // Get username and password from command line or prompt
        if (args.length >= 5) {
            username = args[3];
            password = args[4];
        } else {
            Scanner scanner = new Scanner(System.in);
            System.out.print("Enter username: ");
            username = scanner.nextLine();
            System.out.print("Enter password: ");
            password = scanner.nextLine();
        }

        ExecutorService executor = Executors.newFixedThreadPool(MAX_THREADS);
        try (DatagramSocket socket = new DatagramSocket()) {
            InetAddress server = InetAddress.getByName(serverAddress);
            socket.setSoTimeout(5000);

            if (action.equals("get")) {
                sendReadRequest(socket, server, filename, username, password);
            } else if (action.equals("put")) {
                sendWriteRequest(socket, server, filename, username, password, executor);
            } else {
                System.out.println("Unknown action: " + action);
            }
        } catch (IOException e) {
            e.printStackTrace();
        } finally {
            executor.shutdown();
        }
    }

    private static void sendReadRequest(DatagramSocket socket, InetAddress server, String filename, String username, String password) throws IOException {
        // Create RRQ packet with username and password
        ByteBuffer buffer = ByteBuffer.allocate(512);
        buffer.putShort((short) 1); // RRQ opcode
        
        // Add username and password
        buffer.put(username.getBytes());
        buffer.put((byte) 0);
        buffer.put(password.getBytes());
        buffer.put((byte) 0);
        
        // Add filename and mode
        buffer.put(filename.getBytes());
        buffer.put((byte) 0);
        buffer.put("octet".getBytes());
        buffer.put((byte) 0);

        DatagramPacket packet = new DatagramPacket(buffer.array(), buffer.position(), server, PORT);
        socket.send(packet);
        System.out.println("Sent RRQ for file: " + filename + " with username: " + username);

        // Start receiving data blocks
        receiveFileData(socket, filename);
    }

    private static void sendWriteRequest(DatagramSocket socket, InetAddress server, String filename, 
            String username, String password, ExecutorService executor) throws IOException {
        // Create WRQ packet with username and password
        ByteBuffer buffer = ByteBuffer.allocate(512);
        buffer.putShort((short) 2); // WRQ opcode
        
        // Add username and password
        buffer.put(username.getBytes());
        buffer.put((byte) 0);
        buffer.put(password.getBytes());
        buffer.put((byte) 0);
        
        // Add filename and mode
        buffer.put(filename.getBytes());
        buffer.put((byte) 0);
        buffer.put("octet".getBytes());
        buffer.put((byte) 0);

        DatagramPacket packet = new DatagramPacket(buffer.array(), buffer.position(), server, PORT);
        socket.send(packet);
        System.out.println("Sent WRQ for file: " + filename + " with username: " + username);

        // Wait for initial ACK (block 0)
        boolean ackReceived = false;
        int retries = 0;
        while (!ackReceived && retries < MAX_RETRIES) {
            try {
                DatagramPacket ackPacket = new DatagramPacket(new byte[4], 4);
                socket.receive(ackPacket);
                
                ByteBuffer ackBuffer = ByteBuffer.wrap(ackPacket.getData());
                short opcode = ackBuffer.getShort();
                short blockNum = ackBuffer.getShort();
                
                if (opcode == 4 && blockNum == 0) {
                    ackReceived = true;
                    System.out.println("Received initial ACK, starting parallel transmission");
                }
            } catch (SocketTimeoutException e) {
                retries++;
                System.out.println("Timeout waiting for initial ACK, retrying...");
                socket.send(packet);
            }
        }

        if (!ackReceived) {
            System.out.println("Failed to receive initial ACK after " + MAX_RETRIES + " retries");
            return;
        }

        // Start sending data blocks in parallel
        sendFileData(socket, server, filename, executor);
    }

    private static void sendFileData(DatagramSocket socket, InetAddress server, String filename, ExecutorService executor) {
        try {
            // Read file from clientfiles directory
            Path filePath = Paths.get(CLIENT_FILES_DIR, filename);
            byte[] fileData = Files.readAllBytes(filePath);
            int totalBlocks = (int) Math.ceil((double) fileData.length / BLOCK_SIZE);
            AtomicInteger lastBlockSize = new AtomicInteger(fileData.length % BLOCK_SIZE);
            if (lastBlockSize.get() == 0) lastBlockSize.set(BLOCK_SIZE);

            System.out.println("File size: " + fileData.length + " bytes");
            System.out.println("Total blocks: " + totalBlocks);
            System.out.println("Last block size: " + lastBlockSize.get() + " bytes");

            List<Future<Void>> futures = new ArrayList<>();
            CountDownLatch completionLatch = new CountDownLatch(totalBlocks);
            AtomicBoolean transferComplete = new AtomicBoolean(false);

            // Submit tasks to send each block in parallel
            for (int block = 1; block <= totalBlocks; block++) {
                final int blockNum = block;
                final int start = (block - 1) * BLOCK_SIZE;
                final int end = Math.min(start + BLOCK_SIZE, fileData.length);
                final byte[] blockData = Arrays.copyOfRange(fileData, start, end);
                final boolean isLastBlock = (block == totalBlocks);

                System.out.println("Preparing block " + blockNum + ": " + start + " to " + end + 
                    " (size: " + blockData.length + " bytes)");

                futures.add(executor.submit(() -> {
                    String threadName = Thread.currentThread().getName();
                    try {
                        sendDataBlock(socket, server, blockNum, blockData, threadName, isLastBlock, lastBlockSize.get());
                        completionLatch.countDown();
                    } catch (IOException e) {
                        e.printStackTrace();
                    }
                    return null;
                }));
            }

            // Wait for all blocks to be sent or timeout after 30 seconds
            if (!completionLatch.await(30, TimeUnit.SECONDS)) {
                System.out.println("Timeout waiting for all blocks to be sent");
            } else {
                System.out.println("All blocks sent successfully");
            }

            // Shutdown the executor and wait for all threads to finish
            executor.shutdown();
            if (!executor.awaitTermination(5, TimeUnit.SECONDS)) {
                executor.shutdownNow();
            }

        } catch (IOException | InterruptedException e) {
            e.printStackTrace();
        }
    }

    private static void sendDataBlock(DatagramSocket socket, InetAddress server, int blockNum, byte[] blockData, 
            String threadName, boolean isLastBlock, int lastBlockSize) throws IOException {
        ByteBuffer dataBuffer = ByteBuffer.allocate(4 + blockData.length);
        dataBuffer.putShort((short) 3); // DATA opcode
        dataBuffer.putShort((short) blockNum);
        dataBuffer.put(blockData);

        DatagramPacket dataPacket = new DatagramPacket(dataBuffer.array(), dataBuffer.array().length, server, PORT);
        boolean ackReceived = false;
        int retries = 0;

        while (!ackReceived && retries < MAX_RETRIES) {
            try {
                socket.send(dataPacket);
                System.out.println("Thread " + threadName + " sent block " + blockNum + 
                    (isLastBlock ? " (last block)" : "") + " (size: " + blockData.length + " bytes)");

                DatagramPacket ackPacket = new DatagramPacket(new byte[4], 4);
                socket.receive(ackPacket);

                ByteBuffer ackBuffer = ByteBuffer.wrap(ackPacket.getData());
                short opcode = ackBuffer.getShort();
                short receivedBlockNum = ackBuffer.getShort();

                if (opcode == 4 && receivedBlockNum == blockNum) {
                    System.out.println("Thread " + threadName + " received ACK for block " + blockNum);
                    ackReceived = true;
                }
            } catch (SocketTimeoutException e) {
                retries++;
                System.out.println("Thread " + threadName + " timeout for block " + blockNum + ", retry " + retries);
            }
        }

        if (!ackReceived) {
            System.out.println("Thread " + threadName + " failed to receive ACK for block " + blockNum + 
                " after " + MAX_RETRIES + " retries");
        }
    }

    private static void receiveFileData(DatagramSocket socket, String filename) throws IOException {
        FileOutputStream fileOutputStream = new FileOutputStream(new File(CLIENT_FILES_DIR, filename));
        int blockNum = 1;
        boolean lastBlock = false;

        while (!lastBlock) {
            DatagramPacket dataPacket = new DatagramPacket(new byte[516], 516);
            socket.receive(dataPacket);

            ByteBuffer buffer = ByteBuffer.wrap(dataPacket.getData(), 0, dataPacket.getLength());
            short opcode = buffer.getShort();
            short receivedBlockNum = buffer.getShort();
            byte[] data = new byte[dataPacket.getLength() - 4];
            buffer.get(data);

            if (opcode == 3 && receivedBlockNum == blockNum) {
                // Write received data to the file
                fileOutputStream.write(data);
                System.out.println("Received block " + blockNum + " (size: " + data.length + " bytes)");

                // Send ACK
                ByteBuffer ackBuffer = ByteBuffer.allocate(4);
                ackBuffer.putShort((short) 4); // ACK opcode
                ackBuffer.putShort(receivedBlockNum);
                DatagramPacket ackPacket = new DatagramPacket(ackBuffer.array(), ackBuffer.array().length,
                    dataPacket.getAddress(), dataPacket.getPort());
                socket.send(ackPacket);
                System.out.println("Sent ACK for block " + blockNum);

                // If the received block is less than BLOCK_SIZE, it's the last block
                if (data.length < BLOCK_SIZE) {
                    lastBlock = true;
                    System.out.println("Received last block");
                }

                blockNum++;
            }
        }

        fileOutputStream.close();
        System.out.println("File transfer complete");
    }
}