From 86e979d0b3eb641df53f2e232d0fe84af40aea50 Mon Sep 17 00:00:00 2001
From: Yonghao Yu <yyhao1@gmail.com>
Date: Fri, 6 Dec 2024 15:55:29 -0500
Subject: [PATCH 01/10] draft

---
 ...log-246-add-is-vwb-egress-event-column.xml | 14 ++++++
 api/db/changelog/db.changelog-master.xml      |  1 +
 .../auditors/EgressEventAuditor.java          |  7 ++-
 .../auditors/EgressEventAuditorImpl.java      | 33 ++++++++++++++
 .../api/VwbEgressAdminController.java         | 10 +++++
 .../workbench/db/model/DbEgressEvent.java     | 43 ++++++++++++++++++-
 .../exfiltration/EgressEventService.java      |  2 +
 .../exfiltration/EgressEventServiceImpl.java  | 25 +++++++++++
 8 files changed, 132 insertions(+), 3 deletions(-)
 create mode 100644 api/db/changelog/db.changelog-246-add-is-vwb-egress-event-column.xml

diff --git a/api/db/changelog/db.changelog-246-add-is-vwb-egress-event-column.xml b/api/db/changelog/db.changelog-246-add-is-vwb-egress-event-column.xml
new file mode 100644
index 00000000000..40356fe88a8
--- /dev/null
+++ b/api/db/changelog/db.changelog-246-add-is-vwb-egress-event-column.xml
@@ -0,0 +1,14 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<databaseChangeLog
+        xmlns="http://www.liquibase.org/xml/ns/dbchangelog"
+        xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
+        xsi:schemaLocation="http://www.liquibase.org/xml/ns/dbchangelog
+                      http://www.liquibase.org/xml/ns/dbchangelog/dbchangelog-3.1.xsd">
+    <changeSet author="yonghao" id="db.changelog-246-add-is-vwb-egress-event-column">
+        <addColumn tableName="egress_event">
+            <column name="is_vwb" type="boolean" />
+            <column name="vwb_workspace_id" type="varchar(80)" />
+            <column name="vwb_vm_name" type="varchar(80)" />
+        </addColumn>
+    </changeSet>
+</databaseChangeLog>
diff --git a/api/db/changelog/db.changelog-master.xml b/api/db/changelog/db.changelog-master.xml
index e5ea740e533..2fd01988180 100644
--- a/api/db/changelog/db.changelog-master.xml
+++ b/api/db/changelog/db.changelog-master.xml
@@ -253,6 +253,7 @@
   <include file="changelog/db.changelog-243-replace-extension-count.xml"/>
   <include file="changelog/db.changelog-244-add-is-vwb-workspace-column.xml"/>
   <include file="changelog/db.changelog-245-add-vwb-template-id-cdr-version.xml"/>
+  <include file="changelog/db.changelog-246-add-is-vwb-egress-event-column.xml"/>
   <!--
    Note: to update the DB locally, do the following:
    - Migrate schema changes: `./project.rb run-local-all-migrations`
diff --git a/api/src/main/java/org/pmiops/workbench/actionaudit/auditors/EgressEventAuditor.java b/api/src/main/java/org/pmiops/workbench/actionaudit/auditors/EgressEventAuditor.java
index 798a645d223..3da09118ffd 100644
--- a/api/src/main/java/org/pmiops/workbench/actionaudit/auditors/EgressEventAuditor.java
+++ b/api/src/main/java/org/pmiops/workbench/actionaudit/auditors/EgressEventAuditor.java
@@ -5,6 +5,7 @@
 import org.pmiops.workbench.db.model.DbUser;
 import org.pmiops.workbench.model.SumologicEgressEvent;
 import org.pmiops.workbench.model.SumologicEgressEventRequest;
+import org.pmiops.workbench.model.VwbEgressEventRequest;
 
 /**
  * Auditor service which handles collecting audit logs for high-egress events. These events are
@@ -17,7 +18,11 @@ public interface EgressEventAuditor {
    * event log in the target workspace.
    */
   void fireEgressEvent(SumologicEgressEvent event);
-
+  /**
+   * Decorates a Verily Workbench high-egress event with Workbench metadata and fires an audit
+   * event log in the target workspace.
+   */
+  void fireVwbEgressEvent(VwbEgressEventRequest event, DbUser dbUser);
   /**
    * Decorates a Sumologic-reported high-egress event with Workbench metadata and fires an audit
    * event log in the target workspace for the specified user.
diff --git a/api/src/main/java/org/pmiops/workbench/actionaudit/auditors/EgressEventAuditorImpl.java b/api/src/main/java/org/pmiops/workbench/actionaudit/auditors/EgressEventAuditorImpl.java
index ee45dfd93ee..fb3a2519588 100644
--- a/api/src/main/java/org/pmiops/workbench/actionaudit/auditors/EgressEventAuditorImpl.java
+++ b/api/src/main/java/org/pmiops/workbench/actionaudit/auditors/EgressEventAuditorImpl.java
@@ -29,6 +29,7 @@
 import org.pmiops.workbench.model.SumologicEgressEvent;
 import org.pmiops.workbench.model.SumologicEgressEventRequest;
 import org.pmiops.workbench.model.UserRole;
+import org.pmiops.workbench.model.VwbEgressEventRequest;
 import org.pmiops.workbench.workspaces.WorkspaceService;
 import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.beans.factory.annotation.Qualifier;
@@ -109,6 +110,38 @@ public void fireEgressEvent(SumologicEgressEvent event) {
     fireEvent(agentId, agentEmail, dbWorkspace, event);
   }
 
+  @Override
+  public void fireVwbEgressEvent(VwbEgressEventRequest event, DbUser user) {
+    String actionId = actionIdProvider.get();
+    Builder baseEventBuilder =
+        ActionAuditEvent.builder()
+            .timestamp(clock.millis())
+            .actionId(actionId)
+            .actionType(ActionType.DETECT_HIGH_EGRESS_EVENT)
+            .agentType(AgentType.USER)
+            .agentIdMaybe(user.getUserId())
+            .agentEmailMaybe(user.getUsername())
+            .targetType(TargetType.WORKSPACE);
+    var events = new ArrayList<ActionAuditEvent>();
+    // File 3 events: Workspace, VM, GCP projects. Can add more if needed in the future
+    events.add(
+        baseEventBuilder
+            .targetPropertyMaybe("vwbWorkspaceId")
+            .newValueMaybe(event.getWorkspaceId())
+            .build());
+    events.add(
+        baseEventBuilder
+            .targetPropertyMaybe("gcpProjectId")
+            .newValueMaybe(event.getGcpProjectId())
+            .build());
+    events.add(
+        baseEventBuilder
+            .targetPropertyMaybe("vmName")
+            .newValueMaybe(event.getVmName())
+            .build());
+    actionAuditService.send(events);
+  }
+
   @Override
   public void fireEgressEventForUser(SumologicEgressEvent event, DbUser user) {
     DbWorkspace dbWorkspace = getDbWorkspace(event);
diff --git a/api/src/main/java/org/pmiops/workbench/api/VwbEgressAdminController.java b/api/src/main/java/org/pmiops/workbench/api/VwbEgressAdminController.java
index 7e2d5bff55b..753fc035585 100644
--- a/api/src/main/java/org/pmiops/workbench/api/VwbEgressAdminController.java
+++ b/api/src/main/java/org/pmiops/workbench/api/VwbEgressAdminController.java
@@ -1,13 +1,23 @@
 package org.pmiops.workbench.api;
 
+import org.pmiops.workbench.exfiltration.EgressEventService;
 import org.pmiops.workbench.model.VwbEgressEventRequest;
+import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.http.ResponseEntity;
 import org.springframework.web.bind.annotation.RestController;
 
 @RestController
 public class VwbEgressAdminController implements VwbEgressAdminApiDelegate {
+  private final EgressEventService egressEventService;
+
+  @Autowired
+  public VwbEgressAdminController(EgressEventService egressEventService) {
+    this.egressEventService = egressEventService;
+  }
+
   @Override
   public ResponseEntity<Void> createVwbEgressEvent(VwbEgressEventRequest body) {
+    egressEventService.handleVwbEvent(body);
     return ResponseEntity.ok().build();
   }
 }
diff --git a/api/src/main/java/org/pmiops/workbench/db/model/DbEgressEvent.java b/api/src/main/java/org/pmiops/workbench/db/model/DbEgressEvent.java
index 902c771dd42..32fe7c17fa5 100644
--- a/api/src/main/java/org/pmiops/workbench/db/model/DbEgressEvent.java
+++ b/api/src/main/java/org/pmiops/workbench/db/model/DbEgressEvent.java
@@ -42,6 +42,10 @@ public enum DbEgressEventStatus {
 
   private String bucketAuditEvent;
 
+  private String vwbWorkspaceId;
+  private String vwbVmName;
+  private boolean isVwb;
+
   @Id
   @GeneratedValue(strategy = GenerationType.IDENTITY)
   @Column(name = "egress_event_id", nullable = false)
@@ -149,6 +153,35 @@ public DbEgressEvent setStatus(DbEgressEventStatus s) {
     return this;
   }
 
+  @Column(name = "vwb_workspace_id")
+  public String getVwbWorkspaceId() {
+    return vwbWorkspaceId;
+  }
+
+  public DbEgressEvent setVwbWorkspaceId(String vwbWorkspaceId) {
+    this.vwbWorkspaceId = vwbWorkspaceId;
+    return this;
+  }
+
+  @Column(name = "vwb_vm_name")
+  public String getVwbVmName() {
+    return vwbVmName;
+  }
+
+  public DbEgressEvent setVwbVmName(String vwbVmName) {
+    this.vwbVmName = vwbVmName;
+    return this;
+  }
+
+  @Column(name = "is_vwb")
+  public boolean getIsVwb() {
+    return isVwb;
+  }
+  public DbEgressEvent setIsVwb(boolean isVwb) {
+    this.isVwb = isVwb;
+    return this;
+  }
+
   public boolean equals(Object o) {
     if (this == o) {
       return true;
@@ -165,7 +198,10 @@ public boolean equals(Object o) {
         && Objects.equals(creationTime, that.creationTime)
         && Objects.equals(lastModifiedTime, that.lastModifiedTime)
         && status == that.status
-        && Objects.equals(sumologicEvent, that.sumologicEvent);
+        && Objects.equals(sumologicEvent, that.sumologicEvent)
+        && Objects.equals(vwbWorkspaceId, that.vwbWorkspaceId)
+        && Objects.equals(vwbVmName, that.vwbVmName)
+        && Objects.equals(isVwb, that.isVwb);
   }
 
   @Override
@@ -179,6 +215,9 @@ public int hashCode() {
         creationTime,
         lastModifiedTime,
         status,
-        sumologicEvent);
+        sumologicEvent,
+        vwbWorkspaceId,
+        vwbVmName,
+        isVwb);
   }
 }
diff --git a/api/src/main/java/org/pmiops/workbench/exfiltration/EgressEventService.java b/api/src/main/java/org/pmiops/workbench/exfiltration/EgressEventService.java
index 81153ada202..0d151eccb03 100644
--- a/api/src/main/java/org/pmiops/workbench/exfiltration/EgressEventService.java
+++ b/api/src/main/java/org/pmiops/workbench/exfiltration/EgressEventService.java
@@ -1,7 +1,9 @@
 package org.pmiops.workbench.exfiltration;
 
 import org.pmiops.workbench.model.SumologicEgressEvent;
+import org.pmiops.workbench.model.VwbEgressEventRequest;
 
 public interface EgressEventService {
   void handleEvent(SumologicEgressEvent egressEvent);
+  void handleVwbEvent(VwbEgressEventRequest egressEvent);
 }
diff --git a/api/src/main/java/org/pmiops/workbench/exfiltration/EgressEventServiceImpl.java b/api/src/main/java/org/pmiops/workbench/exfiltration/EgressEventServiceImpl.java
index 06c4545b534..b2cd8ebd08e 100644
--- a/api/src/main/java/org/pmiops/workbench/exfiltration/EgressEventServiceImpl.java
+++ b/api/src/main/java/org/pmiops/workbench/exfiltration/EgressEventServiceImpl.java
@@ -24,6 +24,7 @@
 import org.pmiops.workbench.db.model.DbUser;
 import org.pmiops.workbench.db.model.DbWorkspace;
 import org.pmiops.workbench.model.SumologicEgressEvent;
+import org.pmiops.workbench.model.VwbEgressEventRequest;
 import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.stereotype.Service;
 
@@ -104,6 +105,30 @@ public void handleEvent(SumologicEgressEvent event) {
     maybeEvent.ifPresent(e -> taskQueueService.pushEgressEventTask(e.getEgressEventId()));
   }
 
+  @Override
+  public void handleVwbEvent(VwbEgressEventRequest egressEvent) {
+    Optional<DbUser> egressUser = userService.getByUsername(egressEvent.getUserEmail());
+    if (egressUser.isEmpty()) {
+      logger.severe(
+          String.format(
+              "An egress event happened for workspace: %s, but we're unable to find a user %s.",
+              egressEvent.getWorkspaceId(), egressEvent.getUserEmail()));
+    }
+    this.egressEventAuditor.fireVwbEgressEvent(egressEvent, egressUser.get());
+    egressEventDao.save(
+        new DbEgressEvent()
+            .setUser(egressUser.orElse(null))
+            .setWorkspace(egressEvent.getWorkspaceId())
+            .setEgressMegabytes(
+                Optional.ofNullable(egressEvent.getEgressMib())
+                    // Mebibytes (2^20 bytes) -> Megabytes (10^6 bytes)
+                    .map(mib -> (float) (mib * ((1 << 20) / 1e6)))
+                    .orElse(null))
+            .setEgressWindowSeconds(
+                Optional.ofNullable(egressEvent.getTimeWindowDuration()).orElse(null))
+            .setStatus(DbEgressEventStatus.PENDING));
+  }
+
   @NotNull
   private Optional<DbUser> findEgressUsers(SumologicEgressEvent event) {
 

From aaf4806f1b126e5013a28b1ffd9baf5ab190773a Mon Sep 17 00:00:00 2001
From: Yonghao Yu <yyhao1@gmail.com>
Date: Mon, 9 Dec 2024 17:38:20 -0500
Subject: [PATCH 02/10] add tests

---
 ...log-246-add-is-vwb-egress-event-column.xml |  2 +
 .../auditors/EgressEventAuditor.java          |  6 +-
 .../auditors/EgressEventAuditorImpl.java      |  5 +-
 .../workbench/db/model/DbEgressEvent.java     | 27 +++++++
 .../exfiltration/EgressEventService.java      |  1 +
 .../exfiltration/EgressEventServiceImpl.java  | 34 +++++----
 .../workbench/db/dao/EgressEventDaoTest.java  | 71 +++++++++++++++++++
 .../exfiltration/EgressEventServiceTest.java  | 39 ++++++++++
 8 files changed, 165 insertions(+), 20 deletions(-)
 create mode 100644 api/src/test/java/org/pmiops/workbench/db/dao/EgressEventDaoTest.java

diff --git a/api/db/changelog/db.changelog-246-add-is-vwb-egress-event-column.xml b/api/db/changelog/db.changelog-246-add-is-vwb-egress-event-column.xml
index 40356fe88a8..2a99b062bc3 100644
--- a/api/db/changelog/db.changelog-246-add-is-vwb-egress-event-column.xml
+++ b/api/db/changelog/db.changelog-246-add-is-vwb-egress-event-column.xml
@@ -9,6 +9,8 @@
             <column name="is_vwb" type="boolean" />
             <column name="vwb_workspace_id" type="varchar(80)" />
             <column name="vwb_vm_name" type="varchar(80)" />
+            <column name="gcp_project_id" type="varchar(80)" />
+            <column name="vwb_incident_count" type="bigint" />
         </addColumn>
     </changeSet>
 </databaseChangeLog>
diff --git a/api/src/main/java/org/pmiops/workbench/actionaudit/auditors/EgressEventAuditor.java b/api/src/main/java/org/pmiops/workbench/actionaudit/auditors/EgressEventAuditor.java
index 3da09118ffd..fbc3b5dee21 100644
--- a/api/src/main/java/org/pmiops/workbench/actionaudit/auditors/EgressEventAuditor.java
+++ b/api/src/main/java/org/pmiops/workbench/actionaudit/auditors/EgressEventAuditor.java
@@ -18,11 +18,13 @@ public interface EgressEventAuditor {
    * event log in the target workspace.
    */
   void fireEgressEvent(SumologicEgressEvent event);
+
   /**
-   * Decorates a Verily Workbench high-egress event with Workbench metadata and fires an audit
-   * event log in the target workspace.
+   * Decorates a Verily Workbench high-egress event with Workbench metadata and fires an audit event
+   * log in the target workspace.
    */
   void fireVwbEgressEvent(VwbEgressEventRequest event, DbUser dbUser);
