From 5568136d86689352dc8ee2fd49bb470c23510b56 Mon Sep 17 00:00:00 2001
From: Andrew Lalis <andrewlalisofficial@gmail.com>
Date: Sat, 11 Sep 2021 23:43:15 +0200
Subject: [PATCH] Cleaned up client encryption logic.

---
 .../andrewl/concord_client/ConcordClient.java | 32 ++++++-------------
 .../andrewl/concord_core/msg/Encryption.java  |  3 ++
 2 files changed, 12 insertions(+), 23 deletions(-)

diff --git a/client/src/main/java/nl/andrewl/concord_client/ConcordClient.java b/client/src/main/java/nl/andrewl/concord_client/ConcordClient.java
index c4b3f2f..377afc2 100644
--- a/client/src/main/java/nl/andrewl/concord_client/ConcordClient.java
+++ b/client/src/main/java/nl/andrewl/concord_client/ConcordClient.java
@@ -34,8 +34,8 @@ import java.util.Map;
 
 public class ConcordClient implements Runnable {
 	private final Socket socket;
-	private InputStream in;
-	private OutputStream out;
+	private final InputStream in;
+	private final OutputStream out;
 	private final Serializer serializer;
 
 	@Getter
@@ -48,9 +48,14 @@ public class ConcordClient implements Runnable {
 	public ConcordClient(String host, int port, String nickname, Path tokensFile) throws IOException {
 		this.eventManager = new EventManager(this);
 		this.socket = new Socket(host, port);
-		this.in = this.socket.getInputStream();
-		this.out = this.socket.getOutputStream();
 		this.serializer = new Serializer();
+		try {
+			var streams = Encryption.upgrade(socket.getInputStream(), socket.getOutputStream(), this.serializer);
+			this.in = streams.first();
+			this.out = streams.second();
+		} catch (GeneralSecurityException e) {
+			throw new IOException("Could not establish secure connection to the server.", e);
+		}
 		this.model = this.initializeConnectionToServer(nickname, tokensFile);
 
 		// Add event listeners.
@@ -75,7 +80,6 @@ public class ConcordClient implements Runnable {
 	 * messages, or if the server sends an unexpected response.
 	 */
 	private ClientModel initializeConnectionToServer(String nickname, Path tokensFile) throws IOException {
-		this.establishEncryption();
 		String token = this.getSessionToken(tokensFile);
 		this.serializer.writeMessage(new Identification(nickname, token), this.out);
 		Message reply = this.serializer.readMessage(this.in);
@@ -90,24 +94,6 @@ public class ConcordClient implements Runnable {
 		}
 	}
 
-	/**
-	 * Establishes an encrypted connection to the server. This should be the
-	 * first method which interacts with the server, since it sends and receives
-	 * specific key information, and all subsequent traffic should be encrypted.
-	 * @throws IOException If encryption could not be established.
-	 */
-	private void establishEncryption() throws IOException {
-		try {
-			System.out.println("Initializing end-to-end encryption with the server...");
-			var streams = Encryption.upgrade(this.in, this.out, this.serializer);
-			this.in = streams.first();
-			this.out = streams.second();
-			System.out.println("Successfully established cipher streams.");
-		} catch (GeneralSecurityException e) {
-			throw new IOException(e);
-		}
-	}
-
 	public void sendMessage(Message message) throws IOException {
 		this.serializer.writeMessage(message, this.out);
 	}
diff --git a/core/src/main/java/nl/andrewl/concord_core/msg/Encryption.java b/core/src/main/java/nl/andrewl/concord_core/msg/Encryption.java
index d2c1b18..ebc5610 100644
--- a/core/src/main/java/nl/andrewl/concord_core/msg/Encryption.java
+++ b/core/src/main/java/nl/andrewl/concord_core/msg/Encryption.java
@@ -19,6 +19,9 @@ import java.util.Arrays;
 import java.util.Collections;
 import java.util.List;
 
+/**
+ * Utility class for handling the establishment of encrypted communication.
+ */
 public class Encryption {
 	public static Pair<CipherInputStream, CipherOutputStream> upgrade(
 			InputStream in,