From 022bb272e6f3b43a66ac7ec6a65f5ec9fca245b6 Mon Sep 17 00:00:00 2001
From: shamoon <4887959+shamoon@users.noreply.github.com>
Date: Tue, 8 Aug 2023 16:27:40 -0700
Subject: [PATCH] Restrict status messages by owner if set

---
 .../data/websocket-consumer-status-message.ts |  1 +
 .../services/consumer-status.service.spec.ts  | 45 ++++++++++++++++++-
 .../app/services/consumer-status.service.ts   | 14 +++++-
 src/documents/consumer.py                     |  3 +-
 4 files changed, 60 insertions(+), 3 deletions(-)

diff --git a/src-ui/src/app/data/websocket-consumer-status-message.ts b/src-ui/src/app/data/websocket-consumer-status-message.ts
index aecdda7c0..d1ac590b1 100644
--- a/src-ui/src/app/data/websocket-consumer-status-message.ts
+++ b/src-ui/src/app/data/websocket-consumer-status-message.ts
@@ -6,4 +6,5 @@ export interface WebsocketConsumerStatusMessage {
   status?: string
   message?: string
   document_id: number
+  owner_id?: number
 }
diff --git a/src-ui/src/app/services/consumer-status.service.spec.ts b/src-ui/src/app/services/consumer-status.service.spec.ts
index 3725f847d..d3867e889 100644
--- a/src-ui/src/app/services/consumer-status.service.spec.ts
+++ b/src-ui/src/app/services/consumer-status.service.spec.ts
@@ -12,6 +12,7 @@ import { environment } from 'src/environments/environment'
 import { DocumentService } from './rest/document.service'
 import { HttpEventType, HttpResponse } from '@angular/common/http'
 import WS from 'jest-websocket-mock'
+import { SettingsService } from './settings.service'
 
 describe('ConsumerStatusService', () => {
   let httpTestingController: HttpTestingController
@@ -24,7 +25,21 @@ describe('ConsumerStatusService', () => {
 
   beforeEach(() => {
     TestBed.configureTestingModule({
-      providers: [ConsumerStatusService, DocumentService],
+      providers: [
+        ConsumerStatusService,
+        DocumentService,
+        SettingsService,
+        {
+          provide: SettingsService,
+          useValue: {
+            currentUser: {
+              id: 1,
+              username: 'testuser',
+              is_superuser: false,
+            },
+          },
+        },
+      ],
       imports: [HttpClientTestingModule],
     })
 
@@ -275,4 +290,32 @@ describe('ConsumerStatusService', () => {
       1
     )
   })
+
+  it('should not notify current user if document has different expected owner', () => {
+    consumerStatusService.connect()
+    server.send({
+      task_id: '1234',
+      filename: 'file1.pdf',
+      current_progress: 50,
+      max_progress: 100,
+      docuement_id: 12,
+      owner_id: 1,
+      status: 'WORKING',
+    })
+
+    server.send({
+      task_id: '5678',
+      filename: 'file2.pdf',
+      current_progress: 50,
+      max_progress: 100,
+      docuement_id: 13,
+      owner_id: 2,
+      status: 'WORKING',
+    })
+
+    consumerStatusService.disconnect()
+    expect(consumerStatusService.getConsumerStatusNotCompleted()).toHaveLength(
+      1
+    )
+  })
 })
diff --git a/src-ui/src/app/services/consumer-status.service.ts b/src-ui/src/app/services/consumer-status.service.ts
index 394975333..3e21da138 100644
--- a/src-ui/src/app/services/consumer-status.service.ts
+++ b/src-ui/src/app/services/consumer-status.service.ts
@@ -2,6 +2,7 @@ import { Injectable } from '@angular/core'
 import { Subject } from 'rxjs'
 import { environment } from 'src/environments/environment'
 import { WebsocketConsumerStatusMessage } from '../data/websocket-consumer-status-message'
+import { SettingsService } from './settings.service'
 
 // see ConsumerFilePhase in src/documents/consumer.py
 export enum FileStatusPhase {
@@ -44,6 +45,8 @@ export class FileStatus {
 
   documentId: number
 
+  ownerId: number
+
   getProgress(): number {
     switch (this.phase) {
       case FileStatusPhase.STARTED:
@@ -81,7 +84,7 @@ export class FileStatus {
   providedIn: 'root',
 })
 export class ConsumerStatusService {
-  constructor() {}
+  constructor(private settingsService: SettingsService) {}
 
   private statusWebSocket: WebSocket
 
@@ -143,6 +146,15 @@ export class ConsumerStatusService {
     this.statusWebSocket.onmessage = (ev) => {
       let statusMessage: WebsocketConsumerStatusMessage = JSON.parse(ev['data'])
 
+      // tasks are async so we rely on checking user id
+      if (
+        statusMessage.owner_id &&
+        statusMessage.owner_id !== this.settingsService.currentUser?.id &&
+        !this.settingsService.currentUser?.is_superuser
+      ) {
+        return
+      }
+
       let statusMessageGet = this.get(
         statusMessage.task_id,
         statusMessage.filename
diff --git a/src/documents/consumer.py b/src/documents/consumer.py
index 6fa830101..0ec6090c2 100644
--- a/src/documents/consumer.py
+++ b/src/documents/consumer.py
@@ -90,6 +90,7 @@ class Consumer(LoggingMixin):
             "status": status,
             "message": message,
             "document_id": document_id,
+            "owner_id": self.override_owner_id if self.override_owner_id else None,
         }
         async_to_sync(self.channel_layer.group_send)(
             "status_updates",
@@ -118,7 +119,7 @@ class Consumer(LoggingMixin):
         self.override_document_type_id = None
         self.override_asn = None
         self.task_id = None
-        self.owner_id = None
+        self.override_owner_id = None
 
         self.channel_layer = get_channel_layer()