Miniserv Server should be ready to rock

src/main/java/net/montoyo/wd/miniserv/AbstractClient.java

@@ -9,6 +9,7 @@ import net.montoyo.wd.utilities.Log;
 import java.io.ByteArrayInputStream;
 import java.io.DataInputStream;
 import java.io.IOException;
+import java.io.UnsupportedEncodingException;
 import java.lang.reflect.InvocationTargetException;
 import java.lang.reflect.Method;
 import java.nio.ByteBuffer;
@@ -132,4 +133,20 @@ public abstract class AbstractClient {
         }
     }
 
+    protected static String readString(DataInputStream dis) throws IOException {
+        int len = dis.readShort() & 0xFFFF;
+        byte[] str = new byte[len];
+        dis.readFully(str);
+
+        try {
+            return new String(str, "UTF-8");
+        } catch(UnsupportedEncodingException ex) {
+            throw new RuntimeException(ex); //Shouldn't happen
+        }
+    }
+
+    protected byte[] getCurrentPacketRawData() {
+        return packetReader.getPacketData();
+    }
+
 }

src/main/java/net/montoyo/wd/miniserv/PacketID.java

@@ -6,10 +6,12 @@ package net.montoyo.wd.miniserv;
 
 public enum PacketID {
 
+    INIT_CONN,          //C->S
     AUTHENTICATE,       //C->S and S->C
     PING,               //C->S and S->C
     BEGIN_FILE_UPLOAD,  //C->S
     FILE_PART,          //C->S and S->C
+    FILE_STATUS,        //S->C
     GET_FILE;           //C->S
 
     public static PacketID fromInt(int i) {

src/main/java/net/montoyo/wd/miniserv/server/ClientManager.java

@@ -24,7 +24,7 @@ public class ClientManager {
     private final HashMap<UUID, byte[]> keys = new HashMap<>();
     private final ReentrantReadWriteLock keyLock = new ReentrantReadWriteLock();
 
-    public byte[] getClientKey(UUID uuid) {
+    public byte[] getOrGenClientKey(UUID uuid) {
         keyLock.readLock().lock();
         byte[] key = keys.get(uuid);
         keyLock.readLock().unlock();
@@ -41,6 +41,10 @@ public class ClientManager {
         return key;
     }
 
+    public byte[] getClientKey(UUID uuid) {
+        return keys.get(uuid);
+    }
+
     public void revokeClientKey(UUID id) {
         keyLock.writeLock().lock();
         keys.remove(id);
@@ -84,7 +88,7 @@ public class ClientManager {
             Cipher cipher = Cipher.getInstance(KeyParameters.RSA_CIPHER);
 
             cipher.init(Cipher.ENCRYPT_MODE, key, random);
-            return cipher.doFinal(getClientKey(client));
+            return cipher.doFinal(getOrGenClientKey(client));
         } catch(NoSuchAlgorithmException | NoSuchPaddingException ex) {
             Log.warningEx("%s is not supported?!?!", ex, KeyParameters.RSA_CIPHER);
         } catch(InvalidKeySpecException | InvalidKeyException ex) {

src/main/java/net/montoyo/wd/miniserv/server/Server.java

@@ -7,6 +7,7 @@ package net.montoyo.wd.miniserv.server;
 import net.montoyo.wd.utilities.Log;
 import net.montoyo.wd.utilities.Util;
 
+import java.io.File;
 import java.net.InetSocketAddress;
 import java.nio.ByteBuffer;
 import java.nio.channels.SelectionKey;
@@ -18,12 +19,24 @@ import java.util.HashMap;
 
 public class Server extends Thread {
 
+    private static Server instance;
+
+    public static Server getInstance() {
+        if(instance == null)
+            instance = new Server();
+
+        return instance;
+    }
+
     private ServerSocketChannel server;
     private Selector selector;
     private int port = 25566;
     private final ArrayList<ServerClient> clientList = new ArrayList<>();
     private final HashMap<SocketChannel, ServerClient> clientMap = new HashMap<>();
     private final ByteBuffer readBuffer = ByteBuffer.allocateDirect(8192);
+    private final ClientManager clientMgr = new ClientManager();
+    private File directory;
+    private long maxQuota = 1024 * 1024; //1 MiB max
 
     public Server() {
         setDaemon(true);
@@ -81,6 +94,7 @@ public class Server extends Thread {
                     ServerClient toAdd = new ServerClient(chan, selector);
                     clientMap.put(chan, toAdd);
                     clientList.add(toAdd);
+                    toAdd.onConnect();
                 }
             }
 
@@ -138,4 +152,25 @@ public class Server extends Thread {
         Util.silentClose(cli.getChannel());
     }
 
+    public ClientManager getClientManager() {
+        return clientMgr;
+    }
+
+    public void setDirectory(File dir) {
+        if(!dir.exists()) {
+            if(!dir.mkdir())
+                Log.warning("Could not create miniserv storage directory %s", dir.getAbsolutePath());
+        }
+
+        directory = dir;
+    }
+
+    public File getDirectory() {
+        return directory;
+    }
+
+    public long getMaxQuota() {
+        return maxQuota;
+    }
+
 }

src/main/java/net/montoyo/wd/miniserv/server/ServerClient.java

@@ -5,16 +5,29 @@
 package net.montoyo.wd.miniserv.server;
 
 import net.montoyo.wd.miniserv.*;
+import net.montoyo.wd.utilities.Log;
+import net.montoyo.wd.utilities.Util;
 
-import java.io.DataInputStream;
+import java.io.*;
-import java.io.IOException;
 import java.nio.channels.Selector;
 import java.nio.channels.SocketChannel;
+import java.util.UUID;
+import java.util.function.Consumer;
 
 public class ServerClient extends AbstractClient {
 
+    private static final byte[] FILE_UPLOAD_BUFFER = new byte[65536];
+
     private boolean remove;
     private boolean isAuthenticated;
+    private UUID uuid;
+    private byte[] challenge;
+    private File userDir;
+    private long quota;
+    private FileOutputStream currentFile;
+    private long currentFileSize;
+    private long currentFileExpectedSize;
+    private boolean sendingFile; //!= receiving, which is handled by currentFile
 
     ServerClient(SocketChannel s, Selector ss) {
         socket = s;
@@ -26,6 +39,9 @@ public class ServerClient extends AbstractClient {
         remove = true;
     }
 
+    public void onConnect() {
+    }
+
     void setShouldRemove() {
         remove = true;
     }
@@ -38,16 +54,248 @@ public class ServerClient extends AbstractClient {
         return socket;
     }
 
+    @PacketHandler(PacketID.INIT_CONN)
+    public void handleInitConnPacket(DataInputStream dis) throws IOException {
+        if(uuid != null) {
+            Log.warning("A client tried to change his UUID?");
+            return;
+        }
+
+        long msb = dis.readLong();
+        long lsb = dis.readLong();
+        UUID uuid = new UUID(msb, lsb);
+        byte[] key = Server.getInstance().getClientManager().getClientKey(uuid);
+
+        if(key == null) {
+            Log.warning("Unkown client with UUID %s wanted to connect", uuid.toString());
+            remove = true;
+        } else {
+            this.uuid = uuid;
+            challenge = Server.getInstance().getClientManager().generateChallenge();
+
+            OutgoingPacket pkt = new OutgoingPacket();
+            pkt.writeByte(PacketID.AUTHENTICATE.ordinal());
+            pkt.writeByte(challenge.length);
+            pkt.writeBytes(challenge);
+
+            sendPacket(pkt);
+        }
+    }
+
     @PacketHandler(PacketID.AUTHENTICATE)
     public void handleAuthPacket(DataInputStream dis) throws IOException {
-        //TODO: Do some stuff
+        if(uuid == null) {
+            Log.warning("A client tried to authenticate, but he didn't send the connection packet");
+            remove = true;
+            return;
+        }
+
+        int len = dis.readByte() & 0xFF;
+        byte[] mac = new byte[len];
+
+        if(Server.getInstance().getClientManager().verifyClient(uuid, challenge, mac)) {
+            Log.info("Client with UUID %s authenticated successfully", uuid.toString());
+
+            userDir = new File(Server.getInstance().getDirectory(), uuid.toString());
+            if(!userDir.exists() && !userDir.mkdir())
+                Log.warning("Could not create storage directory for user %s, things may go wrong...", uuid.toString());
+
+            try {
+                DataInputStream quotaDis = new DataInputStream(new FileInputStream(new File(userDir, ".quota")));
+                quota = quotaDis.readLong();
+                Util.silentClose(quotaDis);
+            } catch(FileNotFoundException ex) {
+                quota = 0;
+            } catch(IOException ex) {
+                Log.warningEx("Couldn't read quota for user %s, things may go wrong...", ex, uuid.toString());
+                quota = Server.getInstance().getMaxQuota();
+            }
+
+            isAuthenticated = true;
+        } else {
+            Log.warning("Client with UUID %s failed to authenticate", uuid.toString());
+            remove = true;
+        }
     }
 
     @PacketHandler(PacketID.PING)
     public void handlePing(DataInputStream dis) {
-        OutgoingPacket pkt = new OutgoingPacket();
+        if(isAuthenticated) {
-        pkt.writeByte(PacketID.PING.ordinal());
+            OutgoingPacket pkt = new OutgoingPacket();
-        sendPacket(pkt);
+            pkt.writeByte(PacketID.PING.ordinal());
+            sendPacket(pkt);
+        }
+    }
+
+    @PacketHandler(PacketID.BEGIN_FILE_UPLOAD)
+    public void handleBeginUpload(DataInputStream dis) throws IOException {
+        if(isAuthenticated) {
+            String fname = readString(dis);
+            long size = dis.readLong();
+
+            OutgoingPacket rep = new OutgoingPacket();
+            rep.writeByte(PacketID.BEGIN_FILE_UPLOAD.ordinal());
+
+            if(isFileNameInvalid(fname)) {
+                Log.warning("Client %s tried to upload a file with a bad name", uuid.toString());
+                rep.writeByte(1);
+            } else if(size <= 0) {
+                Log.warning("Client %s tried to upload a file an invalid size", uuid.toString());
+                rep.writeByte(2);
+            } else if(quota + size > Server.getInstance().getMaxQuota())
+                rep.writeByte(3);
+            else if(currentFile != null || sendingFile)
+                rep.writeByte(4);
+            else {
+                File fle = new File(userDir, fname);
+
+                if(fle.exists())
+                    rep.writeByte(5);
+                else {
+                    try {
+                        currentFile = new FileOutputStream(fle);
+                        currentFileSize = 0;
+                        currentFileExpectedSize = size;
+
+                        rep.writeByte(0); //OK
+                    } catch(IOException ex) {
+                        Log.warningEx("IOException while uploading file %s from user %s", ex, fname, uuid.toString());
+                        rep.writeByte(6);
+                    }
+                }
+            }
+
+            sendPacket(rep);
+        }
+    }
+
+    @PacketHandler(PacketID.FILE_PART)
+    public void handleFilePart(DataInputStream dis) throws IOException {
+        if(isAuthenticated && currentFile != null) {
+            int len = dis.readShort() & 0xFFFF;
+            if(len <= 0) {
+                //Aborted by user
+                finishUpload(1);
+                return;
+            }
+
+            currentFileSize += (long) len;
+            if(currentFileSize > currentFileExpectedSize) {
+                //Exceeded expected size
+                finishUpload(2);
+                return;
+            }
+
+            try {
+                currentFile.write(getCurrentPacketRawData(), 3, len);
+            } catch(IOException ex) {
+                Log.warningEx("Client %s encountered an IOException while uploading some file", ex, uuid.toString());
+                finishUpload(3);
+                currentFileSize -= (long) len;
+                return;
+            }
+
+            if(currentFileSize >= currentFileExpectedSize)
+                finishUpload(0); //No error
+        }
+    }
+
+    @PacketHandler(PacketID.GET_FILE)
+    public void handleGetFile(DataInputStream dis) throws IOException {
+        if(isAuthenticated && currentFile == null) {
+            long msb = dis.readLong();
+            long lsb = dis.readLong();
+            String fname = readString(dis);
+
+            OutgoingPacket rep = new OutgoingPacket();
+            rep.writeByte(PacketID.GET_FILE.ordinal());
+
+            if(isFileNameInvalid(fname))
+                rep.writeByte(1);
+            else {
+                UUID user = new UUID(msb, lsb);
+                File fle = new File(Server.getInstance().getDirectory(), user.toString() + File.separatorChar + fname);
+
+                try {
+                    rep.setOnFinishAction(new SendFileCallback(fle));
+                    sendingFile = true;
+                } catch(FileNotFoundException ex) {
+                    rep.writeByte(2);
+                }
+            }
+
+            sendPacket(rep);
+        }
+    }
+
+    private static boolean isFileNameInvalid(String fname) {
+        return fname.isEmpty() || fname.length() > 64 || fname.charAt(0) == '.' || fname.indexOf('/') >= 0 || fname.indexOf('\\') >= 0;
+    }
+
+    private void finishUpload(int status) {
+        if(currentFile != null) {
+            OutgoingPacket pkt = new OutgoingPacket();
+            pkt.writeByte(PacketID.FILE_STATUS.ordinal());
+            pkt.writeByte(status);
+            sendPacket(pkt);
+
+            Util.silentClose(currentFile);
+            currentFile = null;
+
+            quota += currentFileSize;
+
+            try {
+                DataOutputStream dos = new DataOutputStream(new FileOutputStream(new File(userDir, ".quota")));
+                dos.writeLong(quota);
+                Util.silentClose(dos);
+            } catch(IOException ex) {
+                Log.errorEx("Could not save quota data for user %s", ex, uuid.toString());
+            }
+        }
+    }
+
+    private class SendFileCallback implements Consumer<OutgoingPacket> {
+
+        private final FileInputStream fis;
+
+        private SendFileCallback(File fle) throws FileNotFoundException {
+            fis = new FileInputStream(fle);
+        }
+
+        @Override
+        public void accept(OutgoingPacket nocare) {
+            int rd;
+
+            do {
+                try {
+                    rd = fis.read(FILE_UPLOAD_BUFFER);
+                } catch(IOException ex) {
+                    Log.warningEx("Caught IOException while sending some file", ex);
+
+                    OutgoingPacket pkt = new OutgoingPacket();
+                    pkt.writeByte(PacketID.GET_FILE.ordinal());
+                    pkt.writeByte(3); //Read error
+                    sendPacket(pkt);
+
+                    Util.silentClose(fis);
+                    sendingFile = false;
+                    return;
+                }
+            } while(rd == 0);
+
+            OutgoingPacket pkt = new OutgoingPacket();
+            if(rd < 0) {
+                rd = 0; //EOF
+                sendingFile = false;
+            } else
+                pkt.setOnFinishAction(this);
+
+            pkt.writeByte(PacketID.FILE_PART.ordinal());
+            pkt.writeShort(rd);
+            pkt.writeBytes(FILE_UPLOAD_BUFFER);
+            sendPacket(pkt);
+        }
+
     }
 
 }