1 files changed, 28 insertions, 13 deletions
diff --git a/src/api/streaming.ts b/src/api/streaming.ts
index 38068d1e3d..7a8d2d4354 100644
--- a/src/api/streaming.ts
+++ b/src/api/streaming.ts
@@ -2,6 +2,8 @@ import * as http from 'http';
 import * as websocket from 'websocket';
 import * as redis from 'redis';
 import User from './models/user';
+import AccessToken from './models/access-token';
+import isNativeToken from './common/is-native-token';
 
 import homeStream from './stream/home';
 import messagingStream from './stream/messaging';
@@ -17,7 +19,13 @@ module.exports = (server: http.Server) => {
 	ws.on('request', async (request) => {
 		const connection = request.accept();
 
-		const user = await authenticate(connection);
+		const user = await authenticate(connection, request.resourceURL.query.i);
+
+		if (user == null) {
+			connection.send('authentication-failed');
+			connection.close();
+			return;
+		}
 
 		// Connect to Redis
 		const subscriber = redis.createClient(
@@ -41,29 +49,36 @@ module.exports = (server: http.Server) => {
 	});
 };
 
-function authenticate(connection: websocket.connection): Promise<any> {
-	return new Promise((resolve, reject) => {
-		// Listen first message
-		connection.once('message', async (data) => {
-			const msg = JSON.parse(data.utf8Data);
-
+function authenticate(connection: websocket.connection, token: string): Promise<any> {
+	return new Promise(async (resolve, reject) => {
+		if (isNativeToken(token)) {
 			// Fetch user
 			// SELECT _id
 			const user = await User
 				.findOne({
-					token: msg.i
+					token: token
 				}, {
 					_id: true
 				});
 
-			if (user === null) {
-				connection.close();
-				return;
+			resolve(user);
+		} else {
+			const accessToken = await AccessToken.findOne({
+				hash: token
+			});
+
+			if (accessToken == null) {
+				return reject('invalid signature');
 			}
 
-			connection.send('authenticated');
+			// Fetch user
+			// SELECT _id
+			const user = await User
+				.findOne({ _id: accessToken.user_id }, {
+					_id: true
+				});
 
 			resolve(user);
-		});
+		}
 	});
 }