From d63c3c0cef9ab4e3713c6e4bea28a02f1c1e11a6 Mon Sep 17 00:00:00 2001
From: ThibG <thib@sitedethib.com>
Date: Fri, 24 May 2019 15:21:42 +0200
Subject: [PATCH] Improve streaming server security (#10818)

* Check OAuth token scopes in the streaming API

* Use Sec-WebSocket-Protocol instead of query string to pass WebSocket token

Inspired by https://github.com/kubevirt/kubevirt/issues/1242
---
 app/javascript/mastodon/stream.js |  6 +---
 streaming/index.js                | 50 ++++++++++++++++++++++++++-----
 2 files changed, 44 insertions(+), 12 deletions(-)

diff --git a/app/javascript/mastodon/stream.js b/app/javascript/mastodon/stream.js
index 306a068b7..c4642344f 100644
--- a/app/javascript/mastodon/stream.js
+++ b/app/javascript/mastodon/stream.js
@@ -71,11 +71,7 @@ export function connectStream(path, pollingRefresh = null, callbacks = () => ({
 export default function getStream(streamingAPIBaseURL, accessToken, stream, { connected, received, disconnected, reconnected }) {
   const params = [ `stream=${stream}` ];
 
-  if (accessToken !== null) {
-    params.push(`access_token=${accessToken}`);
-  }
-
-  const ws = new WebSocketClient(`${streamingAPIBaseURL}/api/v1/streaming/?${params.join('&')}`);
+  const ws = new WebSocketClient(`${streamingAPIBaseURL}/api/v1/streaming/?${params.join('&')}`, accessToken);
 
   ws.onopen      = connected;
   ws.onmessage   = e => received(JSON.parse(e.data));
diff --git a/streaming/index.js b/streaming/index.js
index 558267700..639867b28 100644
--- a/streaming/index.js
+++ b/streaming/index.js
@@ -195,14 +195,14 @@ const startWorker = (workerId) => {
     next();
   };
 
-  const accountFromToken = (token, req, next) => {
+  const accountFromToken = (token, allowedScopes, req, next) => {
     pgPool.connect((err, client, done) => {
       if (err) {
         next(err);
         return;
       }
 
-      client.query('SELECT oauth_access_tokens.resource_owner_id, users.account_id, users.chosen_languages FROM oauth_access_tokens INNER JOIN users ON oauth_access_tokens.resource_owner_id = users.id WHERE oauth_access_tokens.token = $1 AND oauth_access_tokens.revoked_at IS NULL LIMIT 1', [token], (err, result) => {
+      client.query('SELECT oauth_access_tokens.resource_owner_id, users.account_id, users.chosen_languages, oauth_access_tokens.scopes FROM oauth_access_tokens INNER JOIN users ON oauth_access_tokens.resource_owner_id = users.id WHERE oauth_access_tokens.token = $1 AND oauth_access_tokens.revoked_at IS NULL LIMIT 1', [token], (err, result) => {
         done();
 
         if (err) {
@@ -218,18 +218,29 @@ const startWorker = (workerId) => {
           return;
         }
 
+        const scopes = result.rows[0].scopes.split(' ');
+
+        if (allowedScopes.size > 0 && !scopes.some(scope => allowedScopes.includes(scope))) {
+          err = new Error('Access token does not cover required scopes');
+          err.statusCode = 401;
+
+          next(err);
+          return;
+        }
+
         req.accountId = result.rows[0].account_id;
         req.chosenLanguages = result.rows[0].chosen_languages;
+        req.allowNotifications = scopes.some(scope => ['read', 'read:notifications'].includes(scope));
 
         next();
       });
     });
   };
 
-  const accountFromRequest = (req, next, required = true) => {
+  const accountFromRequest = (req, next, required = true, allowedScopes = ['read']) => {
     const authorization = req.headers.authorization;
     const location = url.parse(req.url, true);
-    const accessToken = location.query.access_token;
+    const accessToken = location.query.access_token || req.headers['sec-websocket-protocol'];
 
     if (!authorization && !accessToken) {
       if (required) {
@@ -246,7 +257,7 @@ const startWorker = (workerId) => {
 
     const token = authorization ? authorization.replace(/^Bearer /, '') : accessToken;
 
-    accountFromToken(token, req, next);
+    accountFromToken(token, allowedScopes, req, next);
   };
 
   const PUBLIC_STREAMS = [
@@ -261,6 +272,16 @@ const startWorker = (workerId) => {
   const wsVerifyClient = (info, cb) => {
     const location = url.parse(info.req.url, true);
     const authRequired = !PUBLIC_STREAMS.some(stream => stream === location.query.stream);
+    const allowedScopes = [];
+
+    if (authRequired) {
+      allowedScopes.push('read');
+      if (location.query.stream === 'user:notification') {
+        allowedScopes.push('read:notifications');
+      } else {
+        allowedScopes.push('read:statuses');
+      }
+    }
 
     accountFromRequest(info.req, err => {
       if (!err) {
@@ -269,7 +290,7 @@ const startWorker = (workerId) => {
         log.error(info.req.requestId, err.toString());
         cb(false, 401, 'Unauthorized');
       }
-    }, authRequired);
+    }, authRequired, allowedScopes);
   };
 
   const PUBLIC_ENDPOINTS = [
@@ -286,7 +307,18 @@ const startWorker = (workerId) => {
     }
 
     const authRequired = !PUBLIC_ENDPOINTS.some(endpoint => endpoint === req.path);
-    accountFromRequest(req, next, authRequired);
+    const allowedScopes = [];
+
+    if (authRequired) {
+      allowedScopes.push('read');
+      if (req.path === '/api/v1/streaming/user/notification') {
+        allowedScopes.push('read:notifications');
+      } else {
+        allowedScopes.push('read:statuses');
+      }
+    }
+
+    accountFromRequest(req, next, authRequired, allowedScopes);
   };
 
   const errorMiddleware = (err, req, res, {}) => {
@@ -339,6 +371,10 @@ const startWorker = (workerId) => {
         return;
       }
 
+      if (event === 'notification' && !req.allowNotifications) {
+        return;
+      }
+
       // Only messages that may require filtering are statuses, since notifications
       // are already personalized and deletes do not matter
       if (!needsFiltering || event !== 'update') {