+
   /**
    * Decorates a Sumologic-reported high-egress event with Workbench metadata and fires an audit
    * event log in the target workspace for the specified user.
diff --git a/api/src/main/java/org/pmiops/workbench/actionaudit/auditors/EgressEventAuditorImpl.java b/api/src/main/java/org/pmiops/workbench/actionaudit/auditors/EgressEventAuditorImpl.java
index fb3a2519588..23d62543281 100644
--- a/api/src/main/java/org/pmiops/workbench/actionaudit/auditors/EgressEventAuditorImpl.java
+++ b/api/src/main/java/org/pmiops/workbench/actionaudit/auditors/EgressEventAuditorImpl.java
@@ -135,10 +135,7 @@ public void fireVwbEgressEvent(VwbEgressEventRequest event, DbUser user) {
             .newValueMaybe(event.getGcpProjectId())
             .build());
     events.add(
-        baseEventBuilder
-            .targetPropertyMaybe("vmName")
-            .newValueMaybe(event.getVmName())
-            .build());
+        baseEventBuilder.targetPropertyMaybe("vmName").newValueMaybe(event.getVmName()).build());
     actionAuditService.send(events);
   }
 
diff --git a/api/src/main/java/org/pmiops/workbench/db/model/DbEgressEvent.java b/api/src/main/java/org/pmiops/workbench/db/model/DbEgressEvent.java
index 32fe7c17fa5..22f0bbf1f6c 100644
--- a/api/src/main/java/org/pmiops/workbench/db/model/DbEgressEvent.java
+++ b/api/src/main/java/org/pmiops/workbench/db/model/DbEgressEvent.java
@@ -44,7 +44,9 @@ public enum DbEgressEventStatus {
 
   private String vwbWorkspaceId;
   private String vwbVmName;
+  private String gcpProjectId;
   private boolean isVwb;
+  private int vwbIncidentCount;
 
   @Id
   @GeneratedValue(strategy = GenerationType.IDENTITY)
@@ -173,10 +175,31 @@ public DbEgressEvent setVwbVmName(String vwbVmName) {
     return this;
   }
 
+  @Column(name = "gcp_project_id")
+  public String getGcpProjectId() {
+    return gcpProjectId;
+  }
+
+  public DbEgressEvent setGcpProjectId(String gcpProjectId) {
+    this.vwbVmName = gcpProjectId;
+    return this;
+  }
+
+  @Column(name = "vwb_incident_count")
+  public int getVwbIncidentCount() {
+    return vwbIncidentCount;
+  }
+
+  public DbEgressEvent setVwbIncidentCount(int vwbIncidentCount) {
+    this.vwbIncidentCount = vwbIncidentCount;
+    return this;
+  }
+
   @Column(name = "is_vwb")
   public boolean getIsVwb() {
     return isVwb;
   }
+
   public DbEgressEvent setIsVwb(boolean isVwb) {
     this.isVwb = isVwb;
     return this;
@@ -201,6 +224,8 @@ public boolean equals(Object o) {
         && Objects.equals(sumologicEvent, that.sumologicEvent)
         && Objects.equals(vwbWorkspaceId, that.vwbWorkspaceId)
         && Objects.equals(vwbVmName, that.vwbVmName)
+        && Objects.equals(gcpProjectId, that.gcpProjectId)
+        && Objects.equals(vwbIncidentCount, that.vwbIncidentCount)
         && Objects.equals(isVwb, that.isVwb);
   }
 
@@ -218,6 +243,8 @@ public int hashCode() {
         sumologicEvent,
         vwbWorkspaceId,
         vwbVmName,
+        gcpProjectId,
+        vwbIncidentCount,
         isVwb);
   }
 }
diff --git a/api/src/main/java/org/pmiops/workbench/exfiltration/EgressEventService.java b/api/src/main/java/org/pmiops/workbench/exfiltration/EgressEventService.java
index 0d151eccb03..9161515a795 100644
--- a/api/src/main/java/org/pmiops/workbench/exfiltration/EgressEventService.java
+++ b/api/src/main/java/org/pmiops/workbench/exfiltration/EgressEventService.java
@@ -5,5 +5,6 @@
 
 public interface EgressEventService {
   void handleEvent(SumologicEgressEvent egressEvent);
+
   void handleVwbEvent(VwbEgressEventRequest egressEvent);
 }
diff --git a/api/src/main/java/org/pmiops/workbench/exfiltration/EgressEventServiceImpl.java b/api/src/main/java/org/pmiops/workbench/exfiltration/EgressEventServiceImpl.java
index b2cd8ebd08e..e0c78ce78bb 100644
--- a/api/src/main/java/org/pmiops/workbench/exfiltration/EgressEventServiceImpl.java
+++ b/api/src/main/java/org/pmiops/workbench/exfiltration/EgressEventServiceImpl.java
@@ -101,7 +101,7 @@ public void handleEvent(SumologicEgressEvent event) {
     this.egressEventAuditor.fireEgressEventForUser(event, egressUser.get());
 
     Optional<DbEgressEvent> maybeEvent =
-        this.maybePersistEgressEvent(event, Optional.of(egressUser.get()), dbWorkspaceMaybe);
+        this.maybePersistSumoEgressEvent(event, Optional.of(egressUser.get()), dbWorkspaceMaybe);
     maybeEvent.ifPresent(e -> taskQueueService.pushEgressEventTask(e.getEgressEventId()));
   }
 
@@ -115,18 +115,23 @@ public void handleVwbEvent(VwbEgressEventRequest egressEvent) {
               egressEvent.getWorkspaceId(), egressEvent.getUserEmail()));
     }
     this.egressEventAuditor.fireVwbEgressEvent(egressEvent, egressUser.get());
-    egressEventDao.save(
-        new DbEgressEvent()
-            .setUser(egressUser.orElse(null))
-            .setWorkspace(egressEvent.getWorkspaceId())
-            .setEgressMegabytes(
-                Optional.ofNullable(egressEvent.getEgressMib())
-                    // Mebibytes (2^20 bytes) -> Megabytes (10^6 bytes)
-                    .map(mib -> (float) (mib * ((1 << 20) / 1e6)))
-                    .orElse(null))
-            .setEgressWindowSeconds(
-                Optional.ofNullable(egressEvent.getTimeWindowDuration()).orElse(null))
-            .setStatus(DbEgressEventStatus.PENDING));
+    DbEgressEvent event =
+        egressEventDao.save(
+            new DbEgressEvent()
+                .setUser(egressUser.orElse(null))
+                .setIsVwb(true)
+                .setVwbWorkspaceId(egressEvent.getWorkspaceId())
+                .setVwbVmName(egressEvent.getVmName())
+                .setVwbIncidentCount(egressEvent.getIncidentCount().intValue())
+                .setEgressMegabytes(
+                    Optional.ofNullable(egressEvent.getEgressMib())
+                        // Mebibytes (2^20 bytes) -> Megabytes (10^6 bytes)
+                        .map(mib -> (float) (mib * ((1 << 20) / 1e6)))
+                        .orElse(null))
+                .setEgressWindowSeconds(
+                    Optional.ofNullable(egressEvent.getTimeWindowDuration()).orElse(null))
+                .setStatus(DbEgressEventStatus.PENDING));
+    taskQueueService.pushEgressEventTask(event.getEgressEventId());
   }
 
   @NotNull
@@ -164,7 +169,7 @@ private boolean isEventStale(SumologicEgressEvent egressEvent) {
     return windowStart.isBefore(clock.instant().minus(windowDuration.multipliedBy(2L)));
   }
 
-  private Optional<DbEgressEvent> maybePersistEgressEvent(
+  private Optional<DbEgressEvent> maybePersistSumoEgressEvent(
       SumologicEgressEvent event,
       Optional<DbUser> userMaybe,
       Optional<DbWorkspace> workspaceMaybe) {
@@ -176,6 +181,7 @@ private Optional<DbEgressEvent> maybePersistEgressEvent(
         egressEventDao.save(
             new DbEgressEvent()
                 .setUser(userMaybe.orElse(null))
+                .setIsVwb(false)
                 .setWorkspace(workspaceMaybe.orElse(null))
                 .setEgressMegabytes(
                     Optional.ofNullable(event.getEgressMib())
diff --git a/api/src/test/java/org/pmiops/workbench/db/dao/EgressEventDaoTest.java b/api/src/test/java/org/pmiops/workbench/db/dao/EgressEventDaoTest.java
new file mode 100644
index 00000000000..ba3d4fd1568
--- /dev/null
+++ b/api/src/test/java/org/pmiops/workbench/db/dao/EgressEventDaoTest.java
@@ -0,0 +1,71 @@
+package org.pmiops.workbench.db.dao;
+
+import static com.google.common.truth.Truth.assertThat;
+
+import java.sql.Timestamp;
+import java.time.Instant;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+import org.pmiops.workbench.FakeClockConfiguration;
+import org.pmiops.workbench.config.CommonConfig;
+import org.pmiops.workbench.db.model.DbEgressEvent;
+import org.pmiops.workbench.db.model.DbEgressEvent.DbEgressEventStatus;
+import org.pmiops.workbench.db.model.DbUser;
+import org.pmiops.workbench.db.model.DbWorkspace;
+import org.springframework.beans.factory.annotation.Autowired;
+import org.springframework.boot.test.autoconfigure.orm.jpa.DataJpaTest;
+import org.springframework.context.annotation.Import;
+
+@Import({FakeClockConfiguration.class, CommonConfig.class})
+@DataJpaTest
+public class EgressEventDaoTest {
+  @Autowired private UserDao userDao;
+
+  @Autowired private EgressEventDao egressEventDao;
+
+  private static final Timestamp CURRENT_TIMESTAMP = Timestamp.from(Instant.now());
+  private DbUser user;
+
+  @BeforeEach
+  public void setUp() {
+    user = userDao.save(new DbUser());
+  }
+
+  @Test
+  public void testCreateAndGet() {
+    // Create
+    DbEgressEvent event = createValidDbEgressEvent();
+    event = egressEventDao.save(event);
+    assertThat(event.getEgressEventId()).isNotNull();
+
+    // Read
+    DbEgressEvent foundEvent = egressEventDao.findById(event.getEgressEventId()).orElse(null);
+    assertThat(foundEvent).isNotNull();
+    assertThat(foundEvent.getUser()).isEqualTo(user);
+    assertThat(foundEvent.getVwbWorkspaceId()).isEqualTo("testWorkspaceId");
+    assertThat(foundEvent.getVwbVmName()).isEqualTo("testVmName");
+    assertThat(foundEvent.getGcpProjectId()).isEqualTo("test-project");
+    assertThat(foundEvent.getIsVwb()).isTrue();
+
+    // Update
+    foundEvent.setStatus(DbEgressEventStatus.VERIFIED_FALSE_POSITIVE);
+    egressEventDao.save(foundEvent);
+    DbEgressEvent updatedEvent = egressEventDao.findById(event.getEgressEventId()).orElse(null);
+    assertThat(updatedEvent.getStatus())
+        .isEqualTo(DbEgressEvent.DbEgressEventStatus.VERIFIED_FALSE_POSITIVE);
+  }
+
+  private DbEgressEvent createValidDbEgressEvent() {
+    DbEgressEvent event = new DbEgressEvent();
+    event.setUser(user);
+    event.setWorkspace(new DbWorkspace());
+    event.setCreationTime(CURRENT_TIMESTAMP);
+    event.setStatus(DbEgressEvent.DbEgressEventStatus.PENDING);
+    event.setEgressWindowSeconds(3600L);
+    event.setVwbWorkspaceId("testWorkspaceId");
+    event.setVwbVmName("testVmName");
+    event.setGcpProjectId("test-project");
+    event.setIsVwb(true);
+    return event;
+  }
+}
diff --git a/api/src/test/java/org/pmiops/workbench/exfiltration/EgressEventServiceTest.java b/api/src/test/java/org/pmiops/workbench/exfiltration/EgressEventServiceTest.java
index b033c5f92b9..6bcd84f13a0 100644
--- a/api/src/test/java/org/pmiops/workbench/exfiltration/EgressEventServiceTest.java
+++ b/api/src/test/java/org/pmiops/workbench/exfiltration/EgressEventServiceTest.java
@@ -288,6 +288,45 @@ public void testCreateEgressEventAlert_staleEventShortWindowPersisted() {
     assertThat(dbEvents).hasSize(2);
   }
 
+  @Test
+  public void testHandleVwbEgressEvent() {
+    VwbEgressEventRequest vwbEvent =
+        new VwbEgressEventRequest()
+            .userEmail("testuser@example.com")
+            .workspaceId("testWorkspaceId")
+            .vmName("testVmName")
+            .incidentCount(1L)
+            .egressMib(500.0)
+            .egressMibThreshold(100.0)
+            .gcpProjectId("test-gcp-project")
+            .timeWindowDuration(600L)
+            .timeWindowStart(NOW.toEpochMilli());
+
+    DbUser dbUser = new DbUser();
+    dbUser.setUsername("testuser@example.com");
+    dbUser.setUserId(123L);
+
+    doReturn(Optional.of(dbUser)).when(mockUserService).getByUsername("testuser@example.com");
+
+    egressEventService.handleVwbEvent(vwbEvent);
+
+    verify(mockEgressEventAuditor).fireVwbEgressEvent(vwbEvent, dbUser);
+    verify(mockTaskQueueService).pushEgressEventTask(anyLong());
+
+    List<DbEgressEvent> dbEvents = ImmutableList.copyOf(egressEventDao.findAll());
+    assertThat(dbEvents).hasSize(1);
+    DbEgressEvent dbEvent = Iterables.getOnlyElement(dbEvents);
+    assertThat(dbEvent.getUser()).isEqualTo(dbUser);
+    assertThat(dbEvent.getVwbWorkspaceId()).isEqualTo("testWorkspaceId");
+    assertThat(dbEvent.getVwbVmName()).isEqualTo("testVmName");
+    assertThat(dbEvent.getVwbIncidentCount()).isEqualTo(1);
+    assertThat(dbEvent.getEgressMegabytes()).isEqualTo(500.0);
+    assertThat(dbEvent.getEgressMegabytes()).isEqualTo(500.0 * 1.04858); // Convert MiB
+    assertThat(dbEvent.getGcpProjectId()).isEqualTo("test-gcp-project");
+    assertThat(dbEvent.getEgressWindowSeconds()).isEqualTo(600L);
+    assertThat(dbEvent.getIsVwb()).isTrue();
+  }
+
   private static SumologicEgressEvent recentEgressEventForUser(DbUser user) {
     return new SumologicEgressEvent()
         .projectName(DEFAULT_GOOGLE_PROJECT)

From edd7ab00eb0f470748fcf73e907df10e868c5ffb Mon Sep 17 00:00:00 2001
From: Yonghao Yu <yyhao1@gmail.com>
Date: Mon, 9 Dec 2024 18:00:18 -0500
Subject: [PATCH 03/10] fix test

---
 .../api/VwbEgressAdminController.java         | 27 +++++-
 .../workbench/db/model/DbEgressEvent.java     |  2 +-
 .../api/VwbEgressAdminControllerTest.java     | 86 +++++++++++++++++++
 .../exfiltration/EgressEventServiceTest.java  | 12 +--
 4 files changed, 117 insertions(+), 10 deletions(-)
 create mode 100644 api/src/test/java/org/pmiops/workbench/api/VwbEgressAdminControllerTest.java

diff --git a/api/src/main/java/org/pmiops/workbench/api/VwbEgressAdminController.java b/api/src/main/java/org/pmiops/workbench/api/VwbEgressAdminController.java
index 753fc035585..0a0458cb337 100644
--- a/api/src/main/java/org/pmiops/workbench/api/VwbEgressAdminController.java
+++ b/api/src/main/java/org/pmiops/workbench/api/VwbEgressAdminController.java
@@ -1,5 +1,10 @@
 package org.pmiops.workbench.api;
 
+import jakarta.inject.Provider;
+import java.util.logging.Logger;
+import org.pmiops.workbench.config.WorkbenchConfig;
+import org.pmiops.workbench.db.model.DbUser;
+import org.pmiops.workbench.exceptions.ForbiddenException;
 import org.pmiops.workbench.exfiltration.EgressEventService;
 import org.pmiops.workbench.model.VwbEgressEventRequest;
 import org.springframework.beans.factory.annotation.Autowired;
@@ -8,15 +13,35 @@
 
 @RestController
 public class VwbEgressAdminController implements VwbEgressAdminApiDelegate {
+  private static final Logger log = Logger.getLogger(NotebooksController.class.getName());
   private final EgressEventService egressEventService;
 
+  private final Provider<WorkbenchConfig> workbenchConfigProvider;
+
+  private final Provider<DbUser> userProvider;
+
   @Autowired
-  public VwbEgressAdminController(EgressEventService egressEventService) {
+  public VwbEgressAdminController(
+      EgressEventService egressEventService,
+      Provider<WorkbenchConfig> workbenchConfigProvider,
+      Provider<DbUser> userProvider) {
     this.egressEventService = egressEventService;
+    this.workbenchConfigProvider = workbenchConfigProvider;
+    this.userProvider = userProvider;
   }
 
   @Override
   public ResponseEntity<Void> createVwbEgressEvent(VwbEgressEventRequest body) {
+    if (!workbenchConfigProvider
+        .get()
+        .vwb
+        .exfilManagerServiceAccount
+        .equals(userProvider.get().getUsername())) {
+      throw new ForbiddenException(
+          String.format(
+              "User %s is not authorized to create VwbEgressEvent",
+              userProvider.get().getUsername()));
+    }
     egressEventService.handleVwbEvent(body);
     return ResponseEntity.ok().build();
   }
diff --git a/api/src/main/java/org/pmiops/workbench/db/model/DbEgressEvent.java b/api/src/main/java/org/pmiops/workbench/db/model/DbEgressEvent.java
index 22f0bbf1f6c..0e2fe5494d1 100644
--- a/api/src/main/java/org/pmiops/workbench/db/model/DbEgressEvent.java
+++ b/api/src/main/java/org/pmiops/workbench/db/model/DbEgressEvent.java
@@ -181,7 +181,7 @@ public String getGcpProjectId() {
   }
 
   public DbEgressEvent setGcpProjectId(String gcpProjectId) {
-    this.vwbVmName = gcpProjectId;
+    this.gcpProjectId = gcpProjectId;
     return this;
   }
 
diff --git a/api/src/test/java/org/pmiops/workbench/api/VwbEgressAdminControllerTest.java b/api/src/test/java/org/pmiops/workbench/api/VwbEgressAdminControllerTest.java
new file mode 100644
index 00000000000..d13a46c2428
--- /dev/null
+++ b/api/src/test/java/org/pmiops/workbench/api/VwbEgressAdminControllerTest.java
@@ -0,0 +1,86 @@
+package org.pmiops.workbench.api;
+
+import static com.google.common.truth.Truth.assertThat;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+
+import jakarta.inject.Provider;
+import java.time.Instant;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+import org.mockito.Mock;
+import org.pmiops.workbench.FakeClockConfiguration;
+import org.pmiops.workbench.config.WorkbenchConfig;
+import org.pmiops.workbench.db.model.DbUser;
+import org.pmiops.workbench.exceptions.ForbiddenException;
+import org.pmiops.workbench.exfiltration.EgressEventService;
+import org.pmiops.workbench.model.VwbEgressEventRequest;
+import org.springframework.boot.test.autoconfigure.orm.jpa.DataJpaTest;
+import org.springframework.context.annotation.Import;
+import org.springframework.http.HttpStatus;
+import org.springframework.http.ResponseEntity;
+import org.springframework.test.annotation.DirtiesContext;
+
+@DataJpaTest
+@DirtiesContext(classMode = DirtiesContext.ClassMode.BEFORE_EACH_TEST_METHOD)
+@Import(FakeClockConfiguration.class)
+public class VwbEgressAdminControllerTest {
+  @Mock private Provider<DbUser> mockUserProvider;
+  @Mock private Provider<WorkbenchConfig> mockConfigProvider;
+  @Mock private EgressEventService mockEgressEventService;
+
+  private VwbEgressAdminController vwbEgressAdminController;
+
+  private WorkbenchConfig workbenchConfig = new WorkbenchConfig();
+  private VwbEgressEventRequest vwbEvent;
+
+  @BeforeEach
+  public void setUp() {
+    this.vwbEgressAdminController =
+        new VwbEgressAdminController(mockEgressEventService, mockConfigProvider, mockUserProvider);
+    workbenchConfig.vwb = new WorkbenchConfig.VwbConfig();
+    workbenchConfig.vwb.exfilManagerServiceAccount = "authorized-user@example.com";
+
+    vwbEvent =
+        new VwbEgressEventRequest()
+            .userEmail("testuser@example.com")
+            .workspaceId("testWorkspaceId")
+            .vmName("testVmName")
+            .incidentCount(1L)
+            .egressMib(500.0)
+            .egressMibThreshold(100.0)
+            .gcpProjectId("test-gcp-project")
+            .timeWindowDuration(600L)
+            .timeWindowStart(Instant.now().toEpochMilli());
+  }
+
+  @Test
+  public void testCreateVwbEgressEvent() {
+    DbUser dbUser = new DbUser();
+    dbUser.setUsername("test-user@example.com");
+
+    when(mockConfigProvider.get()).thenReturn(workbenchConfig);
+    when(mockUserProvider.get()).thenReturn(dbUser);
+
+    ResponseEntity<Void> response = vwbEgressAdminController.createVwbEgressEvent(vwbEvent);
+
+    verify(mockEgressEventService).handleVwbEvent(vwbEvent);
+    assertThat(response.getStatusCode()).isEqualTo(HttpStatus.OK);
+  }
+
+  @Test
+  public void testCreateVwbEgressEvent_unauthorizedUser() {
+    DbUser dbUser = new DbUser();
+    dbUser.setUsername("unauthorized-user@example.com");
+
+    when(mockConfigProvider.get()).thenReturn(workbenchConfig);
+    when(mockUserProvider.get()).thenReturn(dbUser);
+
+    assertThrows(
+        ForbiddenException.class,
+        () -> {
+          vwbEgressAdminController.createVwbEgressEvent(vwbEvent);
+        });
+  }
+}
diff --git a/api/src/test/java/org/pmiops/workbench/exfiltration/EgressEventServiceTest.java b/api/src/test/java/org/pmiops/workbench/exfiltration/EgressEventServiceTest.java
index 6bcd84f13a0..f16d474145b 100644
--- a/api/src/test/java/org/pmiops/workbench/exfiltration/EgressEventServiceTest.java
+++ b/api/src/test/java/org/pmiops/workbench/exfiltration/EgressEventServiceTest.java
@@ -292,7 +292,7 @@ public void testCreateEgressEventAlert_staleEventShortWindowPersisted() {
   public void testHandleVwbEgressEvent() {
     VwbEgressEventRequest vwbEvent =
         new VwbEgressEventRequest()
-            .userEmail("testuser@example.com")
+            .userEmail(dbUser1.getUsername())
             .workspaceId("testWorkspaceId")
             .vmName("testVmName")
             .incidentCount(1L)
@@ -302,21 +302,17 @@ public void testHandleVwbEgressEvent() {
             .timeWindowDuration(600L)
             .timeWindowStart(NOW.toEpochMilli());
 
-    DbUser dbUser = new DbUser();
-    dbUser.setUsername("testuser@example.com");
-    dbUser.setUserId(123L);
-
-    doReturn(Optional.of(dbUser)).when(mockUserService).getByUsername("testuser@example.com");
+    doReturn(Optional.of(dbUser1)).when(mockUserService).getByDatabaseId(dbUser1.getUserId());
 
     egressEventService.handleVwbEvent(vwbEvent);
 
-    verify(mockEgressEventAuditor).fireVwbEgressEvent(vwbEvent, dbUser);
+    verify(mockEgressEventAuditor).fireVwbEgressEvent(vwbEvent, dbUser1);
     verify(mockTaskQueueService).pushEgressEventTask(anyLong());
 
     List<DbEgressEvent> dbEvents = ImmutableList.copyOf(egressEventDao.findAll());
     assertThat(dbEvents).hasSize(1);
     DbEgressEvent dbEvent = Iterables.getOnlyElement(dbEvents);
-    assertThat(dbEvent.getUser()).isEqualTo(dbUser);
+    assertThat(dbEvent.getUser()).isEqualTo(dbUser1);
     assertThat(dbEvent.getVwbWorkspaceId()).isEqualTo("testWorkspaceId");
     assertThat(dbEvent.getVwbVmName()).isEqualTo("testVmName");
     assertThat(dbEvent.getVwbIncidentCount()).isEqualTo(1);

From 397444d312bac984a5d0702649fc24ae8fdbda3d Mon Sep 17 00:00:00 2001
From: Yonghao Yu <yyhao1@gmail.com>
Date: Tue, 10 Dec 2024 10:00:31 -0500
Subject: [PATCH 04/10] fix test

---
 .../workbench/api/VwbEgressAdminControllerTest.java | 13 ++++++-------
 .../exfiltration/EgressEventServiceTest.java        |  3 +--
 2 files changed, 7 insertions(+), 9 deletions(-)

diff --git a/api/src/test/java/org/pmiops/workbench/api/VwbEgressAdminControllerTest.java b/api/src/test/java/org/pmiops/workbench/api/VwbEgressAdminControllerTest.java
index d13a46c2428..d4fc8a1f547 100644
--- a/api/src/test/java/org/pmiops/workbench/api/VwbEgressAdminControllerTest.java
+++ b/api/src/test/java/org/pmiops/workbench/api/VwbEgressAdminControllerTest.java
@@ -34,13 +34,17 @@ public class VwbEgressAdminControllerTest {
 
   private WorkbenchConfig workbenchConfig = new WorkbenchConfig();
   private VwbEgressEventRequest vwbEvent;
+  private DbUser authorizedUser;
 
   @BeforeEach
   public void setUp() {
     this.vwbEgressAdminController =
         new VwbEgressAdminController(mockEgressEventService, mockConfigProvider, mockUserProvider);
+    authorizedUser = new DbUser();
+    authorizedUser.setUsername("test-user@example.com");
     workbenchConfig.vwb = new WorkbenchConfig.VwbConfig();
-    workbenchConfig.vwb.exfilManagerServiceAccount = "authorized-user@example.com";
+    workbenchConfig.vwb.exfilManagerServiceAccount = authorizedUser.getUsername();
+    when(mockConfigProvider.get()).thenReturn(workbenchConfig);
 
     vwbEvent =
         new VwbEgressEventRequest()
@@ -57,11 +61,7 @@ public void setUp() {
 
   @Test
   public void testCreateVwbEgressEvent() {
-    DbUser dbUser = new DbUser();
-    dbUser.setUsername("test-user@example.com");
-
-    when(mockConfigProvider.get()).thenReturn(workbenchConfig);
-    when(mockUserProvider.get()).thenReturn(dbUser);
+    when(mockUserProvider.get()).thenReturn(authorizedUser);
 
     ResponseEntity<Void> response = vwbEgressAdminController.createVwbEgressEvent(vwbEvent);
 
@@ -74,7 +74,6 @@ public void testCreateVwbEgressEvent_unauthorizedUser() {
     DbUser dbUser = new DbUser();
     dbUser.setUsername("unauthorized-user@example.com");
 
-    when(mockConfigProvider.get()).thenReturn(workbenchConfig);
     when(mockUserProvider.get()).thenReturn(dbUser);
 
     assertThrows(
diff --git a/api/src/test/java/org/pmiops/workbench/exfiltration/EgressEventServiceTest.java b/api/src/test/java/org/pmiops/workbench/exfiltration/EgressEventServiceTest.java
index f16d474145b..e11a8f32121 100644
--- a/api/src/test/java/org/pmiops/workbench/exfiltration/EgressEventServiceTest.java
+++ b/api/src/test/java/org/pmiops/workbench/exfiltration/EgressEventServiceTest.java
@@ -316,8 +316,7 @@ public void testHandleVwbEgressEvent() {
     assertThat(dbEvent.getVwbWorkspaceId()).isEqualTo("testWorkspaceId");
     assertThat(dbEvent.getVwbVmName()).isEqualTo("testVmName");
     assertThat(dbEvent.getVwbIncidentCount()).isEqualTo(1);
-    assertThat(dbEvent.getEgressMegabytes()).isEqualTo(500.0);
-    assertThat(dbEvent.getEgressMegabytes()).isEqualTo(500.0 * 1.04858); // Convert MiB
+    assertThat(dbEvent.getEgressMegabytes()).isEqualTo(524.288); // Convert MiB
     assertThat(dbEvent.getGcpProjectId()).isEqualTo("test-gcp-project");
     assertThat(dbEvent.getEgressWindowSeconds()).isEqualTo(600L);
     assertThat(dbEvent.getIsVwb()).isTrue();

From 59f37c15fa7a14fac3f46010d2ab716cb5c8baa4 Mon Sep 17 00:00:00 2001
From: Yonghao Yu <yyhao1@gmail.com>
Date: Mon, 16 Dec 2024 10:05:15 -0500
Subject: [PATCH 05/10] fix test

---
 ...elog-246-add-vwb-egress-event-columns.xml} |   1 +
 api/db/changelog/db.changelog-master.xml      |   2 +-
 .../workbench/db/dao/EgressEventDao.java      |   3 +-
 .../workbench/db/model/DbEgressEvent.java     |  11 ++
 .../exfiltration/EgressEventServiceImpl.java  |   3 +-
 .../EgressRemediationService.java             |   5 +-
 ...EgressObjectLengthsRemediationService.java |   5 +
 .../EgressSumologicRemediationService.java    |   6 +
 .../impl/EgressVwbRemediationService.java     | 112 ++++++++++++++++++
 api/src/main/resources/workbench-api.yaml     |   5 +-
 .../workbench/db/dao/EgressEventDaoTest.java  |  22 ++++
 11 files changed, 170 insertions(+), 5 deletions(-)
 rename api/db/changelog/{db.changelog-246-add-is-vwb-egress-event-column.xml => db.changelog-246-add-vwb-egress-event-columns.xml} (92%)
 create mode 100644 api/src/main/java/org/pmiops/workbench/exfiltration/impl/EgressVwbRemediationService.java

diff --git a/api/db/changelog/db.changelog-246-add-is-vwb-egress-event-column.xml b/api/db/changelog/db.changelog-246-add-vwb-egress-event-columns.xml
similarity index 92%
rename from api/db/changelog/db.changelog-246-add-is-vwb-egress-event-column.xml
rename to api/db/changelog/db.changelog-246-add-vwb-egress-event-columns.xml
index 2a99b062bc3..617ae7e062b 100644
--- a/api/db/changelog/db.changelog-246-add-is-vwb-egress-event-column.xml
+++ b/api/db/changelog/db.changelog-246-add-vwb-egress-event-columns.xml
@@ -11,6 +11,7 @@
             <column name="vwb_vm_name" type="varchar(80)" />
             <column name="gcp_project_id" type="varchar(80)" />
             <column name="vwb_incident_count" type="bigint" />
+            <column name="vwb_egress_event_id" type="varchar(80)" />
         </addColumn>
     </changeSet>
 </databaseChangeLog>
diff --git a/api/db/changelog/db.changelog-master.xml b/api/db/changelog/db.changelog-master.xml
index 2fd01988180..9738c248791 100644
--- a/api/db/changelog/db.changelog-master.xml
+++ b/api/db/changelog/db.changelog-master.xml
@@ -253,7 +253,7 @@
   <include file="changelog/db.changelog-243-replace-extension-count.xml"/>
   <include file="changelog/db.changelog-244-add-is-vwb-workspace-column.xml"/>
   <include file="changelog/db.changelog-245-add-vwb-template-id-cdr-version.xml"/>
-  <include file="changelog/db.changelog-246-add-is-vwb-egress-event-column.xml"/>
+  <include file="changelog/db.db.changelog-246-add-vwb-egress-event-columns.xml"/>
   <!--
    Note: to update the DB locally, do the following:
    - Migrate schema changes: `./project.rb run-local-all-migrations`
diff --git a/api/src/main/java/org/pmiops/workbench/db/dao/EgressEventDao.java b/api/src/main/java/org/pmiops/workbench/db/dao/EgressEventDao.java
index 6bda388cf77..4b469d338c8 100644
--- a/api/src/main/java/org/pmiops/workbench/db/dao/EgressEventDao.java
+++ b/api/src/main/java/org/pmiops/workbench/db/dao/EgressEventDao.java
@@ -1,6 +1,7 @@
 package org.pmiops.workbench.db.dao;
 
 import java.sql.Timestamp;
+import java.util.Collection;
 import java.util.List;
 import org.pmiops.workbench.db.model.DbEgressEvent;
 import org.pmiops.workbench.db.model.DbEgressEvent.DbEgressEventStatus;
@@ -23,7 +24,7 @@ public interface EgressEventDao
   Page<DbEgressEvent> findAllByUserAndWorkspaceOrderByCreationTimeDesc(
       DbUser user, DbWorkspace workspace, Pageable p);
 
-  List<DbEgressEvent> findAllByUserAndStatusNot(DbUser user, DbEgressEventStatus status);
+  List<DbEgressEvent> findAllByUserAndStatusNotIn(DbUser user, Collection<DbEgressEventStatus> status);
 
   List<DbEgressEvent> findAllByUserAndWorkspaceAndCreationTimeBetweenAndCreationTimeNot(
       DbUser user,
diff --git a/api/src/main/java/org/pmiops/workbench/db/model/DbEgressEvent.java b/api/src/main/java/org/pmiops/workbench/db/model/DbEgressEvent.java
index 0e2fe5494d1..277b7688d43 100644
--- a/api/src/main/java/org/pmiops/workbench/db/model/DbEgressEvent.java
+++ b/api/src/main/java/org/pmiops/workbench/db/model/DbEgressEvent.java
@@ -45,6 +45,7 @@ public enum DbEgressEventStatus {
   private String vwbWorkspaceId;
   private String vwbVmName;
   private String gcpProjectId;
+  private String vwbEgressEventId;
   private boolean isVwb;
   private int vwbIncidentCount;
 
@@ -185,6 +186,16 @@ public DbEgressEvent setGcpProjectId(String gcpProjectId) {
     return this;
   }
 
+  @Column(name = "vwb_egress_event_id")
+  public String getVwbEgressEventId() {
+    return vwbEgressEventId;
+  }
+
+  public DbEgressEvent setVwbEgressEventId(String vwbEgressEventId) {
+    this.vwbEgressEventId = vwbEgressEventId;
+    return this;
+  }
+
   @Column(name = "vwb_incident_count")
   public int getVwbIncidentCount() {
     return vwbIncidentCount;
diff --git a/api/src/main/java/org/pmiops/workbench/exfiltration/EgressEventServiceImpl.java b/api/src/main/java/org/pmiops/workbench/exfiltration/EgressEventServiceImpl.java
index e0c78ce78bb..06f976de81b 100644
--- a/api/src/main/java/org/pmiops/workbench/exfiltration/EgressEventServiceImpl.java
+++ b/api/src/main/java/org/pmiops/workbench/exfiltration/EgressEventServiceImpl.java
@@ -120,7 +120,8 @@ public void handleVwbEvent(VwbEgressEventRequest egressEvent) {
             new DbEgressEvent()
                 .setUser(egressUser.orElse(null))
                 .setIsVwb(true)
-                .setVwbWorkspaceId(egressEvent.getWorkspaceId())
+                .setVwbEgressEventId(egressEvent.getVwbEgressEventId())
+                .setVwbWorkspaceId(egressEvent.getVwbWorkspaceId())
                 .setVwbVmName(egressEvent.getVmName())
                 .setVwbIncidentCount(egressEvent.getIncidentCount().intValue())
                 .setEgressMegabytes(
diff --git a/api/src/main/java/org/pmiops/workbench/exfiltration/EgressRemediationService.java b/api/src/main/java/org/pmiops/workbench/exfiltration/EgressRemediationService.java
index 14d525bb020..292815bf3c0 100644
--- a/api/src/main/java/org/pmiops/workbench/exfiltration/EgressRemediationService.java
+++ b/api/src/main/java/org/pmiops/workbench/exfiltration/EgressRemediationService.java
@@ -170,7 +170,7 @@ protected abstract void logEvent(DbEgressEvent event, EgressRemediationAction ac
    */
   private int getEgressIncidentCountForUser(DbUser user) {
     List<DbEgressEvent> events =
-        egressEventDao.findAllByUserAndStatusNot(user, DbEgressEventStatus.VERIFIED_FALSE_POSITIVE);
+        egressEventDao.findAllByUserAndStatusNotIn(user, List.of(DbEgressEventStatus.VERIFIED_FALSE_POSITIVE, DbEgressEventStatus.BYPASSED));
 
     // If any egress alerts are missing workspace metadata (this should not happen), consider each
     // as unique incidents.
@@ -288,4 +288,7 @@ private void stopUserRuntimesAndApps(String userEmail) {
             "stopped %d runtimes and %d apps for user %s",
             stoppedRuntimeCount, stoppedAppCount, userEmail));
   }
