Upload files to "src/main/java/com/miniapps/messaging/controller"

src/main/java/com/miniapps/messaging/controller/MessageServer.java

@@ -0,0 +1,239 @@
+package com.miniapps.messaging.controller;
+
+import org.springframework.stereotype.Controller;
+import org.springframework.web.bind.annotation.PostMapping;
+import org.springframework.web.bind.annotation.RequestBody;
+import org.springframework.web.bind.annotation.ResponseBody;
+
+import java.security.KeyFactory;
+import java.security.PublicKey;
+import java.security.Signature;
+import java.security.spec.X509EncodedKeySpec;
+import java.time.Instant;
+import java.util.*;
+import java.util.stream.Collectors;
+
+
+@Controller("/")
+public class MessageServer {
+
+    private final static String KEY_ALGORITHM = "RSA";
+    private final static String SIGN_ALGORITHM = "SHA256withRSA";
+
+    // Instance Data
+    private final Set<Message> storedMessages = new HashSet<>();
+    private final Set<String> allowedPeers = new HashSet<>();
+    private String masterKey = "";
+
+    // --- Data Transfer Objects (DTOs) ---
+    // Public fields are kept for simplicity as requested.
+
+    public static class Message {
+        public String encryptedMessage;
+        public String senderPublicKey;
+        public String recipientPublicKey;
+        public String timestamp; // Should ideally be Instant or long
+
+        /**
+         * The canonical string representation of the message content for signing.
+         * The order is crucial.
+         */
+        public String getStringToSign() {
+            // Note: timestamp is included here. On the client, this is the creation time.
+            return this.encryptedMessage + ";" +
+                    this.senderPublicKey + ";" +
+                    this.recipientPublicKey + ";" +
+                    this.timestamp;
+        }
+    }
+
+    public static class MessageEnvelope {
+        public Set<Message> messages;
+        public String publicKey;
+        public String signature;
+
+        /**
+         * The canonical string representation of the entire envelope content for signing.
+         * Order: All messages (sorted and joined) + sender's public key.
+         */
+        public String getStringToSign() {
+            // Sort messages to ensure deterministic signing (crucial for Sets)
+            String messageData = messages.stream()
+                    .map(Message::getStringToSign)
+                    .sorted() // Ensure a consistent signing order
+                    .collect(Collectors.joining("|"));
+
+            return messageData + ";" + this.publicKey;
+        }
+    }
+
+    // --- Server Initialization ---
+
+    public void init(String keysStr) {
+        // FIX: Corrected fixed-size list issue and used get(0) for compatibility
+        List<String> keys = new ArrayList<>(Arrays.asList(keysStr.split(";")));
+        if (keys.isEmpty()) {
+            throw new IllegalArgumentException("Initialization string cannot be empty.");
+        }
+        masterKey = keys.get(0);
+        allowedPeers.addAll(keys);
+    }
+
+    // --- Core Utility Methods ---
+
+    /**
+     * Verifies a digital signature against a message using a public key.
+     * @param messageToVerify The canonical string of the data that was signed.
+     * @param publicKey The base64-encoded X.509 public key of the signer.
+     * @param signature The base64-encoded signature.
+     * @return true if the signature is valid, false otherwise.
+     */
+    public static boolean verifySignature(String messageToVerify, String publicKey, String signature) {
+        // Renamed variables for clarity
+        try {
+            KeyFactory keyFactory = KeyFactory.getInstance(KEY_ALGORITHM);
+            // Remove whitespace before decoding
+            X509EncodedKeySpec keySpec = new X509EncodedKeySpec(
+                    Base64.getDecoder().decode(publicKey.replaceAll("\\s+", ""))
+            );
+            PublicKey pubKey = keyFactory.generatePublic(keySpec);
+
+            Signature sig = Signature.getInstance(SIGN_ALGORITHM);
+            sig.initVerify(pubKey);
+            sig.update(messageToVerify.getBytes());
+
+            return sig.verify(Base64.getDecoder().decode(signature.replaceAll("\\s+", "")));
+
+        } catch (Exception e) {
+            // BEST PRACTICE: Log the exception to aid debugging security failures
+            System.err.println("Signature verification failed: " + e.getMessage());
+            // e.printStackTrace(); // Use proper logging in a real app
+            return false;
+        }
+    }
+
+    /**
+     * Checks if the timestamps in the messages are fresh (e.g., within 5 minutes).
+     * This is a critical step to prevent replay attacks.
+     */
+    private boolean isFresh(MessageEnvelope envelope) {
+        Instant fiveMinutesAgo = Instant.now().minusSeconds(300);
+        for (Message m : envelope.messages) {
+            try {
+                if (Instant.parse(m.timestamp).isBefore(fiveMinutesAgo)) {
+                    return false; // Message is too old
+                }
+            } catch (Exception e) {
+                System.err.println("Invalid message timestamp format: " + m.timestamp);
+                return false; // Invalid timestamp format
+            }
+        }
+        return true;
+    }
+
+    // --- Peer List Operations ---
+
+    public String addPeerList(MessageEnvelope envelope) {
+        // FIX: Signature verification now uses the envelope's content string
+        if (!envelope.publicKey.equals(masterKey)) {
+            return "permission denied: only master key can add peers";
+        }
+
+        if (!verifySignature(envelope.getStringToSign(), envelope.publicKey, envelope.signature)) {
+            return "invalid signature";
+        }
+
+        allowedPeers.addAll(
+                envelope.messages.stream().map(m -> m.recipientPublicKey).collect(Collectors.toSet())
+        );
+        return "OK";
+    }
+
+    public String removePeers(MessageEnvelope envelope) {
+        // FIX: Corrected master key and signature logic
+        if (!envelope.publicKey.equals(masterKey)) {
+            return "permission denied: only master key can remove peers";
+        }
+
+        if (!verifySignature(envelope.getStringToSign(), envelope.publicKey, envelope.signature)) {
+            return "invalid signature";
+        }
+
+        allowedPeers.removeAll(
+                envelope.messages.stream().map(m -> m.recipientPublicKey).collect(Collectors.toSet())
+        );
+        return "OK";
+    }
+
+    public String listPeers() {
+        return allowedPeers.stream().sorted().collect(Collectors.joining("|"));
+    }
+
+    // --- Server Message Operations ---
+    @ResponseBody
+    @PostMapping("/messages/receiveMessage")
+    public String receiveMessage(@RequestBody MessageEnvelope envelope)
+    {
+
+        if(!masterKey.isEmpty() && !allowedPeers.contains(envelope.publicKey))
+        {
+            return null;
+        }
+
+        // FIX: Signature verification now uses the envelope's content string
+        if (!verifySignature(envelope.getStringToSign(), envelope.publicKey, envelope.signature)) {
+            return "invalid signature";
+        }
+
+
+
+        // SECURITY CHECK: Replay protection
+        if (!isFresh(envelope)) {
+            return "invalid timestamp: message is too old";
+        }
+
+        for (Message message : envelope.messages) {
+            // Check: Ensure the sender's public key inside the message matches the envelope signer
+            if (message.senderPublicKey.equals(envelope.publicKey)) {
+
+                // FIX: Removed side effect. Creating a new Message object to store.
+                Message storedMessage = new Message();
+                storedMessage.encryptedMessage = message.encryptedMessage;
+                storedMessage.senderPublicKey = message.senderPublicKey;
+                storedMessage.recipientPublicKey = message.recipientPublicKey;
+                // Add the server's time (optional, but good for tracking delivery)
+                storedMessage.timestamp = Instant.now().toString();
+
+                storedMessages.add(storedMessage);
+            }
+        }
+        return "OK";
+    }
+
+    @ResponseBody
+    @PostMapping("/messages/retrieveMessages")
+    public Set<Message> retrieveMessages(@RequestBody MessageEnvelope envelope)
+    {
+
+        if(!masterKey.isEmpty() && !allowedPeers.contains(envelope.publicKey))
+        {
+            return null;
+        }
+        // FIX: Signature verification now uses the envelope's content string
+        if (!verifySignature(envelope.getStringToSign(), envelope.publicKey, envelope.signature)) {
+            // Using null to clearly signal a verification/authorization failure
+            return null;
+        }
+
+        // Filter messages for the requesting recipient
+        Set<Message> retrieved = storedMessages.stream()
+                .filter(m -> m.recipientPublicKey.equals(envelope.publicKey))
+                .collect(Collectors.toSet());
+
+        // Remove the retrieved messages from storage
+        //comment it if messages get cancelled every day at midnight
+        storedMessages.removeAll(retrieved);
+
+        return retrieved;
+    }
+}