Refactor UDUP module

1 files changed, 16 insertions, 14 deletions

diff --git a/src/main/java/pl/edu/mimuw/cloudatlas/agent/modules/UDUPServer.java b/src/main/java/pl/edu/mimuw/cloudatlas/agent/modules/UDUPServer.java
index f89f986..9a64bbd 100644
--- a/src/main/java/pl/edu/mimuw/cloudatlas/agent/modules/UDUPServer.java
+++ b/src/main/java/pl/edu/mimuw/cloudatlas/agent/modules/UDUPServer.java
@@ -3,41 +3,43 @@ package pl.edu.mimuw.cloudatlas.agent.modules;
 import com.google.common.primitives.Bytes;
 import pl.edu.mimuw.cloudatlas.agent.messages.UDUPMessage;
 
-import java.io.ByteArrayInputStream;
 import java.io.IOException;
 import java.net.*;
 import java.util.HashMap;
 
 public class UDUPServer {
-    UDUP udp;
+    private UDUP udp;
+    private UDUPSerializer serializer;
     private DatagramSocket socket;
     private InetAddress address;
-    private byte[] buf;
     private HashMap<InetAddress, byte[]> partialPackets;
+    private int bufSize;
 
-    public UDUPServer(UDUP udp, InetAddress addr, int port, int bufSize) throws SocketException, UnknownHostException {
+    public UDUPServer(UDUP udp, InetAddress addr, int port, int bufSize) throws SocketException {
         this.udp = udp;
         this.socket = new DatagramSocket(port, addr);
         this.address = addr;
-        this.buf = new byte[bufSize];
+        this.bufSize = bufSize;
         this.partialPackets = new HashMap<>();
+        this.serializer = new UDUPSerializer();
     }
 
     public void acceptMessage() throws IOException, InterruptedException {
+        byte[] buf = new byte[bufSize];
         DatagramPacket packet = new DatagramPacket(buf, buf.length);
         this.socket.receive(packet);
-        System.out.println("UDP received packet: " + packet.getData());
+        System.out.println("UDP " + this.address + " received packet from " + packet.getAddress());
 
         if (packet.getOffset() == 0) {
-            UDUPMessage msg = this.udp.deserialize(new ByteArrayInputStream(packet.getData()));
-            // TODO check if not full packet anyway
-            // TODO check for errors if it's not the end od transmission
-
-            this.udp.addConversation(msg);
+            UDUPMessage msg = this.serializer.deserialize(packet.getData());
             System.out.println("UDP received message " + msg.getContent().getMessageId());
 
-            if (msg.getDestinationModule() != ModuleType.UDP) {
-                this.udp.sendMessage(msg.getContent());
+            if (packet.getLength() == this.bufSize) {
+                this.addPartialMessageAndCheckSerialization(packet.getAddress(), packet.getData());
+            } else  {
+                if (msg.getDestinationModule() != ModuleType.UDP) {
+                    this.udp.sendMessage(msg.getContent());
+                }
             }
         } else {
             this.addPartialMessageAndCheckSerialization(packet.getAddress(), packet.getData());
@@ -49,7 +51,7 @@ public class UDUPServer {
             byte[] previousPacketData = this.partialPackets.get(senderAddress);
             byte[] allPacketData = Bytes.concat(previousPacketData, packetData);
             try {
-                UDUPMessage msg = this.udp.deserialize(new ByteArrayInputStream(allPacketData));
+                UDUPMessage msg = this.serializer.deserialize(allPacketData);
                 this.udp.sendMessage(msg.getContent());
                 this.partialPackets.remove(senderAddress);
             } catch (Error | Exception e) {