TFTPServer.java

import java.io.*;
import java.net.*;
import java.nio.*;
import java.nio.file.*;
import java.util.*;
import java.util.concurrent.*;
import java.text.SimpleDateFormat;
import java.util.Date;

public class TFTPServer {
    private static final int PORT = 1069;
    private static final int BLOCK_SIZE = 512;
    private static final int MAX_RETRIES = 5;
    private static final String SERVER_FILES_DIR = "serverfiles";
    private static final Map<String, FileTransferState> activeTransfers = new ConcurrentHashMap<>();
    private static final Map<String, String> userCredentials = new HashMap<>();
    private static final int MAX_THREADS = 10;

    static {
        // Initialize with some test credentials
        userCredentials.put("admin", "password123");
        userCredentials.put("user", "pass123");
    }

    private static boolean authenticate(String username, String password) {
        String storedPassword = userCredentials.get(username);
        return storedPassword != null && storedPassword.equals(password);
    }

    private static class FileTransferState {
        FileOutputStream fileOutputStream;
        Set<Integer> receivedBlocks;
        int expectedBlock;
        String filename;
        InetAddress clientAddress;
        int clientPort;
        int totalBlocks;

        public FileTransferState(String filename, InetAddress clientAddress, int clientPort) throws FileNotFoundException {
            this.filename = filename;
            this.clientAddress = clientAddress;
            this.clientPort = clientPort;
            // Create serverfiles directory if it doesn't exist
            File serverDir = new File(SERVER_FILES_DIR);
            if (!serverDir.exists()) {
                serverDir.mkdirs();
            }
            this.fileOutputStream = new FileOutputStream(new File(serverDir, filename));
            this.receivedBlocks = new HashSet<>();
            this.expectedBlock = 1;
            this.totalBlocks = 0;
        }
    }

    public static void main(String[] args) {
        try (DatagramSocket socket = new DatagramSocket(PORT)) {
            System.out.println("TFTP Server running on port " + PORT);

            while (true) {
                DatagramPacket packet = new DatagramPacket(new byte[516], 516);
                socket.receive(packet);
                handleRequest(socket, packet);
            }
        } catch (Exception e) {
            e.printStackTrace();
        }
    }

    private static void handleRequest(DatagramSocket socket, DatagramPacket packet) throws Exception {
        ByteBuffer buffer = ByteBuffer.wrap(packet.getData(), 0, packet.getLength());
        short opcode = buffer.getShort();

        if (opcode == 3 || opcode == 4) {
            // Handle DATA and ACK packets directly without authentication
            if (opcode == 3) {
                handleDataPacket(socket, packet);
            } else {
                handleAckPacket(socket, packet);
            }
            return;
        }

        // For RRQ and WRQ, extract and verify credentials
        String username = extractString(buffer);
        String password = extractString(buffer);

        if (!authenticate(username, password)) {
            System.out.println("Authentication failed for user: " + username);
            return;
        }

        System.out.println("Authentication successful for user: " + username);

        // Extract filename after authentication
        String filename = extractString(buffer);
        String mode = extractString(buffer); // Skip mode

        if (opcode == 1) { // RRQ
            handleReadRequest(socket, packet, filename, username);
        } else if (opcode == 2) { // WRQ
            handleWriteRequest(socket, packet, filename, username);
        }
    }

    private static String extractString(ByteBuffer buffer) {
        StringBuilder sb = new StringBuilder();
        byte b;
        while ((b = buffer.get()) != 0) {
            sb.append((char) b);
        }
        return sb.toString();
    }

    private static void handleReadRequest(DatagramSocket socket, DatagramPacket packet, String filename, String username) throws Exception {
        File file = new File(SERVER_FILES_DIR, filename);
        
        if (!file.exists()) {
            System.out.println("Thread " + Thread.currentThread().getName() + 
                " - File not found: " + filename);
            return;
        }

        System.out.println("Thread " + Thread.currentThread().getName() + 
            " received RRQ for file: " + filename + " from user: " + username);
        
        // Start sending file in blocks
        sendFileBlocks(socket, file, packet.getAddress(), packet.getPort());
    }

