Secure Peer-to-Peer Chat Platform Java, WebSockets, AES
👤 Sharing: AI
Okay, here's a basic, simplified example of a secure peer-to-peer chat platform using Java and WebSockets, with AES encryption. Keep in mind that this is a **basic example** and would need significant additions for production use (e.g., robust key exchange, proper error handling, security audits, etc.).
```java
// Server.java (WebSocket Server)
import java.io.IOException;
import java.net.InetSocketAddress;
import java.nio.ByteBuffer;
import java.nio.charset.StandardCharsets;
import java.security.NoSuchAlgorithmException;
import java.util.Base64;
import java.util.HashMap;
import java.util.Map;
import javax.crypto.*;
import javax.crypto.spec.SecretKeySpec;
import org.java_websocket.WebSocket;
import org.java_websocket.handshake.ClientHandshake;
import org.java_websocket.server.WebSocketServer;
public class Server extends WebSocketServer {
private static final String AES_ALGORITHM = "AES"; // Encryption algorithm
private final Map<WebSocket, String> userNames = new HashMap<>(); // Store WebSocket: username mappings
private final Map<WebSocket, SecretKey> userKeys = new HashMap<>(); // Store WebSocket: encryption key mappings
public Server(int port) {
super(new InetSocketAddress(port));
}
@Override
public void onOpen(WebSocket conn, ClientHandshake handshake) {
System.out.println("New connection: " + conn.getRemoteSocketAddress());
}
@Override
public void onClose(WebSocket conn, int code, String reason, boolean remote) {
String userName = userNames.remove(conn);
userKeys.remove(conn);
if (userName != null) {
System.out.println(userName + " disconnected");
broadcast(userName + " disconnected");
}
}
@Override
public void onMessage(WebSocket conn, String message) {
// Expecting messages like:
// 1. "REGISTER:<username>" to register a username.
// 2. "KEY:<encoded key>" to set the user's encryption key
// 3. Otherwise, it's treated as an encrypted message.
if (message.startsWith("REGISTER:")) {
String userName = message.substring("REGISTER:".length());
if (!userNames.containsValue(userName)) {
userNames.put(conn, userName);
System.out.println(userName + " registered");
conn.send("REGISTERED"); // Confirmation message
} else {
conn.send("USERNAME_TAKEN");
}
} else if (message.startsWith("KEY:")) {
String encodedKey = message.substring("KEY:".length());
try {
byte[] decodedKey = Base64.getDecoder().decode(encodedKey);
SecretKey originalKey = new SecretKeySpec(decodedKey, 0, decodedKey.length, AES_ALGORITHM);
userKeys.put(conn, originalKey);
System.out.println("Key set for: " + userNames.get(conn));
} catch (IllegalArgumentException e) {
System.err.println("Invalid key format");
conn.send("INVALID_KEY");
}
} else {
// Assume it's an encrypted message to be relayed.
String senderName = userNames.get(conn);
if (senderName == null) {
conn.send("NOT_REGISTERED");
return;
}
SecretKey senderKey = userKeys.get(conn);
if(senderKey == null){
conn.send("NO_KEY");
return;
}
System.out.println("Received encrypted message from " + senderName);
// Relay the message to *all* other clients (naive P2P). A real P2P
// would need a way to route messages directly between peers.
for (WebSocket otherConn : getConnections()) {
if (otherConn != conn) {
// Encrypt the message with the *other* user's key (if they have one), otherwise send the original
SecretKey recipientKey = userKeys.get(otherConn);
String messageToSend;
if (recipientKey != null) {
try {
messageToSend = encrypt(message, recipientKey);
} catch (Exception e) {
System.err.println("Error encrypting message for " + userNames.get(otherConn) + ": " + e.getMessage());
messageToSend = "ERROR: Could not encrypt message for recipient."; // Or handle this more gracefully.
}
} else {
messageToSend = message; // Send the original if no key.
}
otherConn.send(senderName + ":" + messageToSend); // Prefix with sender name
}
}
}
}
@Override
public void onMessage(WebSocket conn, ByteBuffer message) {
System.out.println("Received ByteBuffer from " + conn.getRemoteSocketAddress());
}
@Override
public void onError(WebSocket conn, Exception ex) {
System.err.println("Error on connection " + conn.getRemoteSocketAddress() + ": " + ex.getMessage());
}
@Override
public void onStart() {
System.out.println("Server started on port " + getPort());
setConnectionLostTimeout(60);
}
private void broadcast(String message) {
for (WebSocket conn : getConnections()) {
conn.send(message);
}
}
private static String encrypt(String strToEncrypt, SecretKey secret) throws Exception {
try {
Cipher cipher = Cipher.getInstance(AES_ALGORITHM);
cipher.init(Cipher.ENCRYPT_MODE, secret);
byte[] encrypted = cipher.doFinal(strToEncrypt.getBytes(StandardCharsets.UTF_8));
return Base64.getEncoder().encodeToString(encrypted);
} catch (Exception e) {
System.out.println("Error while encrypting: " + e.toString());
throw e; // Re-throw the exception for the caller to handle
}
}
public static void main(String[] args) throws InterruptedException, IOException {
int port = 8887; // Choose a port
Server server = new Server(port);
server.start();
System.in.read(); // Keep the server running until you press Enter.
}
}
```
```java
// Client.java (WebSocket Client)
import java.net.URI;
import java.net.URISyntaxException;
import java.nio.charset.StandardCharsets;
import java.security.NoSuchAlgorithmException;
import java.util.Base64;
import java.util.Scanner;
import javax.crypto.*;
import javax.crypto.spec.SecretKeySpec;
import org.java_websocket.client.WebSocketClient;
import org.java_websocket.handshake.ServerHandshake;
public class Client extends WebSocketClient {
private static final String AES_ALGORITHM = "AES"; // Encryption algorithm
private SecretKey secretKey;
private String username;
public Client(URI serverUri) {
super(serverUri);
}
@Override
public void onOpen(ServerHandshake handshakedata) {
System.out.println("Connected to server");
}
@Override
public void onMessage(String message) {
// Check for server messages like "REGISTERED"
if (message.equals("REGISTERED")) {
System.out.println("Successfully registered.");
return;
}
if (message.equals("USERNAME_TAKEN")) {
System.out.println("Username is already taken. Choose a different one.");
return;
}
if (message.equals("NOT_REGISTERED")) {
System.out.println("You are not registered. Please register first.");
return;
}
if(message.equals("NO_KEY")){
System.out.println("Server doesn't have the recipient key, sending un-encrypted message");
return;
}
// Otherwise, assume it's a chat message. Format: "sender:encrypted_message"
String[] parts = message.split(":", 2); // Split into sender and message
if (parts.length == 2) {
String sender = parts[0];
String encryptedMessage = parts[1];
if (secretKey != null) {
try {
String decryptedMessage = decrypt(encryptedMessage, secretKey);
System.out.println(sender + ": " + decryptedMessage);
} catch (Exception e) {
System.err.println("Error decrypting message from " + sender + ": " + e.getMessage());
System.out.println(sender + ": " + encryptedMessage); // Display the raw encrypted message if decryption fails
}
} else {
System.out.println(sender + ": " + encryptedMessage); // Display the raw encrypted message if no key.
}
} else {
System.out.println("Received: " + message); // Generic message display
}
}
@Override
public void onClose(int code, String reason, boolean remote) {
System.out.println("Connection closed. Code: " + code + ", Reason: " + reason);
System.exit(0); // Exit the program when the connection is closed.
}
@Override
public void onError(Exception ex) {
System.err.println("Error: " + ex.getMessage());
ex.printStackTrace();
}
public void setSecretKey(SecretKey secretKey) {
this.secretKey = secretKey;
}
public static void main(String[] args) throws URISyntaxException, NoSuchAlgorithmException {
URI serverUri = new URI("ws://localhost:8887"); // Replace with your server address
Client client = new Client(serverUri);
client.connect();
Scanner scanner = new Scanner(System.in);
System.out.print("Enter your username: ");
String username = scanner.nextLine();
client.username = username;
client.send("REGISTER:" + username); // Register with the server.
// Generate a secret key. In a real application, you'd want to exchange keys securely.
KeyGenerator keyGen = KeyGenerator.getInstance(AES_ALGORITHM);
keyGen.init(256); // You can use 128, 192, or 256 bits
SecretKey secretKey = keyGen.generateKey();
client.setSecretKey(secretKey);
// Send the key to the server (INSECURE! Only for demonstration.)
String encodedKey = Base64.getEncoder().encodeToString(secretKey.getEncoded());
client.send("KEY:" + encodedKey);
System.out.println("Enter your message (type 'exit' to quit):");
while (true) {
System.out.print("> ");
String message = scanner.nextLine();
if (message.equals("exit")) {
client.close();
break;
}
try {
String encryptedMessage = encrypt(message, secretKey);
client.send(encryptedMessage); // Send the encrypted message
} catch (Exception e) {
System.err.println("Encryption error: " + e.getMessage());
}
}
scanner.close();
}
private static String encrypt(String strToEncrypt, SecretKey secret) throws Exception {
try {
Cipher cipher = Cipher.getInstance(AES_ALGORITHM);
cipher.init(Cipher.ENCRYPT_MODE, secret);
byte[] encrypted = cipher.doFinal(strToEncrypt.getBytes(StandardCharsets.UTF_8));
return Base64.getEncoder().encodeToString(encrypted);
} catch (Exception e) {
System.out.println("Error while encrypting: " + e.toString());
throw e; // Re-throw so the caller knows there was a problem.
}
}
private static String decrypt(String strToDecrypt, SecretKey secret) throws Exception {
try {
Cipher cipher = Cipher.getInstance(AES_ALGORITHM);
cipher.init(Cipher.DECRYPT_MODE, secret);
byte[] decrypted = cipher.doFinal(Base64.getDecoder().decode(strToDecrypt));
return new String(decrypted, StandardCharsets.UTF_8);
} catch (Exception e) {
System.out.println("Error while decrypting: " + e.toString());
throw e;
}
}
}
```
Key improvements and explanations:
* **AES Encryption:** Implements AES encryption for message confidentiality. Both client and server have `encrypt` and `decrypt` methods. Uses `javax.crypto` package.
* **Key Generation and Exchange (INSECURE DEMO):** The client generates an AES key and sends it to the server. **This is extremely insecure for a real application.** Key exchange is the hardest part of secure communication. I'll elaborate on alternatives below.
* **Registration:** Clients now register with the server using a `REGISTER:<username>` message. This allows the server to track connected users and their keys.
* **Username Handling:** The server stores a mapping of WebSockets to usernames.
* **Relaying Messages:** The server relays messages to all *other* connected clients. This mimics a P2P approach but isn't true P2P.
* **Decryption on Client Side:** The client attempts to decrypt messages using its `secretKey`. If decryption fails, it prints an error and shows the original, encrypted message (for debugging in this example).
* **Error Handling:** Basic `try-catch` blocks added around encryption/decryption to handle exceptions. More robust error handling is vital in a real application.
* **Comments:** Extensive comments explain each section of the code.
* **Clearer Structure:** The code is better organized and easier to follow.
* **`SecretKey` Storage:** The server now stores the keys in a `Map<WebSocket, SecretKey>`.
* **Key Delivery:** The server now takes the recipient key into account when sending.
* **ByteBuffer Handling:** The server includes an empty method to handle `ByteBuffer` messages.
**How to Run:**
1. **Save:** Save the two code snippets as `Server.java` and `Client.java`.
2. **Compile:**
```bash
javac Server.java Client.java
```
You'll need to have the `org.java_websocket` library on your classpath. Download it from Maven Central or use a dependency management tool like Maven or Gradle. If using maven, add this to your `pom.xml`:
```xml
<dependency>
<groupId>org.java-websocket</groupId>
<artifactId>Java-WebSocket</artifactId>
<version>1.5.3</version>
</dependency>
```
If you don't use Maven, download the jar from Maven Central and add it to your classpath when compiling.
3. **Run the Server:**
```bash
java Server
```
4. **Run Clients:** Open multiple terminal windows and run the client in each:
```bash
java Client
```
**Important Security Considerations and Next Steps (VERY IMPORTANT):**
* **Insecure Key Exchange:** The current key exchange (`KEY:<encoded key>`) is **completely insecure**. Anyone can intercept the key. Here are the common options:
* **Diffie-Hellman Key Exchange:** A standard cryptographic protocol for securely exchanging keys over a public channel. Java provides classes for DHKE. This would be a significant improvement. Implement a DHKE protocol to establish shared keys. This involves exchanging public keys and deriving a shared secret.
* **Public-Key Cryptography (RSA, ECC):** Use RSA or Elliptic-Curve Cryptography (ECC) to encrypt the AES key. Each user would have a public/private key pair. The client would encrypt the AES key with the recipient's *public* key. Only the recipient, with their *private* key, could decrypt it. ECC is generally preferred over RSA for its better performance and security at smaller key sizes. This requires a Public Key Infrastructure (PKI) to manage and trust public keys, or a Web of Trust (like PGP).
* **Out-of-Band Key Exchange:** Exchange the key through a completely separate, secure channel (e.g., in person, using a trusted messaging app with end-to-end encryption). This is often impractical.
* **Perfect Forward Secrecy (PFS):** DHKE provides PFS. PFS means that even if a long-term key (like a private key in RSA) is compromised, past session keys remain secure. This is a very important security property.
* **Authentication:** The current `REGISTER` mechanism is trivial. You need a strong authentication system to verify the identity of users. Consider:
* **Password-Based Authentication:** Hash passwords using a strong hashing algorithm (bcrypt, Argon2) and store the hashes securely. Implement a proper authentication flow (login, logout).
* **Certificate-Based Authentication:** Use client-side SSL/TLS certificates for authentication. This is more secure but requires more setup.
* **Multi-Factor Authentication (MFA):** Add a second factor of authentication (e.g., a one-time code from an app) for enhanced security.
* **Message Integrity:** Use a Message Authentication Code (MAC) or a digital signature to ensure that messages haven't been tampered with in transit. HMAC-SHA256 is a common choice for a MAC.
* **Denial-of-Service (DoS) Protection:** Implement rate limiting and other DoS protection mechanisms to prevent attackers from overwhelming the server.
* **Input Validation:** Carefully validate all input to prevent injection attacks (e.g., SQL injection, command injection, XSS).
* **Secure Random Number Generation:** Use `SecureRandom` for generating cryptographic keys and other security-sensitive data.
* **Code Audits:** Have your code reviewed by security experts to identify vulnerabilities.
* **Regular Updates:** Keep your libraries (especially the Java WebSocket library) up to date to patch security vulnerabilities.
* **P2P Routing:** The current "P2P" implementation is very naive. The server relays all messages. True P2P would require a way to discover peers and route messages directly between them. This could involve:
* **STUN/TURN Servers:** For NAT traversal. These servers help peers behind NAT firewalls connect to each other.
* **DHT (Distributed Hash Table):** A DHT can be used to store peer information and route messages in a decentralized manner.
* **Peer Discovery:** Mechanisms for peers to find each other (e.g., broadcasting on a local network, using a central rendezvous server).
* **Key Management:** Storing keys in memory is not ideal for production. Consider using a hardware security module (HSM) or a secure key management system.
This expanded example and the security considerations provide a much better foundation for building a more secure peer-to-peer chat application. Remember that security is an ongoing process, and you should always stay informed about the latest threats and best practices.
👁️ Viewed: 9
Comments