+
+  protected abstract void getEgressIncidentCountForUser(DbEgressEvent event,
+      EgressRemediationAction action) throws ApiException;
 }
diff --git a/api/src/main/java/org/pmiops/workbench/exfiltration/impl/EgressObjectLengthsRemediationService.java b/api/src/main/java/org/pmiops/workbench/exfiltration/impl/EgressObjectLengthsRemediationService.java
index 5b68c0ca1b1..140474b6cc5 100644
--- a/api/src/main/java/org/pmiops/workbench/exfiltration/impl/EgressObjectLengthsRemediationService.java
+++ b/api/src/main/java/org/pmiops/workbench/exfiltration/impl/EgressObjectLengthsRemediationService.java
@@ -72,4 +72,9 @@ protected void logEvent(DbEgressEvent event, EgressRemediationAction action) thr
   protected boolean shouldSkipEgressEvent(DbEgressEvent event) {
     return false;
   }
+
+  @Override
+  protected void getEgressIncidentCountForUser(DbEgressEvent event, EgressRemediationAction action)
+      throws ApiException {
+  }
 }
diff --git a/api/src/main/java/org/pmiops/workbench/exfiltration/impl/EgressSumologicRemediationService.java b/api/src/main/java/org/pmiops/workbench/exfiltration/impl/EgressSumologicRemediationService.java
index 3112d6978d7..c77f5d243f2 100644
--- a/api/src/main/java/org/pmiops/workbench/exfiltration/impl/EgressSumologicRemediationService.java
+++ b/api/src/main/java/org/pmiops/workbench/exfiltration/impl/EgressSumologicRemediationService.java
@@ -108,6 +108,12 @@ protected boolean shouldSkipEgressEvent(DbEgressEvent event) {
     return false;
   }
 
