From d078808486b25e9fb2f90f15dd9d9a2bca8250c0 Mon Sep 17 00:00:00 2001 From: Gaa56 Date: Thu, 30 Oct 2025 19:25:27 +0000 Subject: [PATCH] Update SocketConnection --- .../java/sd/protocol/SocketConnection.java | 106 ++++++++++-------- 1 file changed, 61 insertions(+), 45 deletions(-) diff --git a/main/src/main/java/sd/protocol/SocketConnection.java b/main/src/main/java/sd/protocol/SocketConnection.java index f6392a4..446281c 100644 --- a/main/src/main/java/sd/protocol/SocketConnection.java +++ b/main/src/main/java/sd/protocol/SocketConnection.java @@ -1,15 +1,23 @@ package sd.protocol; import java.io.Closeable; +import java.io.DataInputStream; +import java.io.DataOutputStream; import java.io.IOException; -import java.io.ObjectInputStream; -import java.io.ObjectOutputStream; + +import java.io.InputStream; +import java.io.OutputStream; import java.net.ConnectException; import java.net.Socket; import java.net.SocketTimeoutException; import java.net.UnknownHostException; import java.util.concurrent.TimeUnit; +import sd.serialization.MessageSerializer; +import sd.serialization.SerializationException; +import sd.serialization.SerializerFactory; + + /** * Wrapper class that simplifies communication via Sockets. * Includes connection retry logic for robustness. @@ -17,8 +25,9 @@ import java.util.concurrent.TimeUnit; public class SocketConnection implements Closeable { private final Socket socket; - private final ObjectOutputStream outputStream; - private final ObjectInputStream inputStream; + private final OutputStream outputStream; + private final InputStream inputStream; + private final MessageSerializer serializer; // --- Configuration for Retry Logic --- /** Maximum number of connection attempts. */ @@ -88,16 +97,11 @@ public class SocketConnection implements Closeable { // If connection was successful, assign to final variable and create streams this.socket = tempSocket; - try { - // IMPORTANT: The order is crucial. OutputStream first. - this.outputStream = new ObjectOutputStream(socket.getOutputStream()); - this.inputStream = new ObjectInputStream(socket.getInputStream()); - } catch (IOException e) { - // If stream creation fails even after successful socket connection, clean up. - System.err.println("[SocketConnection] Failed to create streams after connection: " + e.getMessage()); - try { socket.close(); } catch (IOException closeEx) { /* ignore */ } - throw e; // Re-throw the stream creation exception - } + + this.outputStream = socket.getOutputStream(); + this.inputStream = socket.getInputStream(); + this.serializer = SerializerFactory.createDefault(); + } @@ -111,12 +115,10 @@ public class SocketConnection implements Closeable { */ public SocketConnection(Socket acceptedSocket) throws IOException { this.socket = acceptedSocket; - - // IMPORTANT: The order is crucial. OutputStream first. - this.outputStream = new ObjectOutputStream(socket.getOutputStream()); - this.inputStream = new ObjectInputStream(socket.getInputStream()); - System.out.printf("[SocketConnection] Connection accepted from %s:%d.%n", - acceptedSocket.getInetAddress().getHostAddress(), acceptedSocket.getPort()); + this.outputStream = socket.getOutputStream(); + this.inputStream = socket.getInputStream(); + this.serializer = SerializerFactory.createDefault(); + } /** @@ -126,52 +128,66 @@ public class SocketConnection implements Closeable { * @throws IOException If writing to the stream fails or socket is not connected. */ public void sendMessage(MessageProtocol message) throws IOException { - if (!isConnected()) { - throw new IOException("Socket is not connected."); + if (socket == null || !socket.isConnected()) { + throw new IOException("Socket is not connected"); } - synchronized (outputStream) { - outputStream.writeObject(message); - outputStream.flush(); // Ensures the message is sent immediately. + + try { + // Serializa para bytes JSON + byte[] data = serializer.serialize(message); + + // Write 4-byte length prefix + DataOutputStream dataOut = new DataOutputStream(outputStream); + dataOut.writeInt(data.length); + dataOut.write(data); + dataOut.flush(); + + } catch (SerializationException e) { + throw new IOException("Failed to serialize message", e); } } /** * Tries to read (deserialize) a MessageProtocol object from the socket. - * This call is "blocked" until an object is received. * * @return The "envelope" (MessageProtocol) that was received. * @throws IOException If the connection is lost, the stream is corrupted, or socket is not connected. * @throws ClassNotFoundException If the received object is unknown. */ public MessageProtocol receiveMessage() throws IOException, ClassNotFoundException { - if (!isConnected()) { - throw new IOException("Socket is not connected."); + if (socket == null || !socket.isConnected()) { + throw new IOException("Socket is not connected"); } - synchronized (inputStream) { - return (MessageProtocol) inputStream.readObject(); + + try { + // Lê um prefixo de 4 bytes - indicador de tamanho + DataInputStream dataIn = new DataInputStream(inputStream); + int length = dataIn.readInt(); + + if (length <= 0 || length > 10_000_000) { // Sanity check (10MB max) + throw new IOException("Invalid message length: " + length); + } + + // Ler dados da mensagem + byte[] data = new byte[length]; + dataIn.readFully(data); + + // Deserialize do JSON + return serializer.deserialize(data, MessageProtocol.class); + + } catch (SerializationException e) { + throw new IOException("Failed to deserialize message", e); } } /** * Closes the socket and all streams (Input and Output). - * It is called automatically if you use 'try-with-resources'. */ @Override public void close() throws IOException { - System.out.printf("[SocketConnection] Closing connection to %s:%d.%n", - socket != null ? socket.getInetAddress().getHostAddress() : "N/A", - socket != null ? socket.getPort() : -1); - try { - if (inputStream != null) inputStream.close(); - } catch (IOException e) { /* ignore */ } - - try { - if (outputStream != null) outputStream.close(); - } catch (IOException e) { /* ignore */ } - - if (socket != null && !socket.isClosed()) { - socket.close(); - } + if (inputStream != null) inputStream.close(); + if (outputStream != null) outputStream.close(); + if (socket != null) socket.close(); } /**