    private static void sendFileBlocks(DatagramSocket socket, File file, InetAddress clientAddress, int clientPort) throws Exception {
        try (FileInputStream fileInputStream = new FileInputStream(file)) {
            byte[] fileData = Files.readAllBytes(file.toPath());
            int totalBlocks = (int) Math.ceil((double) fileData.length / BLOCK_SIZE);
            
            System.out.println("\n=== Starting File Transfer ===");
            System.out.println("Thread " + Thread.currentThread().getName() + 
                " initiating transfer of " + totalBlocks + " blocks");
            System.out.println("File: " + file.getName() + " (" + fileData.length + " bytes)");
            System.out.println("=====================================\n");

            // Use a single thread to send blocks sequentially
            for (int blockNum = 1; blockNum <= totalBlocks; blockNum++) {
                int start = (blockNum - 1) * BLOCK_SIZE;
                int end = Math.min(start + BLOCK_SIZE, fileData.length);
                byte[] blockData = Arrays.copyOfRange(fileData, start, end);
                boolean isLastBlock = (blockNum == totalBlocks);

                System.out.println("Thread " + Thread.currentThread().getName() + 
                    " preparing to send block " + blockNum + 
                    (isLastBlock ? " (last block)" : "") + 
                    " (size: " + blockData.length + " bytes)");

                // Create DATA packet
                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,
                    clientAddress, clientPort);
                
                boolean ackReceived = false;
                int retries = 0;

                while (!ackReceived && retries < MAX_RETRIES) {
                    try {
                        long startTime = System.currentTimeMillis();
                        socket.send(dataPacket);
                        System.out.println("Thread " + Thread.currentThread().getName() + 
                            " sent block " + blockNum + 
                            (isLastBlock ? " (last block)" : "") + 
                            " (size: " + blockData.length + " bytes) at " + 
                            new SimpleDateFormat("HH:mm:ss.SSS").format(new Date(startTime)));

                        // Wait for ACK
                        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) {
                            long endTime = System.currentTimeMillis();
                            System.out.println("Thread " + Thread.currentThread().getName() + 
                                " received ACK for block " + blockNum + 
                                " after " + (endTime - startTime) + "ms");
                            ackReceived = true;
                        }
                    } catch (SocketTimeoutException e) {
                        retries++;
                        System.out.println("Thread " + Thread.currentThread().getName() + 
                            " timeout for block " + blockNum + 
                            ", retry " + retries + "/" + MAX_RETRIES);
                    }
                }

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

            System.out.println("\n=== Transfer Complete ===");
            System.out.println("Thread " + Thread.currentThread().getName() + 
                " completed transfer of " + totalBlocks + " blocks");
            System.out.println("Total file size: " + fileData.length + " bytes");
        }
    }

    private static void handleWriteRequest(DatagramSocket socket, DatagramPacket packet, String filename, String username) throws Exception {
        // Create transfer ID from client address and port
        String transferId = packet.getAddress().getHostAddress() + ":" + packet.getPort();
        
        System.out.println("Thread " + Thread.currentThread().getName() + 
            " received WRQ for file: " + filename + " from user: " + username);

        // Create new transfer state
        activeTransfers.put(transferId, new FileTransferState(filename, packet.getAddress(), packet.getPort()));

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

    private static void handleDataPacket(DatagramSocket socket, DatagramPacket packet) throws Exception {
        ByteBuffer buffer = ByteBuffer.wrap(packet.getData(), 0, packet.getLength());
        short opcode = buffer.getShort();
        short blockNum = buffer.getShort();
        byte[] data = new byte[packet.getLength() - 4];
        buffer.get(data);

        // Try to find the transfer state by matching the client address
        String clientAddress = packet.getAddress().getHostAddress();
        FileTransferState state = null;
        String foundTransferId = null;

        // Search for the transfer state with matching client address
        for (Map.Entry<String, FileTransferState> entry : activeTransfers.entrySet()) {
            if (entry.getValue().clientAddress.getHostAddress().equals(clientAddress)) {
                state = entry.getValue();
                foundTransferId = entry.getKey();
                break;
            }
        }

        if (state != null) {
            synchronized (state) {
                if (!state.receivedBlocks.contains((int) blockNum)) {
                    // Write the data block
                    state.fileOutputStream.write(data);
                    state.fileOutputStream.flush(); // Ensure data is written to disk
                    state.receivedBlocks.add((int) blockNum);
                    System.out.println("Written block " + blockNum + " to file " + state.filename + 
                        " (size: " + data.length + " bytes)");

                    // Send ACK to the original client port
                    ByteBuffer ackBuffer = ByteBuffer.allocate(4);
                    ackBuffer.putShort((short) 4); // ACK opcode
                    ackBuffer.putShort(blockNum);
                    DatagramPacket ackPacket = new DatagramPacket(ackBuffer.array(), ackBuffer.array().length,
                        state.clientAddress, state.clientPort);
                    socket.send(ackPacket);
                    System.out.println("Sent ACK for block " + blockNum);

                    // Check if this was the last block
                    if (data.length < BLOCK_SIZE) {
                        state.totalBlocks = blockNum;
                        System.out.println("Last block received: " + blockNum + " (size: " + data.length + " bytes)");
                    }

                    // Check if all blocks have been received
                    if (state.totalBlocks > 0 && state.receivedBlocks.size() == state.totalBlocks) {
                        state.fileOutputStream.close();
                        activeTransfers.remove(foundTransferId);
                        System.out.println("File transfer complete for " + state.filename);
                        
                        // Verify file size
                        File receivedFile = new File(SERVER_FILES_DIR, state.filename);
                        System.out.println("Final file size: " + receivedFile.length() + " bytes");
                    }
                } else {
                    // Resend ACK for duplicate block
                    ByteBuffer ackBuffer = ByteBuffer.allocate(4);
                    ackBuffer.putShort((short) 4);
                    ackBuffer.putShort(blockNum);
                    DatagramPacket ackPacket = new DatagramPacket(ackBuffer.array(), ackBuffer.array().length,
                        state.clientAddress, state.clientPort);
                    socket.send(ackPacket);
                    System.out.println("Resent ACK for duplicate block " + blockNum);
                }
            }
        } else {
            System.out.println("No active transfer found for client " + clientAddress);
        }
    }

    private static void handleAckPacket(DatagramSocket socket, DatagramPacket packet) throws Exception {
        ByteBuffer buffer = ByteBuffer.wrap(packet.getData(), 0, packet.getLength());
        short opcode = buffer.getShort();
        short blockNum = buffer.getShort();

        // Find the transfer state
        String clientAddress = packet.getAddress().getHostAddress();
        FileTransferState state = null;
        String foundTransferId = null;

        for (Map.Entry<String, FileTransferState> entry : activeTransfers.entrySet()) {
            if (entry.getValue().clientAddress.getHostAddress().equals(clientAddress)) {
                state = entry.getValue();
                foundTransferId = entry.getKey();
                break;
            }
        }

        if (state != null) {
            System.out.println("Received ACK for block " + blockNum + " from " + clientAddress);
        }
    }

    private static String extractFilename(DatagramPacket packet) {
        ByteBuffer buffer = ByteBuffer.wrap(packet.getData(), 0, packet.getLength());
        buffer.getShort(); // Skip the WRQ opcode
        StringBuilder filename = new StringBuilder();
        byte b;
        while ((b = buffer.get()) != 0) {
            filename.append((char) b);
        }
        return filename.toString();
    }
}