+  @Override
+  protected void getEgressIncidentCountForUser(DbEgressEvent event, EgressRemediationAction action)
+      throws ApiException {
+
+  }
+
   private boolean isUserBypassedForLargeFileDownload(DbEgressEvent event) {
     return userAdminService.getCurrentEgressBypassWindow(event.getUser().getUserId()) != null
         && event.getEgressMegabytes() != null
diff --git a/api/src/main/java/org/pmiops/workbench/exfiltration/impl/EgressVwbRemediationService.java b/api/src/main/java/org/pmiops/workbench/exfiltration/impl/EgressVwbRemediationService.java
new file mode 100644
index 00000000000..690047aab11
--- /dev/null
+++ b/api/src/main/java/org/pmiops/workbench/exfiltration/impl/EgressVwbRemediationService.java
@@ -0,0 +1,112 @@
+package org.pmiops.workbench.exfiltration.impl;
+
+import static org.pmiops.workbench.exfiltration.ExfiltrationUtils.SUMOLOGIC_JIRA_HANDLER_QUALIFIER;
+import static org.pmiops.workbench.leonardo.LeonardoAppUtils.appServiceNameToAppType;
+
+import com.google.common.base.Strings;
+import jakarta.inject.Provider;
+import jakarta.mail.MessagingException;
+import java.time.Clock;
+import java.util.logging.Logger;
+import org.pmiops.workbench.actionaudit.auditors.EgressEventAuditor;
+import org.pmiops.workbench.config.WorkbenchConfig;
+import org.pmiops.workbench.db.dao.EgressEventDao;
+import org.pmiops.workbench.db.dao.UserService;
+import org.pmiops.workbench.db.model.DbEgressEvent;
+import org.pmiops.workbench.db.model.DbUser;
+import org.pmiops.workbench.exfiltration.EgressRemediationAction;
+import org.pmiops.workbench.exfiltration.EgressRemediationService;
+import org.pmiops.workbench.exfiltration.ExfiltrationUtils;
+import org.pmiops.workbench.exfiltration.jirahandler.EgressJiraHandler;
+import org.pmiops.workbench.jira.ApiException;
+import org.pmiops.workbench.leonardo.LeonardoApiClient;
+import org.pmiops.workbench.leonardo.LeonardoAppUtils;
+import org.pmiops.workbench.mail.MailService;
+import org.pmiops.workbench.model.SumologicEgressEvent;
+import org.pmiops.workbench.user.UserAdminService;
+import org.pmiops.workbench.utils.mappers.EgressEventMapper;
+import org.springframework.beans.factory.annotation.Autowired;
+import org.springframework.beans.factory.annotation.Qualifier;
+import org.springframework.stereotype.Service;
+
+@Service(ExfiltrationUtils.EGRESS_SUMOLOGIC_SERVICE_QUALIFIER)
+public class EgressVwbRemediationService extends EgressRemediationService {
+  private static final Logger logger =
+      Logger.getLogger(EgressVwbRemediationService.class.getName());
+
+  // The maximum egress user is allowed when bypassed for large file download. If more than this
+  // user will still trigger egress alert.
+  // Current value is 100GB
+  private static final int EGRESS_HARD_LIMIT_MB = 100 * 1024;
+
+  private final EgressJiraHandler egressJiraHandler;
+  private final MailService mailService;
+  private final EgressEventMapper egressEventMapper;
+  private final UserAdminService userAdminService;
+
+  @Autowired
+  public EgressVwbRemediationService(
+      Clock clock,
+      Provider<WorkbenchConfig> workbenchConfigProvider,
+      UserService userService,
+      LeonardoApiClient leonardoNotebooksClient,
+      EgressEventAuditor egressEventAuditor,
+      EgressEventDao egressEventDao,
+      EgressEventMapper egressEventMapper,
+      @Qualifier(SUMOLOGIC_JIRA_HANDLER_QUALIFIER) EgressJiraHandler egressJiraHandler,
+      MailService mailService,
+      UserAdminService userAdminService) {
+    super(
+        clock,
+        workbenchConfigProvider,
+        userService,
+        leonardoNotebooksClient,
+        egressEventAuditor,
+        egressEventDao,
+        egressEventMapper,
+        userAdminService);
+    this.egressJiraHandler = egressJiraHandler;
+    this.mailService = mailService;
+    this.egressEventMapper = egressEventMapper;
+    this.userAdminService = userAdminService;
+  }
+
+  @Override
+  protected void sendEgressRemediationEmail(
+      DbUser user, EgressRemediationAction action, DbEgressEvent event) throws MessagingException {
+    SumologicEgressEvent originalEvent = egressEventMapper.toSumoLogicEvent(event);
+    String environmentType =
+        appServiceNameToAppType(Strings.nullToEmpty(originalEvent.getSrcGkeServiceName()))
+            .map(LeonardoAppUtils::appDisplayName)
+            .orElse("Jupyter");
+    mailService.sendEgressRemediationEmail(user, action, environmentType);
+  }
+
+  @Override
+  protected void logEvent(DbEgressEvent event, EgressRemediationAction action) throws ApiException {
+    egressJiraHandler.logEventToJira(event, action);
+  }
+
+  @Override
+  protected boolean shouldSkipEgressEvent(DbEgressEvent event) {
+    if (isUserBypassedForLargeFileDownload(event)) {
+      logger.info(
+          String.format(
+              "Skip egress event %d because user is bypassed for large file download",
+              event.getEgressEventId()));
+      return true;
+    }
+    return false;
+  }
+
+  @Override
+  protected void getEgressIncidentCountForUser(DbEgressEvent event, EgressRemediationAction action) throws ApiException {
+    egressJiraHandler.logEventToJira(event, action);
+  }
+
+  private boolean isUserBypassedForLargeFileDownload(DbEgressEvent event) {
+    return userAdminService.getCurrentEgressBypassWindow(event.getUser().getUserId()) != null
+        && event.getEgressMegabytes() != null
+        && event.getEgressMegabytes() < EGRESS_HARD_LIMIT_MB;
+  }
+}
\ No newline at end of file
diff --git a/api/src/main/resources/workbench-api.yaml b/api/src/main/resources/workbench-api.yaml
index b2e6052406a..504b7854ce8 100644
--- a/api/src/main/resources/workbench-api.yaml
+++ b/api/src/main/resources/workbench-api.yaml
@@ -13152,9 +13152,12 @@ components:
           type: integer
           format: int64
           description: current incident count for the user in the workspace
-        workspaceId:
+        vwbWorkspaceId:
           type: string
           description: UUID of a VWB workspace
+        vwbEgressEventId:
+          type: string
+          description: UUID of a VWB egress event
         egress_mib:
           type: number
           description: The total amount of egress detected during the window, in mebibytes
diff --git a/api/src/test/java/org/pmiops/workbench/db/dao/EgressEventDaoTest.java b/api/src/test/java/org/pmiops/workbench/db/dao/EgressEventDaoTest.java
index ba3d4fd1568..2778e29f9ac 100644
--- a/api/src/test/java/org/pmiops/workbench/db/dao/EgressEventDaoTest.java
+++ b/api/src/test/java/org/pmiops/workbench/db/dao/EgressEventDaoTest.java
@@ -4,6 +4,7 @@
 
 import java.sql.Timestamp;
 import java.time.Instant;
+import java.util.List;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 import org.pmiops.workbench.FakeClockConfiguration;
@@ -55,6 +56,27 @@ public void testCreateAndGet() {
         .isEqualTo(DbEgressEvent.DbEgressEventStatus.VERIFIED_FALSE_POSITIVE);
   }
 
+  @Test
+  public void testFindAllByUserAndStatusNotIn() {
+    DbEgressEvent event1 = createValidDbEgressEvent();
+    event1.setStatus(DbEgressEventStatus.PENDING);
+    egressEventDao.save(event1);
+
+    DbEgressEvent event2 = createValidDbEgressEvent();
+    event2.setStatus(DbEgressEventStatus.REMEDIATED);
+    egressEventDao.save(event2);
+
+    DbEgressEvent event3 = createValidDbEgressEvent();
+    event3.setStatus(DbEgressEventStatus.VERIFIED_FALSE_POSITIVE);
+    egressEventDao.save(event3);
+
+    List<DbEgressEventStatus> excludedStatuses = List.of(DbEgressEventStatus.REMEDIATED, DbEgressEventStatus.VERIFIED_FALSE_POSITIVE);
+    List<DbEgressEvent> result = egressEventDao.findAllByUserAndStatusNotIn(user, excludedStatuses);
+
+    assertThat(result).hasSize(1);
+    assertThat(result.get(0).getStatus()).isEqualTo(DbEgressEventStatus.PENDING);
+  }
+
   private DbEgressEvent createValidDbEgressEvent() {
     DbEgressEvent event = new DbEgressEvent();
     event.setUser(user);

From 5dd520533e912af7c70fa827705edad8274cf3ca Mon Sep 17 00:00:00 2001
From: Yonghao Yu <yyhao1@gmail.com>
Date: Tue, 17 Dec 2024 11:15:39 -0500
Subject: [PATCH 06/10] fix test

---
 .../auditors/EgressEventAuditorImpl.java      |  2 +-
 .../api/CloudTaskEgressController.java        | 23 ++++++--
 .../api/OfflineEgressController.java          | 19 +++---
 .../api/VwbEgressAdminController.java         |  6 +-
 .../cloudtasks/TaskQueueService.java          |  4 +-
 .../workbench/db/dao/EgressEventDao.java      |  3 +-
 .../exfiltration/EgressEventServiceImpl.java  |  6 +-
 .../EgressRemediationService.java             | 53 +++++++++++------
 .../exfiltration/ExfiltrationUtils.java       |  1 +
 ...EgressObjectLengthsRemediationService.java |  5 --
 .../EgressSumologicRemediationService.java    |  6 --
 .../impl/EgressVwbRemediationService.java     |  9 +--
 api/src/main/resources/workbench-api.yaml     |  4 ++
 .../api/OfflineEgressControllerTest.java      |  5 +-
 .../api/VwbEgressAdminControllerTest.java     | 28 ++++++---
 .../workbench/db/dao/EgressEventDaoTest.java  |  3 +-
 .../exfiltration/EgressEventServiceTest.java  | 17 +++---
 .../EgressRemediationServiceTest.java         | 58 ++++++++++---------
 18 files changed, 150 insertions(+), 102 deletions(-)

diff --git a/api/src/main/java/org/pmiops/workbench/actionaudit/auditors/EgressEventAuditorImpl.java b/api/src/main/java/org/pmiops/workbench/actionaudit/auditors/EgressEventAuditorImpl.java
index 23d62543281..47adf531b69 100644
--- a/api/src/main/java/org/pmiops/workbench/actionaudit/auditors/EgressEventAuditorImpl.java
+++ b/api/src/main/java/org/pmiops/workbench/actionaudit/auditors/EgressEventAuditorImpl.java
@@ -127,7 +127,7 @@ public void fireVwbEgressEvent(VwbEgressEventRequest event, DbUser user) {
     events.add(
         baseEventBuilder
             .targetPropertyMaybe("vwbWorkspaceId")
-            .newValueMaybe(event.getWorkspaceId())
+            .newValueMaybe(event.getVwbWorkspaceId())
             .build());
     events.add(
         baseEventBuilder
diff --git a/api/src/main/java/org/pmiops/workbench/api/CloudTaskEgressController.java b/api/src/main/java/org/pmiops/workbench/api/CloudTaskEgressController.java
index 844b193e5d9..189a869229e 100644
--- a/api/src/main/java/org/pmiops/workbench/api/CloudTaskEgressController.java
+++ b/api/src/main/java/org/pmiops/workbench/api/CloudTaskEgressController.java
@@ -1,10 +1,11 @@
 package org.pmiops.workbench.api;
 
 import static org.pmiops.workbench.exfiltration.ExfiltrationUtils.EGRESS_SUMOLOGIC_SERVICE_QUALIFIER;
+import static org.pmiops.workbench.exfiltration.ExfiltrationUtils.EGRESS_VWB_SERVICE_QUALIFIER;
 
+import org.pmiops.workbench.db.dao.EgressEventDao;
 import org.pmiops.workbench.exfiltration.EgressRemediationService;
 import org.pmiops.workbench.model.ProcessEgressEventRequest;
-import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.beans.factory.annotation.Qualifier;
 import org.springframework.http.ResponseEntity;
 import org.springframework.web.bind.annotation.RestController;
@@ -12,13 +13,25 @@
 @RestController
 public class CloudTaskEgressController implements CloudTaskEgressApiDelegate {
 
-  @Autowired
-  @Qualifier(EGRESS_SUMOLOGIC_SERVICE_QUALIFIER)
-  private EgressRemediationService egressRemediationService;
+  private final EgressRemediationService sumologicEgressRemediationService;
+  private final EgressRemediationService vwbEgressRemediationService;
+
+  public CloudTaskEgressController(
+      @Qualifier(EGRESS_SUMOLOGIC_SERVICE_QUALIFIER)
+          EgressRemediationService sumologicEgressRemediationService,
+      @Qualifier(EGRESS_VWB_SERVICE_QUALIFIER) EgressRemediationService vwbEgressRemediationService,
+      EgressEventDao egressEventDao) {
+    this.sumologicEgressRemediationService = sumologicEgressRemediationService;
+    this.vwbEgressRemediationService = vwbEgressRemediationService;
+  }
 
   @Override
   public ResponseEntity<Void> processEgressEvent(ProcessEgressEventRequest request) {
-    egressRemediationService.remediateEgressEvent(request.getEventId());
+    if (request.isIsVwbEgressEvent()) {
+      vwbEgressRemediationService.remediateEgressEvent(request.getEventId());
+    } else {
+      sumologicEgressRemediationService.remediateEgressEvent(request.getEventId());
+    }
     return ResponseEntity.ok().build();
   }
 }
diff --git a/api/src/main/java/org/pmiops/workbench/api/OfflineEgressController.java b/api/src/main/java/org/pmiops/workbench/api/OfflineEgressController.java
index a1bcc845f1d..c2f0856869a 100644
--- a/api/src/main/java/org/pmiops/workbench/api/OfflineEgressController.java
+++ b/api/src/main/java/org/pmiops/workbench/api/OfflineEgressController.java
@@ -5,7 +5,6 @@
 import java.time.Duration;
 import java.util.List;
 import java.util.logging.Logger;
-import java.util.stream.Collectors;
 import org.pmiops.workbench.cloudtasks.TaskQueueService;
 import org.pmiops.workbench.db.dao.EgressEventDao;
 import org.pmiops.workbench.db.model.DbEgressEvent;
@@ -27,20 +26,16 @@ public class OfflineEgressController implements OfflineEgressApiDelegate {
   public ResponseEntity<Void> checkPendingEgressEvents() {
     Timestamp latestModifiedTime = Timestamp.from(clock.instant().minus(PENDING_EVENT_LIMIT));
 
-    List<Long> oldPendingEventIds =
-        egressEventDao
-            .findAllByStatusAndLastModifiedTimeLessThan(
-                DbEgressEventStatus.PENDING, latestModifiedTime)
-            .stream()
-            .map(DbEgressEvent::getEgressEventId)
-            .collect(Collectors.toList());
-    for (long eventId : oldPendingEventIds) {
-      taskQueueService.pushEgressEventTask(eventId);
+    List<DbEgressEvent> oldPendingEvents =
+        egressEventDao.findAllByStatusAndLastModifiedTimeLessThan(
+            DbEgressEventStatus.PENDING, latestModifiedTime);
+    for (DbEgressEvent event : oldPendingEvents) {
+      taskQueueService.pushEgressEventTask(event.getEgressEventId(), event.getIsVwb());
     }
 
-    if (oldPendingEventIds.size() > 0) {
+    if (oldPendingEvents.size() > 0) {
       log.warning(
-          String.format("found and re-enqueued %d old PENDING events", oldPendingEventIds.size()));
+          String.format("found and re-enqueued %d old PENDING events", oldPendingEvents.size()));
     }
     return ResponseEntity.noContent().build();
   }
diff --git a/api/src/main/java/org/pmiops/workbench/api/VwbEgressAdminController.java b/api/src/main/java/org/pmiops/workbench/api/VwbEgressAdminController.java
index 0a0458cb337..d6e9f324e22 100644
--- a/api/src/main/java/org/pmiops/workbench/api/VwbEgressAdminController.java
+++ b/api/src/main/java/org/pmiops/workbench/api/VwbEgressAdminController.java
@@ -5,6 +5,7 @@
 import org.pmiops.workbench.config.WorkbenchConfig;
 import org.pmiops.workbench.db.model.DbUser;
 import org.pmiops.workbench.exceptions.ForbiddenException;
+import org.pmiops.workbench.exceptions.UnauthorizedException;
 import org.pmiops.workbench.exfiltration.EgressEventService;
 import org.pmiops.workbench.model.VwbEgressEventRequest;
 import org.springframework.beans.factory.annotation.Autowired;
@@ -32,12 +33,15 @@ public VwbEgressAdminController(
 
   @Override
   public ResponseEntity<Void> createVwbEgressEvent(VwbEgressEventRequest body) {
+    if (!workbenchConfigProvider.get().featureFlags.enableVWBEgressMonitor) {
+      throw new ForbiddenException("VWB Egress Monitor is disabled");
+    }
     if (!workbenchConfigProvider
         .get()
         .vwb
         .exfilManagerServiceAccount
         .equals(userProvider.get().getUsername())) {
-      throw new ForbiddenException(
+      throw new UnauthorizedException(
           String.format(
               "User %s is not authorized to create VwbEgressEvent",
               userProvider.get().getUsername()));
diff --git a/api/src/main/java/org/pmiops/workbench/cloudtasks/TaskQueueService.java b/api/src/main/java/org/pmiops/workbench/cloudtasks/TaskQueueService.java
index 557983b578f..3fbbf5aaec5 100644
--- a/api/src/main/java/org/pmiops/workbench/cloudtasks/TaskQueueService.java
+++ b/api/src/main/java/org/pmiops/workbench/cloudtasks/TaskQueueService.java
@@ -175,11 +175,11 @@ private <T> List<List<T>> groupDeleteWorkspaceTasks(List<T> workspacesToDelete)
     return CloudTasksUtils.partitionList(workspacesToDelete, batchSize);
   }
 
-  public void pushEgressEventTask(Long eventId) {
+  public void pushEgressEventTask(Long eventId, boolean isVwbEgressEvent) {
     createAndPushTask(
         EGRESS_EVENT_QUEUE_NAME,
         EGRESS_EVENT_PATH,
-        new ProcessEgressEventRequest().eventId(eventId));
+        new ProcessEgressEventRequest().eventId(eventId).isVwbEgressEvent(isVwbEgressEvent));
   }
 
   public void pushCreateWorkspaceTask(long operationId, Workspace workspace) {
diff --git a/api/src/main/java/org/pmiops/workbench/db/dao/EgressEventDao.java b/api/src/main/java/org/pmiops/workbench/db/dao/EgressEventDao.java
index 4b469d338c8..a67b9534071 100644
--- a/api/src/main/java/org/pmiops/workbench/db/dao/EgressEventDao.java
+++ b/api/src/main/java/org/pmiops/workbench/db/dao/EgressEventDao.java
@@ -24,7 +24,8 @@ public interface EgressEventDao
   Page<DbEgressEvent> findAllByUserAndWorkspaceOrderByCreationTimeDesc(
       DbUser user, DbWorkspace workspace, Pageable p);
 
-  List<DbEgressEvent> findAllByUserAndStatusNotIn(DbUser user, Collection<DbEgressEventStatus> status);
+  List<DbEgressEvent> findAllByUserAndStatusNotIn(
+      DbUser user, Collection<DbEgressEventStatus> status);
 
   List<DbEgressEvent> findAllByUserAndWorkspaceAndCreationTimeBetweenAndCreationTimeNot(
       DbUser user,
diff --git a/api/src/main/java/org/pmiops/workbench/exfiltration/EgressEventServiceImpl.java b/api/src/main/java/org/pmiops/workbench/exfiltration/EgressEventServiceImpl.java
index 06f976de81b..49f9d6b8969 100644
--- a/api/src/main/java/org/pmiops/workbench/exfiltration/EgressEventServiceImpl.java
+++ b/api/src/main/java/org/pmiops/workbench/exfiltration/EgressEventServiceImpl.java
@@ -102,7 +102,7 @@ public void handleEvent(SumologicEgressEvent event) {
 
     Optional<DbEgressEvent> maybeEvent =
         this.maybePersistSumoEgressEvent(event, Optional.of(egressUser.get()), dbWorkspaceMaybe);
-    maybeEvent.ifPresent(e -> taskQueueService.pushEgressEventTask(e.getEgressEventId()));
+    maybeEvent.ifPresent(e -> taskQueueService.pushEgressEventTask(e.getEgressEventId(), false));
   }
 
   @Override
@@ -112,7 +112,7 @@ public void handleVwbEvent(VwbEgressEventRequest egressEvent) {
       logger.severe(
           String.format(
               "An egress event happened for workspace: %s, but we're unable to find a user %s.",
-              egressEvent.getWorkspaceId(), egressEvent.getUserEmail()));
+              egressEvent.getVwbWorkspaceId(), egressEvent.getUserEmail()));
     }
     this.egressEventAuditor.fireVwbEgressEvent(egressEvent, egressUser.get());
     DbEgressEvent event =
@@ -132,7 +132,7 @@ public void handleVwbEvent(VwbEgressEventRequest egressEvent) {
                 .setEgressWindowSeconds(
                     Optional.ofNullable(egressEvent.getTimeWindowDuration()).orElse(null))
                 .setStatus(DbEgressEventStatus.PENDING));
-    taskQueueService.pushEgressEventTask(event.getEgressEventId());
+    taskQueueService.pushEgressEventTask(event.getEgressEventId(), true);
   }
 
   @NotNull
diff --git a/api/src/main/java/org/pmiops/workbench/exfiltration/EgressRemediationService.java b/api/src/main/java/org/pmiops/workbench/exfiltration/EgressRemediationService.java
index 292815bf3c0..6089da6a335 100644
--- a/api/src/main/java/org/pmiops/workbench/exfiltration/EgressRemediationService.java
+++ b/api/src/main/java/org/pmiops/workbench/exfiltration/EgressRemediationService.java
@@ -100,7 +100,7 @@ public void remediateEgressEvent(long egressEventId) {
       return;
     }
     // Gather/merge past alerts into an incident history
-    int egressIncidentCount = getEgressIncidentCountForUser(user);
+    int egressIncidentCount = getEgressIncidentCountForUser(event, user);
 
     // Determine escalating remediation action, if any
     EgressAlertRemediationPolicy egressPolicy =
@@ -110,7 +110,7 @@ public void remediateEgressEvent(long egressEventId) {
     // Execute the action, if any
     escalation.ifPresent(
         e -> {
-          EgressRemediationAction action = getEgressRemediationAction(user, e);
+          EgressRemediationAction action = getEgressRemediationAction(user, e, event);
 
           if (egressPolicy != null && egressPolicy.enableJiraTicketing) {
             try {
@@ -135,18 +135,26 @@ public void remediateEgressEvent(long egressEventId) {
   }
 
   @NotNull
-  private EgressRemediationAction getEgressRemediationAction(DbUser user, Escalation e) {
-    EgressRemediationAction action;
-    if (e.disableUser != null) {
-      disableUser(user);
-      action = EgressRemediationAction.DISABLE_USER;
-    } else if (e.suspendCompute != null) {
-      suspendUserCompute(user, Duration.ofMinutes(e.suspendCompute.durationMinutes));
-      action = EgressRemediationAction.SUSPEND_COMPUTE;
+  private EgressRemediationAction getEgressRemediationAction(
+      DbUser user, Escalation e, DbEgressEvent event) {
+    if (event.getIsVwb()) {
+      if (e.disableUser != null) {
+        disableUserVwbDataAccess(user);
+        return EgressRemediationAction.DISABLE_USER;
+      } else if (e.suspendCompute != null) {
+        // Compute will be suspended by VWB.
+        return EgressRemediationAction.SUSPEND_COMPUTE;
+      }
     } else {
-      throw new ServerErrorException("egress alert policy is invalid: " + e);
+      if (e.disableUser != null) {
+        disableUser(user);
+        return EgressRemediationAction.DISABLE_USER;
+      } else if (e.suspendCompute != null) {
+        suspendUserCompute(user, Duration.ofMinutes(e.suspendCompute.durationMinutes));
+        return EgressRemediationAction.SUSPEND_COMPUTE;
+      }
     }
-    return action;
+    throw new ServerErrorException("egress alert policy is invalid: " + e);
   }
 
   protected abstract void sendEgressRemediationEmail(
@@ -168,9 +176,11 @@ protected abstract void logEvent(DbEgressEvent event, EgressRemediationAction ac
    * @return the count of logical egress incidents for this user for all time, including any events
    *     which are actively being processed
    */
-  private int getEgressIncidentCountForUser(DbUser user) {
+  protected int getEgressIncidentCountForUser(DbEgressEvent event, DbUser user) {
     List<DbEgressEvent> events =
-        egressEventDao.findAllByUserAndStatusNotIn(user, List.of(DbEgressEventStatus.VERIFIED_FALSE_POSITIVE, DbEgressEventStatus.BYPASSED));
+        egressEventDao.findAllByUserAndStatusNotIn(
+            user,
+            List.of(DbEgressEventStatus.VERIFIED_FALSE_POSITIVE, DbEgressEventStatus.BYPASSED));
 
     // If any egress alerts are missing workspace metadata (this should not happen), consider each
     // as unique incidents.
@@ -280,6 +290,18 @@ protected void disableUser(DbUser user) {
     stopUserRuntimesAndApps(user.getUsername());
   }
 
+  protected void disableUserVwbDataAccess(DbUser user) {
+    userService.updateUserWithRetries(
+        u -> {
+          u.setDisabled(true);
+          return u;
+        },
+        user,
+        Agent.asSystem());
+
+    // TODO: Lock user access in VWB.
+  }
+
   private void stopUserRuntimesAndApps(String userEmail) {
     int stoppedRuntimeCount = leonardoNotebooksClient.stopAllUserRuntimesAsService(userEmail);
     int stoppedAppCount = leonardoNotebooksClient.deleteUserAppsAsService(userEmail);
@@ -288,7 +310,4 @@ private void stopUserRuntimesAndApps(String userEmail) {
             "stopped %d runtimes and %d apps for user %s",
             stoppedRuntimeCount, stoppedAppCount, userEmail));
   }
-
-  protected abstract void getEgressIncidentCountForUser(DbEgressEvent event,
-      EgressRemediationAction action) throws ApiException;
 }
diff --git a/api/src/main/java/org/pmiops/workbench/exfiltration/ExfiltrationUtils.java b/api/src/main/java/org/pmiops/workbench/exfiltration/ExfiltrationUtils.java
index 93b7b772901..df5bba41848 100644
--- a/api/src/main/java/org/pmiops/workbench/exfiltration/ExfiltrationUtils.java
+++ b/api/src/main/java/org/pmiops/workbench/exfiltration/ExfiltrationUtils.java
@@ -14,6 +14,7 @@ public class ExfiltrationUtils {
 
   public static final String EGRESS_SUMOLOGIC_SERVICE_QUALIFIER = "sumologicEgressService";
   public static final String SUMOLOGIC_JIRA_HANDLER_QUALIFIER = "sumologicJiraHandler";
+  public static final String EGRESS_VWB_SERVICE_QUALIFIER = "vwbEgressService";
 
   private static final Pattern VM_PREFIX_PATTERN = Pattern.compile("all-of-us-(?<userid>\\d+)");
   private static final String USER_ID_GROUP_NAME = "userid";
diff --git a/api/src/main/java/org/pmiops/workbench/exfiltration/impl/EgressObjectLengthsRemediationService.java b/api/src/main/java/org/pmiops/workbench/exfiltration/impl/EgressObjectLengthsRemediationService.java
index 140474b6cc5..5b68c0ca1b1 100644
--- a/api/src/main/java/org/pmiops/workbench/exfiltration/impl/EgressObjectLengthsRemediationService.java
+++ b/api/src/main/java/org/pmiops/workbench/exfiltration/impl/EgressObjectLengthsRemediationService.java
@@ -72,9 +72,4 @@ protected void logEvent(DbEgressEvent event, EgressRemediationAction action) thr
   protected boolean shouldSkipEgressEvent(DbEgressEvent event) {
     return false;
   }
-
-  @Override
-  protected void getEgressIncidentCountForUser(DbEgressEvent event, EgressRemediationAction action)
-      throws ApiException {
-  }
 }
diff --git a/api/src/main/java/org/pmiops/workbench/exfiltration/impl/EgressSumologicRemediationService.java b/api/src/main/java/org/pmiops/workbench/exfiltration/impl/EgressSumologicRemediationService.java
index c77f5d243f2..3112d6978d7 100644
--- a/api/src/main/java/org/pmiops/workbench/exfiltration/impl/EgressSumologicRemediationService.java
+++ b/api/src/main/java/org/pmiops/workbench/exfiltration/impl/EgressSumologicRemediationService.java
@@ -108,12 +108,6 @@ protected boolean shouldSkipEgressEvent(DbEgressEvent event) {
     return false;
   }
 
-  @Override
-  protected void getEgressIncidentCountForUser(DbEgressEvent event, EgressRemediationAction action)
-      throws ApiException {
-
-  }
-
   private boolean isUserBypassedForLargeFileDownload(DbEgressEvent event) {
     return userAdminService.getCurrentEgressBypassWindow(event.getUser().getUserId()) != null
         && event.getEgressMegabytes() != null
diff --git a/api/src/main/java/org/pmiops/workbench/exfiltration/impl/EgressVwbRemediationService.java b/api/src/main/java/org/pmiops/workbench/exfiltration/impl/EgressVwbRemediationService.java
index 690047aab11..9b7cf21270c 100644
--- a/api/src/main/java/org/pmiops/workbench/exfiltration/impl/EgressVwbRemediationService.java
+++ b/api/src/main/java/org/pmiops/workbench/exfiltration/impl/EgressVwbRemediationService.java
@@ -29,7 +29,7 @@
 import org.springframework.beans.factory.annotation.Qualifier;
 import org.springframework.stereotype.Service;
 
-@Service(ExfiltrationUtils.EGRESS_SUMOLOGIC_SERVICE_QUALIFIER)
+@Service(ExfiltrationUtils.EGRESS_VWB_SERVICE_QUALIFIER)
 public class EgressVwbRemediationService extends EgressRemediationService {
   private static final Logger logger =
       Logger.getLogger(EgressVwbRemediationService.class.getName());
@@ -74,6 +74,7 @@ public EgressVwbRemediationService(
   @Override
   protected void sendEgressRemediationEmail(
       DbUser user, EgressRemediationAction action, DbEgressEvent event) throws MessagingException {
+    // TODO(RW-14138): Vwb email support
     SumologicEgressEvent originalEvent = egressEventMapper.toSumoLogicEvent(event);
     String environmentType =
         appServiceNameToAppType(Strings.nullToEmpty(originalEvent.getSrcGkeServiceName()))
@@ -100,8 +101,8 @@ protected boolean shouldSkipEgressEvent(DbEgressEvent event) {
   }
 
   @Override
-  protected void getEgressIncidentCountForUser(DbEgressEvent event, EgressRemediationAction action) throws ApiException {
-    egressJiraHandler.logEventToJira(event, action);
+  protected int getEgressIncidentCountForUser(DbEgressEvent event, DbUser user) {
+    return event.getVwbIncidentCount();
   }
 
   private boolean isUserBypassedForLargeFileDownload(DbEgressEvent event) {
@@ -109,4 +110,4 @@ private boolean isUserBypassedForLargeFileDownload(DbEgressEvent event) {
         && event.getEgressMegabytes() != null
         && event.getEgressMegabytes() < EGRESS_HARD_LIMIT_MB;
   }
-}
\ No newline at end of file
+}
diff --git a/api/src/main/resources/workbench-api.yaml b/api/src/main/resources/workbench-api.yaml
index 504b7854ce8..d9a315bb9ad 100644
--- a/api/src/main/resources/workbench-api.yaml
+++ b/api/src/main/resources/workbench-api.yaml
@@ -11729,6 +11729,10 @@ components:
           type: integer
           description: Egress event ID to be processed
           format: int64
+        isVwbEgressEvent:
+          type: boolean
+          description: Is this a VWB egress event
+          default: false
     DeleteTestUserWorkspacesRequest:
       type: array
       items:
diff --git a/api/src/test/java/org/pmiops/workbench/api/OfflineEgressControllerTest.java b/api/src/test/java/org/pmiops/workbench/api/OfflineEgressControllerTest.java
index 757c6b58b00..16f8756a822 100644
--- a/api/src/test/java/org/pmiops/workbench/api/OfflineEgressControllerTest.java
+++ b/api/src/test/java/org/pmiops/workbench/api/OfflineEgressControllerTest.java
@@ -1,5 +1,6 @@
 package org.pmiops.workbench.api;
 
+import static org.mockito.ArgumentMatchers.anyBoolean;
 import static org.mockito.ArgumentMatchers.anyLong;
 import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
@@ -56,7 +57,7 @@ public void testCheckPendingEgressEvents() {
     fakeClock.setInstant(FakeClockConfiguration.NOW.toInstant());
     offlineEgressController.checkPendingEgressEvents();
 
-    verify(mockTaskQueueService, times(2)).pushEgressEventTask(anyLong());
+    verify(mockTaskQueueService, times(2)).pushEgressEventTask(anyLong(), anyBoolean());
   }
 
   @Test
@@ -74,6 +75,6 @@ public void testCheckPendingEgressEvents_noMatches() {
   }
 
   private DbEgressEvent newEvent() {
-    return new DbEgressEvent().setStatus(DbEgressEventStatus.PENDING);
+    return new DbEgressEvent().setStatus(DbEgressEventStatus.PENDING).setIsVwb(true);
   }
 }
diff --git a/api/src/test/java/org/pmiops/workbench/api/VwbEgressAdminControllerTest.java b/api/src/test/java/org/pmiops/workbench/api/VwbEgressAdminControllerTest.java
index d4fc8a1f547..9b774d2bcd4 100644
--- a/api/src/test/java/org/pmiops/workbench/api/VwbEgressAdminControllerTest.java
+++ b/api/src/test/java/org/pmiops/workbench/api/VwbEgressAdminControllerTest.java
@@ -10,21 +10,17 @@
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 import org.mockito.Mock;
-import org.pmiops.workbench.FakeClockConfiguration;
 import org.pmiops.workbench.config.WorkbenchConfig;
 import org.pmiops.workbench.db.model.DbUser;
 import org.pmiops.workbench.exceptions.ForbiddenException;
+import org.pmiops.workbench.exceptions.UnauthorizedException;
 import org.pmiops.workbench.exfiltration.EgressEventService;
 import org.pmiops.workbench.model.VwbEgressEventRequest;
-import org.springframework.boot.test.autoconfigure.orm.jpa.DataJpaTest;
-import org.springframework.context.annotation.Import;
 import org.springframework.http.HttpStatus;
 import org.springframework.http.ResponseEntity;
-import org.springframework.test.annotation.DirtiesContext;
+import org.springframework.test.context.junit.jupiter.SpringJUnitConfig;
 
-@DataJpaTest
-@DirtiesContext(classMode = DirtiesContext.ClassMode.BEFORE_EACH_TEST_METHOD)
-@Import(FakeClockConfiguration.class)
+@SpringJUnitConfig
 public class VwbEgressAdminControllerTest {
   @Mock private Provider<DbUser> mockUserProvider;
   @Mock private Provider<WorkbenchConfig> mockConfigProvider;
@@ -42,6 +38,9 @@ public void setUp() {
         new VwbEgressAdminController(mockEgressEventService, mockConfigProvider, mockUserProvider);
     authorizedUser = new DbUser();
     authorizedUser.setUsername("test-user@example.com");
+    workbenchConfig.featureFlags = new WorkbenchConfig.FeatureFlagsConfig();
+    workbenchConfig.featureFlags.enableVWBEgressMonitor = true;
+
     workbenchConfig.vwb = new WorkbenchConfig.VwbConfig();
     workbenchConfig.vwb.exfilManagerServiceAccount = authorizedUser.getUsername();
     when(mockConfigProvider.get()).thenReturn(workbenchConfig);
@@ -49,7 +48,7 @@ public void setUp() {
     vwbEvent =
         new VwbEgressEventRequest()
             .userEmail("testuser@example.com")
-            .workspaceId("testWorkspaceId")
+            .vwbWorkspaceId("testWorkspaceId")
             .vmName("testVmName")
             .incidentCount(1L)
             .egressMib(500.0)
@@ -76,6 +75,19 @@ public void testCreateVwbEgressEvent_unauthorizedUser() {
 
     when(mockUserProvider.get()).thenReturn(dbUser);
 
+    assertThrows(
+        UnauthorizedException.class,
+        () -> {
+          vwbEgressAdminController.createVwbEgressEvent(vwbEvent);
+        });
+  }
+
+  @Test
+  public void testCreateVwbEgressEvent_flagNotEnabled() {
+    workbenchConfig.featureFlags = new WorkbenchConfig.FeatureFlagsConfig();
+    workbenchConfig.featureFlags.enableVWBEgressMonitor = false;
+    when(mockUserProvider.get()).thenReturn(authorizedUser);
+
     assertThrows(
         ForbiddenException.class,
         () -> {
diff --git a/api/src/test/java/org/pmiops/workbench/db/dao/EgressEventDaoTest.java b/api/src/test/java/org/pmiops/workbench/db/dao/EgressEventDaoTest.java
index 2778e29f9ac..7dffab3316d 100644
--- a/api/src/test/java/org/pmiops/workbench/db/dao/EgressEventDaoTest.java
+++ b/api/src/test/java/org/pmiops/workbench/db/dao/EgressEventDaoTest.java
@@ -70,7 +70,8 @@ public void testFindAllByUserAndStatusNotIn() {
     event3.setStatus(DbEgressEventStatus.VERIFIED_FALSE_POSITIVE);
     egressEventDao.save(event3);
 
-    List<DbEgressEventStatus> excludedStatuses = List.of(DbEgressEventStatus.REMEDIATED, DbEgressEventStatus.VERIFIED_FALSE_POSITIVE);
+    List<DbEgressEventStatus> excludedStatuses =
+        List.of(DbEgressEventStatus.REMEDIATED, DbEgressEventStatus.VERIFIED_FALSE_POSITIVE);
     List<DbEgressEvent> result = egressEventDao.findAllByUserAndStatusNotIn(user, excludedStatuses);
 
     assertThat(result).hasSize(1);
diff --git a/api/src/test/java/org/pmiops/workbench/exfiltration/EgressEventServiceTest.java b/api/src/test/java/org/pmiops/workbench/exfiltration/EgressEventServiceTest.java
index e11a8f32121..0d31db6503b 100644
--- a/api/src/test/java/org/pmiops/workbench/exfiltration/EgressEventServiceTest.java
+++ b/api/src/test/java/org/pmiops/workbench/exfiltration/EgressEventServiceTest.java
@@ -1,6 +1,7 @@
 package org.pmiops.workbench.exfiltration;
 
 import static com.google.common.truth.Truth.assertThat;
+import static org.mockito.ArgumentMatchers.anyBoolean;
 import static org.mockito.ArgumentMatchers.anyLong;
 import static org.mockito.Mockito.doReturn;
 import static org.mockito.Mockito.verify;
@@ -132,7 +133,7 @@ public void testCreateEgressEventAlert() throws Exception {
     SumologicEgressEvent event = recentEgressEventForUser(dbUser1);
     egressEventService.handleEvent(event);
     verify(mockEgressEventAuditor).fireEgressEventForUser(event, dbUser1);
-    verify(mockTaskQueueService).pushEgressEventTask(anyLong());
+    verify(mockTaskQueueService).pushEgressEventTask(anyLong(), anyBoolean());
 
     List<DbEgressEvent> dbEvents = ImmutableList.copyOf(egressEventDao.findAll());
     assertThat(dbEvents).hasSize(1);
@@ -154,7 +155,7 @@ public void testAppCreateEgressEventAlert() {
 
     egressEventService.handleEvent(event);
     verify(mockEgressEventAuditor).fireEgressEventForUser(event, dbUser1);
-    verify(mockTaskQueueService).pushEgressEventTask(anyLong());
+    verify(mockTaskQueueService).pushEgressEventTask(anyLong(), anyBoolean());
 
     List<DbEgressEvent> dbEvents = ImmutableList.copyOf(egressEventDao.findAll());
     assertThat(dbEvents).hasSize(1);
@@ -176,7 +177,7 @@ public void testAppCreateEgressEventAlert_skipCromwell() {
 
     egressEventService.handleEvent(event);
     verify(mockEgressEventAuditor).fireEgressEventForUser(event, dbUser1);
-    verify(mockTaskQueueService).pushEgressEventTask(anyLong());
+    verify(mockTaskQueueService).pushEgressEventTask(anyLong(), anyBoolean());
 
     List<DbEgressEvent> dbEvents = ImmutableList.copyOf(egressEventDao.findAll());
     assertThat(dbEvents).hasSize(1);
@@ -231,7 +232,7 @@ public void testCreateEgressEventAlert_staleEventsMultiwindow() {
 
     fakeClock.setInstant(NOW);
     egressEventService.handleEvent(oldEgressEvent);
-    verify(mockTaskQueueService).pushEgressEventTask(anyLong());
+    verify(mockTaskQueueService).pushEgressEventTask(anyLong(), anyBoolean());
 
     Iterable<DbEgressEvent> dbEvents = egressEventDao.findAll();
     assertThat(dbEvents).hasSize(2);
@@ -256,7 +257,7 @@ public void testCreateEgressEventAlert_staleEventsDifferentUsers() {
     fakeClock.setInstant(NOW);
     egressEventService.handleEvent(oldEgressEvent);
     verify(mockEgressEventAuditor).fireEgressEventForUser(oldEgressEvent, dbUser1);
-    verify(mockTaskQueueService).pushEgressEventTask(anyLong());
+    verify(mockTaskQueueService).pushEgressEventTask(anyLong(), anyBoolean());
 
     Iterable<DbEgressEvent> dbEvents = egressEventDao.findAll();
     assertThat(dbEvents).hasSize(2);
@@ -282,7 +283,7 @@ public void testCreateEgressEventAlert_staleEventShortWindowPersisted() {
     fakeClock.setInstant(NOW);
     egressEventService.handleEvent(oldEgressEvent);
     verify(mockEgressEventAuditor).fireEgressEventForUser(oldEgressEvent, dbUser1);
-    verify(mockTaskQueueService).pushEgressEventTask(anyLong());
+    verify(mockTaskQueueService).pushEgressEventTask(anyLong(), anyBoolean());
 
     Iterable<DbEgressEvent> dbEvents = egressEventDao.findAll();
     assertThat(dbEvents).hasSize(2);
@@ -293,7 +294,7 @@ public void testHandleVwbEgressEvent() {
     VwbEgressEventRequest vwbEvent =
         new VwbEgressEventRequest()
             .userEmail(dbUser1.getUsername())
-            .workspaceId("testWorkspaceId")
+            .vwbWorkspaceId("testWorkspaceId")
             .vmName("testVmName")
             .incidentCount(1L)
             .egressMib(500.0)
@@ -307,7 +308,7 @@ public void testHandleVwbEgressEvent() {
     egressEventService.handleVwbEvent(vwbEvent);
 
     verify(mockEgressEventAuditor).fireVwbEgressEvent(vwbEvent, dbUser1);
-    verify(mockTaskQueueService).pushEgressEventTask(anyLong());
+    verify(mockTaskQueueService).pushEgressEventTask(anyLong(), anyBoolean());
 
     List<DbEgressEvent> dbEvents = ImmutableList.copyOf(egressEventDao.findAll());
     assertThat(dbEvents).hasSize(1);
diff --git a/api/src/test/java/org/pmiops/workbench/exfiltration/EgressRemediationServiceTest.java b/api/src/test/java/org/pmiops/workbench/exfiltration/EgressRemediationServiceTest.java
index 7aaad0d4502..4140e31cbca 100644
--- a/api/src/test/java/org/pmiops/workbench/exfiltration/EgressRemediationServiceTest.java
+++ b/api/src/test/java/org/pmiops/workbench/exfiltration/EgressRemediationServiceTest.java
@@ -11,6 +11,7 @@
 import static org.mockito.Mockito.verifyNoInteractions;
 import static org.mockito.Mockito.when;
 import static org.pmiops.workbench.exfiltration.ExfiltrationUtils.EGRESS_SUMOLOGIC_SERVICE_QUALIFIER;
+import static org.pmiops.workbench.exfiltration.ExfiltrationUtils.EGRESS_VWB_SERVICE_QUALIFIER;
 
 import com.google.common.collect.ImmutableList;
 import com.google.gson.Gson;
@@ -98,7 +99,11 @@ public class EgressRemediationServiceTest {
 
   @Autowired
   @Qualifier(EGRESS_SUMOLOGIC_SERVICE_QUALIFIER)
-  private EgressRemediationService egressRemediationService;
+  private EgressRemediationService sumologicEgressRemediationService;
+
+  @Autowired
+  @Qualifier(EGRESS_VWB_SERVICE_QUALIFIER)
+  private EgressRemediationService vwbEgressRemediationService;
 
   @Autowired private EgressJiraHandler egressJiraHandler;
 
@@ -180,13 +185,14 @@ public void tearDown() {
   @Test
   public void testRemediateEgressEvent_notFound() {
     assertThrows(
-        NotFoundException.class, () -> egressRemediationService.remediateEgressEvent(123L));
+        NotFoundException.class,
+        () -> sumologicEgressRemediationService.remediateEgressEvent(123L));
   }
 
   @Test
   public void testRemediateEgressEvent_alreadyRemediated() {
     long eventId = saveNewEvent(newGCEEvent().setStatus(DbEgressEventStatus.REMEDIATED));
-    egressRemediationService.remediateEgressEvent(eventId);
+    sumologicEgressRemediationService.remediateEgressEvent(eventId);
 
     assertThat(getDbUser().getDisabled()).isFalse();
     assertComputeNotSuspended();
@@ -197,7 +203,7 @@ public void testRemediateEgressEvent_alreadyRemediated() {
   public void testRemediateEgressEvent_emptyPolicy() {
     workbenchConfig.egressAlertRemediationPolicy = null;
 
-    egressRemediationService.remediateEgressEvent(saveNewEvent());
+    sumologicEgressRemediationService.remediateEgressEvent(saveNewEvent());
 
     assertThat(getDbUser().getDisabled()).isFalse();
     assertComputeNotSuspended();
@@ -209,7 +215,7 @@ public void testRemediateEgressEvent_noMatchingEscalationPolicy() {
     workbenchConfig.egressAlertRemediationPolicy.escalations =
         ImmutableList.of(disableUserAfter(10));
 
-    egressRemediationService.remediateEgressEvent(saveNewEvent());
+    sumologicEgressRemediationService.remediateEgressEvent(saveNewEvent());
 
     assertThat(getDbUser().getDisabled()).isFalse();
     assertComputeNotSuspended();
@@ -222,7 +228,7 @@ public void testRemediateEgressEvent_activeBypassWindow() {
         .thenReturn(new EgressBypassWindow());
 
     long eventId = saveNewEvent(newGCEEvent().setEgressMegabytes((float) (200)));
-    egressRemediationService.remediateEgressEvent(eventId);
+    sumologicEgressRemediationService.remediateEgressEvent(eventId);
 
     assertThat(getDbUser().getDisabled()).isFalse();
     assertComputeNotSuspended();
@@ -237,7 +243,7 @@ public void testRemediateEgressEvent_activeBypassWindowButExeceedHardLimit() {
 
     long eventId = saveNewEvent(newGCEEvent().setEgressMegabytes((float) (110 * 1024)));
 
-    egressRemediationService.remediateEgressEvent(eventId);
+    sumologicEgressRemediationService.remediateEgressEvent(eventId);
 
     assertComputeSuspended(Duration.ofMinutes(1));
   }
@@ -247,7 +253,7 @@ public void testRemediateEgressEvent_firstIncident() {
     workbenchConfig.egressAlertRemediationPolicy.escalations =
         ImmutableList.of(suspendComputeAfter(1, Duration.ofMinutes(1)), disableUserAfter(2));
 
-    egressRemediationService.remediateEgressEvent(saveNewEvent());
+    sumologicEgressRemediationService.remediateEgressEvent(saveNewEvent());
 
     assertThat(getDbUser().getDisabled()).isFalse();
     assertComputeSuspended(Duration.ofMinutes(1));
@@ -262,7 +268,7 @@ public void testRemediateEgressEvent_intermediateEscalation() {
             disableUserAfter(3));
 
     saveOldEvents(Duration.ofDays(1));
-    egressRemediationService.remediateEgressEvent(saveNewEvent());
+    sumologicEgressRemediationService.remediateEgressEvent(saveNewEvent());
 
     assertThat(getDbUser().getDisabled()).isFalse();
     assertComputeSuspended(Duration.ofMinutes(2));
@@ -282,7 +288,7 @@ public void testRemediateEgressEvent_maxEscalation() {
             .mapToObj(Duration::ofDays)
             .collect(Collectors.toList())
             .toArray(new Duration[] {}));
-    egressRemediationService.remediateEgressEvent(saveNewEvent());
+    sumologicEgressRemediationService.remediateEgressEvent(saveNewEvent());
 
     assertThat(getDbUser().getDisabled()).isTrue();
     assertComputeNotSuspended();
@@ -295,7 +301,7 @@ public void testRemediateEgressEvent_simpleIncidentMerge() {
     // Two events within an hour of each-other should merge
     saveOldEvents(Duration.ofHours(22), Duration.ofHours(23));
 
-    egressRemediationService.remediateEgressEvent(saveNewEvent());
+    sumologicEgressRemediationService.remediateEgressEvent(saveNewEvent());
 
     // 2 "incidents" == 2 minute suspension
     assertComputeSuspended(Duration.ofMinutes(2));
@@ -306,7 +312,7 @@ public void testRemediateEgressEvent_incidentIncludesActiveEvent() {
     workbenchConfig.egressAlertRemediationPolicy = suspendXMinutesOnXIncidentsPolicy();
 
     saveOldEvents(Duration.ofMinutes(1), Duration.ofMinutes(2));
-    egressRemediationService.remediateEgressEvent(saveNewEvent());
+    sumologicEgressRemediationService.remediateEgressEvent(saveNewEvent());
 
     // All 3 events should merge into a single logical incident.
     assertComputeSuspended(Duration.ofMinutes(1));
@@ -321,7 +327,7 @@ public void testRemediateEgressEvent_noIncidentMergeAcrossWorkspaces() {
         oldEvent(Duration.ofHours(3)).setWorkspace(dbWorkspace2),
         oldEvent(Duration.ofHours(4)).setWorkspace(dbWorkspace2));
 
-    egressRemediationService.remediateEgressEvent(saveNewEvent());
+    sumologicEgressRemediationService.remediateEgressEvent(saveNewEvent());
 
     // The two events in workspace 2 merge, the other old event and new event do not merge.
     assertComputeSuspended(Duration.ofMinutes(3));
@@ -336,7 +342,7 @@ public void testRemediateEgressEvent_noIncidentMergeMissingWorkspace() {
         oldEvent(Duration.ofHours(3)).setWorkspace(null),
         oldEvent(Duration.ofHours(3)).setWorkspace(null));
 
-    egressRemediationService.remediateEgressEvent(saveNewEvent());
+    sumologicEgressRemediationService.remediateEgressEvent(saveNewEvent());
 
     // None of the events merge
     assertComputeSuspended(Duration.ofMinutes(4));
@@ -348,7 +354,7 @@ public void testRemediateEgressEvent_incidentMergeIsNonAssociative() {
 
     // events 1, 2 should be merged but 3 should not
     saveOldEvents(Duration.ofMinutes(300), Duration.ofMinutes(250), Duration.ofMinutes(200));
-    egressRemediationService.remediateEgressEvent(saveNewEvent());
+    sumologicEgressRemediationService.remediateEgressEvent(saveNewEvent());
 
     // 1 merged old incident, 1 standalone, and the active event
     assertComputeSuspended(Duration.ofMinutes(3));
@@ -361,7 +367,7 @@ public void testRemediateEgressEvent_incidentsExcludeFalsePostives() {
     saveOldEvents(
         oldEvent(Duration.ofHours(3)).setStatus(DbEgressEventStatus.VERIFIED_FALSE_POSITIVE));
 
-    egressRemediationService.remediateEgressEvent(saveNewEvent());
+    sumologicEgressRemediationService.remediateEgressEvent(saveNewEvent());
 
     // Only the active event is considered
     assertComputeSuspended(Duration.ofMinutes(1));
@@ -373,7 +379,7 @@ public void testRemediateEgressEvent_suspendCompute() throws Exception {
         ImmutableList.of(suspendComputeAfter(1, Duration.ofMinutes(1)));
 
     long eventId = saveNewEvent();
-    egressRemediationService.remediateEgressEvent(eventId);
+    sumologicEgressRemediationService.remediateEgressEvent(eventId);
 
     assertThat(getDbUser().getDisabled()).isFalse();
     assertComputeSuspended(Duration.ofMinutes(1));
@@ -392,7 +398,7 @@ public void testRemediateGKEEgressEvent_suspendCompute() throws Exception {
         ImmutableList.of(suspendComputeAfter(1, Duration.ofMinutes(1)));
 
     long eventId = saveNewGKEEvent(AppType.RSTUDIO);
-    egressRemediationService.remediateEgressEvent(eventId);
+    sumologicEgressRemediationService.remediateEgressEvent(eventId);
 
     assertThat(getDbUser().getDisabled()).isFalse();
     assertComputeSuspended(Duration.ofMinutes(1));
@@ -411,7 +417,7 @@ public void testRemediateGKEEgressEvent_skipCromwellApp() {
         ImmutableList.of(suspendComputeAfter(1, Duration.ofMinutes(1)));
 
     long eventId = saveNewGKEEvent(AppType.CROMWELL);
-    egressRemediationService.remediateEgressEvent(eventId);
+    sumologicEgressRemediationService.remediateEgressEvent(eventId);
 
     assertThat(getDbUser().getDisabled()).isFalse();
     assertComputeNotSuspended();
@@ -424,7 +430,7 @@ public void testRemediateEgressEvent_disableUser() throws Exception {
         ImmutableList.of(disableUserAfter(1));
 
     long eventId = saveNewEvent();
-    egressRemediationService.remediateEgressEvent(eventId);
+    sumologicEgressRemediationService.remediateEgressEvent(eventId);
 
     assertThat(getDbUser().getDisabled()).isTrue();
     assertComputeNotSuspended();
@@ -443,7 +449,7 @@ public void testRemediateEgressEvent_auditOnRemediation() {
     workbenchConfig.egressAlertRemediationPolicy.escalations =
         ImmutableList.of(suspendComputeAfter(1, Duration.ofMinutes(1)));
 
-    egressRemediationService.remediateEgressEvent(saveNewEvent());
+    sumologicEgressRemediationService.remediateEgressEvent(saveNewEvent());
     verify(mockEgressEventAuditor).fireRemediateEgressEvent(any(), notNull());
   }
 
@@ -453,7 +459,7 @@ public void testRemediateEgressEvent_notify_skipIfRecentEvent() {
         ImmutableList.of(suspendComputeAfter(1, Duration.ofMinutes(1)));
 
     saveOldEvents(oldEvent(Duration.ofMinutes(30L)).setStatus(DbEgressEventStatus.REMEDIATED));
-    egressRemediationService.remediateEgressEvent(saveNewEvent());
+    sumologicEgressRemediationService.remediateEgressEvent(saveNewEvent());
 
     verifyNoInteractions(mockMailService);
   }
@@ -468,7 +474,7 @@ public void testRemediateEgressEvent_notify_sendIfRecentUnrelatedEvents() throws
             .setStatus(DbEgressEventStatus.REMEDIATED)
             .setWorkspace(dbWorkspace2),
         oldEvent(Duration.ofHours(2L)).setStatus(DbEgressEventStatus.REMEDIATED));
-    egressRemediationService.remediateEgressEvent(saveNewEvent());
+    sumologicEgressRemediationService.remediateEgressEvent(saveNewEvent());
 
     verify(mockMailService)
         .sendEgressRemediationEmail(
@@ -482,7 +488,7 @@ public void testRemediateEgressEvent_noJiraTicketWhenJiraDisabled() throws Excep
     workbenchConfig.egressAlertRemediationPolicy.escalations =
         ImmutableList.of(suspendComputeAfter(1, Duration.ofMinutes(1)));
 
-    egressRemediationService.remediateEgressEvent(saveNewEvent());
+    sumologicEgressRemediationService.remediateEgressEvent(saveNewEvent());
 
     verifyNoInteractions(mockJiraApi);
   }
@@ -497,7 +503,7 @@ public void testRemediateEgressEvent_createJiraTicketNoExisting() throws Excepti
     workbenchConfig.egressAlertRemediationPolicy.escalations =
         ImmutableList.of(suspendComputeAfter(1, Duration.ofMinutes(1)));
 
-    egressRemediationService.remediateEgressEvent(saveNewEvent());
+    sumologicEgressRemediationService.remediateEgressEvent(saveNewEvent());
 
     ArgumentCaptor<IssueUpdateDetails> captor = ArgumentCaptor.forClass(IssueUpdateDetails.class);
     verify(mockJiraApi).createIssue(captor.capture(), anyBoolean());
@@ -532,7 +538,7 @@ public void testRemediateEgressEvent_commentExistingJiraTicket() throws Exceptio
     workbenchConfig.egressAlertRemediationPolicy.escalations =
         ImmutableList.of(suspendComputeAfter(1, Duration.ofMinutes(1)));
 
-    egressRemediationService.remediateEgressEvent(saveNewEvent());
+    sumologicEgressRemediationService.remediateEgressEvent(saveNewEvent());
 
     ArgumentCaptor<Comment> captor = ArgumentCaptor.forClass(Comment.class);
     verify(mockJiraApi).addComment(captor.capture(), eq("123"));

From c5e9bf9e7387d16386fe75ef8c7a521a4ab45a7f Mon Sep 17 00:00:00 2001
From: Yonghao Yu <yyhao1@gmail.com>
Date: Tue, 17 Dec 2024 14:18:31 -0500
Subject: [PATCH 07/10] fix test

---
 api/db/changelog/db.changelog-master.xml                     | 2 +-
 .../org/pmiops/workbench/api/CloudTaskEgressController.java  | 5 ++---
 .../workbench/exfiltration/EgressEventServiceImpl.java       | 1 +
 .../workbench/exfiltration/EgressEventServiceTest.java       | 3 ++-
 .../workbench/exfiltration/EgressRemediationServiceTest.java | 2 ++
 5 files changed, 8 insertions(+), 5 deletions(-)

diff --git a/api/db/changelog/db.changelog-master.xml b/api/db/changelog/db.changelog-master.xml
index 9738c248791..c7fcec7161f 100644
--- a/api/db/changelog/db.changelog-master.xml
+++ b/api/db/changelog/db.changelog-master.xml
@@ -253,7 +253,7 @@
   <include file="changelog/db.changelog-243-replace-extension-count.xml"/>
   <include file="changelog/db.changelog-244-add-is-vwb-workspace-column.xml"/>
   <include file="changelog/db.changelog-245-add-vwb-template-id-cdr-version.xml"/>
-  <include file="changelog/db.db.changelog-246-add-vwb-egress-event-columns.xml"/>
+  <include file="changelog/db.changelog-246-add-vwb-egress-event-columns.xml"/>
   <!--
    Note: to update the DB locally, do the following:
    - Migrate schema changes: `./project.rb run-local-all-migrations`
diff --git a/api/src/main/java/org/pmiops/workbench/api/CloudTaskEgressController.java b/api/src/main/java/org/pmiops/workbench/api/CloudTaskEgressController.java
index 189a869229e..a209e28e05d 100644
--- a/api/src/main/java/org/pmiops/workbench/api/CloudTaskEgressController.java
+++ b/api/src/main/java/org/pmiops/workbench/api/CloudTaskEgressController.java
@@ -3,7 +3,6 @@
 import static org.pmiops.workbench.exfiltration.ExfiltrationUtils.EGRESS_SUMOLOGIC_SERVICE_QUALIFIER;
 import static org.pmiops.workbench.exfiltration.ExfiltrationUtils.EGRESS_VWB_SERVICE_QUALIFIER;
 
-import org.pmiops.workbench.db.dao.EgressEventDao;
 import org.pmiops.workbench.exfiltration.EgressRemediationService;
 import org.pmiops.workbench.model.ProcessEgressEventRequest;
 import org.springframework.beans.factory.annotation.Qualifier;
@@ -19,8 +18,8 @@ public class CloudTaskEgressController implements CloudTaskEgressApiDelegate {
   public CloudTaskEgressController(
       @Qualifier(EGRESS_SUMOLOGIC_SERVICE_QUALIFIER)
           EgressRemediationService sumologicEgressRemediationService,
-      @Qualifier(EGRESS_VWB_SERVICE_QUALIFIER) EgressRemediationService vwbEgressRemediationService,
-      EgressEventDao egressEventDao) {
+      @Qualifier(EGRESS_VWB_SERVICE_QUALIFIER)
+          EgressRemediationService vwbEgressRemediationService) {
     this.sumologicEgressRemediationService = sumologicEgressRemediationService;
     this.vwbEgressRemediationService = vwbEgressRemediationService;
   }
diff --git a/api/src/main/java/org/pmiops/workbench/exfiltration/EgressEventServiceImpl.java b/api/src/main/java/org/pmiops/workbench/exfiltration/EgressEventServiceImpl.java
index 49f9d6b8969..25333ee85d8 100644
--- a/api/src/main/java/org/pmiops/workbench/exfiltration/EgressEventServiceImpl.java
+++ b/api/src/main/java/org/pmiops/workbench/exfiltration/EgressEventServiceImpl.java
@@ -123,6 +123,7 @@ public void handleVwbEvent(VwbEgressEventRequest egressEvent) {
                 .setVwbEgressEventId(egressEvent.getVwbEgressEventId())
                 .setVwbWorkspaceId(egressEvent.getVwbWorkspaceId())
                 .setVwbVmName(egressEvent.getVmName())
+                .setGcpProjectId(egressEvent.getGcpProjectId())
                 .setVwbIncidentCount(egressEvent.getIncidentCount().intValue())
                 .setEgressMegabytes(
                     Optional.ofNullable(egressEvent.getEgressMib())
diff --git a/api/src/test/java/org/pmiops/workbench/exfiltration/EgressEventServiceTest.java b/api/src/test/java/org/pmiops/workbench/exfiltration/EgressEventServiceTest.java
index 0d31db6503b..24b097dcd6a 100644
--- a/api/src/test/java/org/pmiops/workbench/exfiltration/EgressEventServiceTest.java
+++ b/api/src/test/java/org/pmiops/workbench/exfiltration/EgressEventServiceTest.java
@@ -317,7 +317,8 @@ public void testHandleVwbEgressEvent() {
     assertThat(dbEvent.getVwbWorkspaceId()).isEqualTo("testWorkspaceId");
     assertThat(dbEvent.getVwbVmName()).isEqualTo("testVmName");
     assertThat(dbEvent.getVwbIncidentCount()).isEqualTo(1);
-    assertThat(dbEvent.getEgressMegabytes()).isEqualTo(524.288); // Convert MiB
+    assertThat(dbEvent.getEgressMegabytes())
+        .isEqualTo((float) (500 * ((1 << 20) / 1e6))); // Convert MiB
     assertThat(dbEvent.getGcpProjectId()).isEqualTo("test-gcp-project");
     assertThat(dbEvent.getEgressWindowSeconds()).isEqualTo(600L);
     assertThat(dbEvent.getIsVwb()).isTrue();
diff --git a/api/src/test/java/org/pmiops/workbench/exfiltration/EgressRemediationServiceTest.java b/api/src/test/java/org/pmiops/workbench/exfiltration/EgressRemediationServiceTest.java
index 4140e31cbca..196d7022196 100644
--- a/api/src/test/java/org/pmiops/workbench/exfiltration/EgressRemediationServiceTest.java
+++ b/api/src/test/java/org/pmiops/workbench/exfiltration/EgressRemediationServiceTest.java
@@ -47,6 +47,7 @@
 import org.pmiops.workbench.db.model.DbWorkspace;
 import org.pmiops.workbench.exceptions.NotFoundException;
 import org.pmiops.workbench.exfiltration.impl.EgressSumologicRemediationService;
+import org.pmiops.workbench.exfiltration.impl.EgressVwbRemediationService;
 import org.pmiops.workbench.exfiltration.jirahandler.EgressJiraHandler;
 import org.pmiops.workbench.exfiltration.jirahandler.EgressSumologicJiraHandler;
 import org.pmiops.workbench.jira.JiraContent;
@@ -115,6 +116,7 @@ public class EgressRemediationServiceTest {
   @Import({
     EgressEventMapperImpl.class,
     EgressSumologicRemediationService.class,
+    EgressVwbRemediationService.class,
     FakeClockConfiguration.class,
     FakeJpaDateTimeConfiguration.class,
     JiraService.class,

From 9b97800d973bbe577b2c0509899172cdc980eca3 Mon Sep 17 00:00:00 2001
From: Yonghao Yu <yyhao1@gmail.com>
Date: Tue, 17 Dec 2024 15:54:45 -0500
Subject: [PATCH 08/10] fix test

---
 .../java/org/pmiops/workbench/db/dao/EgressEventDaoTest.java   | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/api/src/test/java/org/pmiops/workbench/db/dao/EgressEventDaoTest.java b/api/src/test/java/org/pmiops/workbench/db/dao/EgressEventDaoTest.java
index 7dffab3316d..dd8c482341f 100644
--- a/api/src/test/java/org/pmiops/workbench/db/dao/EgressEventDaoTest.java
+++ b/api/src/test/java/org/pmiops/workbench/db/dao/EgressEventDaoTest.java
@@ -12,7 +12,6 @@
 import org.pmiops.workbench.db.model.DbEgressEvent;
 import org.pmiops.workbench.db.model.DbEgressEvent.DbEgressEventStatus;
 import org.pmiops.workbench.db.model.DbUser;
-import org.pmiops.workbench.db.model.DbWorkspace;
 import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.boot.test.autoconfigure.orm.jpa.DataJpaTest;
 import org.springframework.context.annotation.Import;
@@ -81,11 +80,11 @@ public void testFindAllByUserAndStatusNotIn() {
   private DbEgressEvent createValidDbEgressEvent() {
     DbEgressEvent event = new DbEgressEvent();
     event.setUser(user);
-    event.setWorkspace(new DbWorkspace());
     event.setCreationTime(CURRENT_TIMESTAMP);
     event.setStatus(DbEgressEvent.DbEgressEventStatus.PENDING);
     event.setEgressWindowSeconds(3600L);
     event.setVwbWorkspaceId("testWorkspaceId");
+    event.setLastModifiedTime(CURRENT_TIMESTAMP);
     event.setVwbVmName("testVmName");
     event.setGcpProjectId("test-project");
     event.setIsVwb(true);

From e9f0717aecdaa2b7856b59b5c0ca22713f4775aa Mon Sep 17 00:00:00 2001
From: Yonghao Yu <yyhao1@gmail.com>
Date: Wed, 18 Dec 2024 09:14:58 -0500
Subject: [PATCH 09/10] fix test

---
 ...angelog-246-add-vwb-egress-event-columns.xml |  1 -
 .../workbench/api/OfflineEgressController.java  |  4 +++-
 .../workbench/db/model/DbEgressEvent.java       | 17 ++---------------
 .../exfiltration/EgressEventServiceImpl.java    |  2 --
 .../exfiltration/EgressRemediationService.java  |  2 +-
 .../exfiltration/ExfiltrationUtils.java         |  7 +++++++
 .../api/OfflineEgressControllerTest.java        |  4 +++-
 .../workbench/db/dao/EgressEventDaoTest.java    |  2 +-
 8 files changed, 17 insertions(+), 22 deletions(-)

diff --git a/api/db/changelog/db.changelog-246-add-vwb-egress-event-columns.xml b/api/db/changelog/db.changelog-246-add-vwb-egress-event-columns.xml
index 617ae7e062b..d1e9ab5ebd8 100644
--- a/api/db/changelog/db.changelog-246-add-vwb-egress-event-columns.xml
+++ b/api/db/changelog/db.changelog-246-add-vwb-egress-event-columns.xml
@@ -6,7 +6,6 @@
                       http://www.liquibase.org/xml/ns/dbchangelog/dbchangelog-3.1.xsd">
     <changeSet author="yonghao" id="db.changelog-246-add-is-vwb-egress-event-column">
         <addColumn tableName="egress_event">
-            <column name="is_vwb" type="boolean" />
             <column name="vwb_workspace_id" type="varchar(80)" />
             <column name="vwb_vm_name" type="varchar(80)" />
             <column name="gcp_project_id" type="varchar(80)" />
diff --git a/api/src/main/java/org/pmiops/workbench/api/OfflineEgressController.java b/api/src/main/java/org/pmiops/workbench/api/OfflineEgressController.java
index c2f0856869a..08f7caa6448 100644
--- a/api/src/main/java/org/pmiops/workbench/api/OfflineEgressController.java
+++ b/api/src/main/java/org/pmiops/workbench/api/OfflineEgressController.java
@@ -9,6 +9,7 @@
 import org.pmiops.workbench.db.dao.EgressEventDao;
 import org.pmiops.workbench.db.model.DbEgressEvent;
 import org.pmiops.workbench.db.model.DbEgressEvent.DbEgressEventStatus;
+import org.pmiops.workbench.exfiltration.ExfiltrationUtils;
 import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.http.ResponseEntity;
 import org.springframework.web.bind.annotation.RestController;
@@ -30,7 +31,8 @@ public ResponseEntity<Void> checkPendingEgressEvents() {
         egressEventDao.findAllByStatusAndLastModifiedTimeLessThan(
             DbEgressEventStatus.PENDING, latestModifiedTime);
     for (DbEgressEvent event : oldPendingEvents) {
-      taskQueueService.pushEgressEventTask(event.getEgressEventId(), event.getIsVwb());
+      taskQueueService.pushEgressEventTask(
+          event.getEgressEventId(), ExfiltrationUtils.isVwbEgressEvent(event));
     }
 
     if (oldPendingEvents.size() > 0) {
diff --git a/api/src/main/java/org/pmiops/workbench/db/model/DbEgressEvent.java b/api/src/main/java/org/pmiops/workbench/db/model/DbEgressEvent.java
index 277b7688d43..f5c07a29a33 100644
--- a/api/src/main/java/org/pmiops/workbench/db/model/DbEgressEvent.java
+++ b/api/src/main/java/org/pmiops/workbench/db/model/DbEgressEvent.java
@@ -46,7 +46,6 @@ public enum DbEgressEventStatus {
   private String vwbVmName;
   private String gcpProjectId;
   private String vwbEgressEventId;
-  private boolean isVwb;
   private int vwbIncidentCount;
 
   @Id
@@ -206,16 +205,6 @@ public DbEgressEvent setVwbIncidentCount(int vwbIncidentCount) {
     return this;
   }
 
-  @Column(name = "is_vwb")
-  public boolean getIsVwb() {
-    return isVwb;
-  }
-
-  public DbEgressEvent setIsVwb(boolean isVwb) {
-    this.isVwb = isVwb;
-    return this;
-  }
-
   public boolean equals(Object o) {
     if (this == o) {
       return true;
@@ -236,8 +225,7 @@ public boolean equals(Object o) {
         && Objects.equals(vwbWorkspaceId, that.vwbWorkspaceId)
         && Objects.equals(vwbVmName, that.vwbVmName)
         && Objects.equals(gcpProjectId, that.gcpProjectId)
-        && Objects.equals(vwbIncidentCount, that.vwbIncidentCount)
-        && Objects.equals(isVwb, that.isVwb);
+        && Objects.equals(vwbIncidentCount, that.vwbIncidentCount);
   }
 
   @Override
@@ -255,7 +243,6 @@ public int hashCode() {
         vwbWorkspaceId,
         vwbVmName,
         gcpProjectId,
-        vwbIncidentCount,
-        isVwb);
+        vwbIncidentCount);
   }
 }
diff --git a/api/src/main/java/org/pmiops/workbench/exfiltration/EgressEventServiceImpl.java b/api/src/main/java/org/pmiops/workbench/exfiltration/EgressEventServiceImpl.java
index 25333ee85d8..7a57130e1fc 100644
--- a/api/src/main/java/org/pmiops/workbench/exfiltration/EgressEventServiceImpl.java
+++ b/api/src/main/java/org/pmiops/workbench/exfiltration/EgressEventServiceImpl.java
@@ -119,7 +119,6 @@ public void handleVwbEvent(VwbEgressEventRequest egressEvent) {
         egressEventDao.save(
             new DbEgressEvent()
                 .setUser(egressUser.orElse(null))
-                .setIsVwb(true)
                 .setVwbEgressEventId(egressEvent.getVwbEgressEventId())
                 .setVwbWorkspaceId(egressEvent.getVwbWorkspaceId())
                 .setVwbVmName(egressEvent.getVmName())
@@ -183,7 +182,6 @@ private Optional<DbEgressEvent> maybePersistSumoEgressEvent(
         egressEventDao.save(
             new DbEgressEvent()
                 .setUser(userMaybe.orElse(null))
-                .setIsVwb(false)
                 .setWorkspace(workspaceMaybe.orElse(null))
                 .setEgressMegabytes(
                     Optional.ofNullable(event.getEgressMib())
diff --git a/api/src/main/java/org/pmiops/workbench/exfiltration/EgressRemediationService.java b/api/src/main/java/org/pmiops/workbench/exfiltration/EgressRemediationService.java
index 6089da6a335..ac7d148388a 100644
--- a/api/src/main/java/org/pmiops/workbench/exfiltration/EgressRemediationService.java
+++ b/api/src/main/java/org/pmiops/workbench/exfiltration/EgressRemediationService.java
@@ -137,7 +137,7 @@ public void remediateEgressEvent(long egressEventId) {
   @NotNull
   private EgressRemediationAction getEgressRemediationAction(
       DbUser user, Escalation e, DbEgressEvent event) {
-    if (event.getIsVwb()) {
+    if (ExfiltrationUtils.isVwbEgressEvent(event)) {
       if (e.disableUser != null) {
         disableUserVwbDataAccess(user);
         return EgressRemediationAction.DISABLE_USER;
diff --git a/api/src/main/java/org/pmiops/workbench/exfiltration/ExfiltrationUtils.java b/api/src/main/java/org/pmiops/workbench/exfiltration/ExfiltrationUtils.java
index df5bba41848..f6b56176380 100644
--- a/api/src/main/java/org/pmiops/workbench/exfiltration/ExfiltrationUtils.java
+++ b/api/src/main/java/org/pmiops/workbench/exfiltration/ExfiltrationUtils.java
@@ -2,6 +2,8 @@
 
 import java.util.Optional;
 import java.util.regex.Pattern;
+import org.javers.common.string.Strings;
+import org.pmiops.workbench.db.model.DbEgressEvent;
 import org.pmiops.workbench.utils.Matchers;
 
 /** Utilities and constants for exfiltration services. */
@@ -27,4 +29,9 @@ public static Optional<Long> gkeServiceNameToUserDatabaseId(String gkeServiceNam
     return Matchers.getGroup(VM_PREFIX_PATTERN, gkeServiceName, USER_ID_GROUP_NAME)
         .map(Long::parseLong);
   }
+
+  /** If egress event has VWB workspace id, it would be an VWB egress event */
+  public static boolean isVwbEgressEvent(DbEgressEvent e) {
+    return Strings.isNonEmpty(e.getVwbWorkspaceId());
+  }
 }
diff --git a/api/src/test/java/org/pmiops/workbench/api/OfflineEgressControllerTest.java b/api/src/test/java/org/pmiops/workbench/api/OfflineEgressControllerTest.java
index 16f8756a822..f73ccc6ee66 100644
--- a/api/src/test/java/org/pmiops/workbench/api/OfflineEgressControllerTest.java
+++ b/api/src/test/java/org/pmiops/workbench/api/OfflineEgressControllerTest.java
@@ -75,6 +75,8 @@ public void testCheckPendingEgressEvents_noMatches() {
   }
 
   private DbEgressEvent newEvent() {
-    return new DbEgressEvent().setStatus(DbEgressEventStatus.PENDING).setIsVwb(true);
+    return new DbEgressEvent()
+        .setStatus(DbEgressEventStatus.PENDING)
+        .setVwbWorkspaceId("wb-test-123");
   }
 }
diff --git a/api/src/test/java/org/pmiops/workbench/db/dao/EgressEventDaoTest.java b/api/src/test/java/org/pmiops/workbench/db/dao/EgressEventDaoTest.java
index dd8c482341f..ab4a4993b8e 100644
--- a/api/src/test/java/org/pmiops/workbench/db/dao/EgressEventDaoTest.java
+++ b/api/src/test/java/org/pmiops/workbench/db/dao/EgressEventDaoTest.java
@@ -87,7 +87,7 @@ private DbEgressEvent createValidDbEgressEvent() {
     event.setLastModifiedTime(CURRENT_TIMESTAMP);
     event.setVwbVmName("testVmName");
     event.setGcpProjectId("test-project");
-    event.setIsVwb(true);
+    ;
     return event;
   }
 }

From 9d197018c97d64fff7c8a85fa355b53d7d9cff72 Mon Sep 17 00:00:00 2001
From: Yonghao Yu <yyhao1@gmail.com>
Date: Wed, 18 Dec 2024 09:32:07 -0500
Subject: [PATCH 10/10] fix test

---
 .../java/org/pmiops/workbench/db/dao/EgressEventDaoTest.java     | 1 -
 .../pmiops/workbench/exfiltration/EgressEventServiceTest.java    | 1 -
 2 files changed, 2 deletions(-)

diff --git a/api/src/test/java/org/pmiops/workbench/db/dao/EgressEventDaoTest.java b/api/src/test/java/org/pmiops/workbench/db/dao/EgressEventDaoTest.java
index ab4a4993b8e..afde58a8943 100644
--- a/api/src/test/java/org/pmiops/workbench/db/dao/EgressEventDaoTest.java
+++ b/api/src/test/java/org/pmiops/workbench/db/dao/EgressEventDaoTest.java
@@ -45,7 +45,6 @@ public void testCreateAndGet() {
     assertThat(foundEvent.getVwbWorkspaceId()).isEqualTo("testWorkspaceId");
     assertThat(foundEvent.getVwbVmName()).isEqualTo("testVmName");
     assertThat(foundEvent.getGcpProjectId()).isEqualTo("test-project");
-    assertThat(foundEvent.getIsVwb()).isTrue();
 
     // Update
     foundEvent.setStatus(DbEgressEventStatus.VERIFIED_FALSE_POSITIVE);
diff --git a/api/src/test/java/org/pmiops/workbench/exfiltration/EgressEventServiceTest.java b/api/src/test/java/org/pmiops/workbench/exfiltration/EgressEventServiceTest.java
index 24b097dcd6a..843daea7af1 100644
--- a/api/src/test/java/org/pmiops/workbench/exfiltration/EgressEventServiceTest.java
+++ b/api/src/test/java/org/pmiops/workbench/exfiltration/EgressEventServiceTest.java
@@ -321,7 +321,6 @@ public void testHandleVwbEgressEvent() {
         .isEqualTo((float) (500 * ((1 << 20) / 1e6))); // Convert MiB
     assertThat(dbEvent.getGcpProjectId()).isEqualTo("test-gcp-project");
     assertThat(dbEvent.getEgressWindowSeconds()).isEqualTo(600L);
-    assertThat(dbEvent.getIsVwb()).isTrue();
   }
 
   private static SumologicEgressEvent recentEgressEventForUser(DbUser user) {