Fork based Review PR - Feature/spanner

.github/workflows/build.yml

@@ -3,7 +3,12 @@ on: [push, pull_request]
 env:
   GCS_BUCKET: autobuilds.grr-response.com
   GCS_BUCKET_OPENAPI: autobuilds-grr-openapi
-  DOCKER_REPOSITORY: ghcr.io/google/grr
+  DOCKER_REPOSITORY: ghcr.io/${{ github.repository_owner }}/grr
+  SPANNER_EMULATOR_HOST: localhost:9010
+  SPANNER_DATABASE: grr-database
+  SPANNER_INSTANCE: grr-instance
+  SPANNER_PROJECT_ID: spanner-emulator-project
+  PROJECT_ID: spanner-emulator-project
 jobs:
   test-devenv:
     runs-on: ubuntu-22.04
@@ -44,6 +49,27 @@ jobs:
           fi
           python3 -m venv --system-site-packages "${HOME}/INSTALL"
           travis/install.sh
+      - name: 'Install Cloud SDK'
+        uses: google-github-actions/setup-gcloud@v2
+        with:
+          install_components: 'beta,pubsub-emulator,cloud-spanner-emulator'
+      - name: 'Start Spanner emulator'
+        run: |
+          gcloud config configurations create emulator
+          gcloud config set auth/disable_credentials true
+          gcloud config set project ${{ env.SPANNER_PROJECT_ID }}
+          gcloud config set api_endpoint_overrides/spanner http://localhost:9020/
+          gcloud emulators spanner start &
+          gcloud spanner instances create ${{ env.SPANNER_INSTANCE }} \
+            --config=emulator-config --description="Spanner Test Instance" --nodes=1
+          sudo apt-get update
+          sudo apt install -y protobuf-compiler libprotobuf-dev
+          # Verify that we can create the Spanner Instance and Database
+          cd grr/server/grr_response_server/databases
+          ./spanner_setup.sh
+          cd -
+          # Remove the Database for a clean test env
+          gcloud spanner databases delete ${{ env.SPANNER_DATABASE }} --instance=${{ env.SPANNER_INSTANCE }} --quiet
       - name: Test
         run: |
           source "${HOME}/INSTALL/bin/activate"

.gitignore

@@ -36,3 +36,5 @@ grr/server/grr_response_server/gui/ui/.angular/
 docker_config_files/*.pem
 compose.watch.yaml
 Dockerfile.client
+
+grr/server/grr_response_server/databases/spanner_grr.pb

compose.testing.yaml

@@ -1,14 +1,14 @@
 
 services:
   grr-admin-ui:
-    image: ghcr.io/google/grr:testing
+    image: ghcr.io/google/grr:test
 
   grr-fleetspeak-frontend:
-    image: ghcr.io/google/grr:testing
+    image: ghcr.io/google/grr:test
 
   grr-worker:
-    image: ghcr.io/google/grr:testing
+    image: ghcr.io/google/grr:test
 
   grr-client:
-    image: ghcr.io/google/grr:testing
+    image: ghcr.io/google/grr:test
     privileged: true

grr/core/grr_response_core/config/data_store.py

@@ -121,3 +121,33 @@
         "Only used when Blobstore.implementation is GCSBlobStore."
     ),
 )
+
+# Spanner configuration.
+config_lib.DEFINE_string(
+    "Spanner.project",
+    default=None,
+    help=(
+        "GCP Project where the Spanner Instance is located."
+    ),
+)
+config_lib.DEFINE_string(
+    "Spanner.instance",
+    default="grr-instance",
+    help="The Spanner Instance for GRR.")
+
+config_lib.DEFINE_string(
+    "Spanner.database",
+    default="grr-database",
+    help="The Spanner Database for GRR.")
+
+config_lib.DEFINE_integer(
+    "Spanner.flow_processing_threads_min",
+    default=1,
+    help="The minimum number of flow-processing worker threads.",
+)
+
+config_lib.DEFINE_integer(
+    "Spanner.flow_processing_threads_max",
+    default=20,
+    help="The maximum number of flow-processing worker threads.",
+)

grr/proto/grr_response_proto/flows.proto

@@ -2287,6 +2287,7 @@ message FlowProcessingRequest {
   optional uint64 creation_time = 4 [(sem_type) = {
     type: "RDFDatetime",
   }];
+  optional string ack_id = 5;
 }
 
 message FlowRequest {

grr/proto/grr_response_proto/objects.proto

@@ -331,6 +331,7 @@ message MessageHandlerRequest {
   optional uint64 leased_until = 5 [(sem_type) = { type: "RDFDatetime" }];
   optional string leased_by = 6;
   optional EmbeddedRDFValue request = 7;
+  optional string ack_id = 8;
 }
 
 message SerializedValueOfUnrecognizedType {

grr/server/grr_response_server/databases/db_blob_references_test.py

@@ -82,7 +82,10 @@ def testMultipleHashBlobReferencesCanBeWrittenAndReadBack(self):
   def testWriteHashBlobHandlesLargeAmountsOfData(self):
     hash_id_blob_refs = {}
 
-    for _ in range(50000):
+    # Limit to 16k records to stay within Spanner 80k mutation/commit limit
+    # https://cloud.google.com/spanner/quotas#limits-for
+    # 16k records * 5 columns = 80k mutations/commit
+    for _ in range(16000):
       hash_id = rdf_objects.SHA256HashID(os.urandom(32))
 
       blob_ref = objects_pb2.BlobReference()

grr/server/grr_response_server/databases/db_cronjob_test.py

@@ -191,7 +191,7 @@ def testCronJobLeasing(self):
     job = self._CreateCronJob()
     self.db.WriteCronJob(job)
 
-    current_time = rdfvalue.RDFDatetime.FromSecondsSinceEpoch(10000)
+    current_time = rdfvalue.RDFDatetime.Now()
     lease_time = rdfvalue.Duration.From(5, rdfvalue.MINUTES)
     with test_lib.FakeTime(current_time):
       leased = self.db.LeaseCronJobs(lease_time=lease_time)
@@ -232,7 +232,7 @@ def testCronJobLeasingByID(self):
     for j in jobs:
       self.db.WriteCronJob(j)
 
-    current_time = rdfvalue.RDFDatetime.FromSecondsSinceEpoch(10000)
+    current_time = rdfvalue.RDFDatetime.Now()
     lease_time = rdfvalue.Duration.From(5, rdfvalue.MINUTES)
     with test_lib.FakeTime(current_time):
       leased = self.db.LeaseCronJobs(
@@ -254,7 +254,7 @@ def testCronJobReturning(self):
     with self.assertRaises(ValueError):
       self.db.ReturnLeasedCronJobs([job])
 
-    current_time = rdfvalue.RDFDatetime.FromSecondsSinceEpoch(10000)
+    current_time = rdfvalue.RDFDatetime.Now()
     with test_lib.FakeTime(current_time):
       leased = self.db.LeaseCronJobs(
           cronjob_ids=[leased_job.cron_job_id],
@@ -276,14 +276,14 @@ def testCronJobReturningMultiple(self):
     for j in jobs:
       self.db.WriteCronJob(j)
 
-    current_time = rdfvalue.RDFDatetime.FromSecondsSinceEpoch(10000)
+    current_time = rdfvalue.RDFDatetime.Now()
     with test_lib.FakeTime(current_time):
       leased = self.db.LeaseCronJobs(
           lease_time=rdfvalue.Duration.From(5, rdfvalue.MINUTES)
       )
       self.assertLen(leased, 3)
 
-    current_time = rdfvalue.RDFDatetime.FromSecondsSinceEpoch(10001)
+    current_time = rdfvalue.RDFDatetime.Now()
     with test_lib.FakeTime(current_time):
       unleased_jobs = self.db.LeaseCronJobs(
           lease_time=rdfvalue.Duration.From(5, rdfvalue.MINUTES)
@@ -292,7 +292,7 @@ def testCronJobReturningMultiple(self):
 
       self.db.ReturnLeasedCronJobs(leased)
 
-    current_time = rdfvalue.RDFDatetime.FromSecondsSinceEpoch(10002)
+    current_time = rdfvalue.RDFDatetime.Now()
     with test_lib.FakeTime(current_time):
       leased = self.db.LeaseCronJobs(
           lease_time=rdfvalue.Duration.From(5, rdfvalue.MINUTES)

grr/server/grr_response_server/databases/db_flows_test.py

@@ -2270,15 +2270,16 @@ def testWritesAndReadsSingleFlowResultOfSingleType(self):
     sample_result = flows_pb2.FlowResult(client_id=client_id, flow_id=flow_id)
     sample_result.payload.Pack(jobs_pb2.ClientSummary(client_id=client_id))
 
-    with test_lib.FakeTime(42):
+    current_time = rdfvalue.RDFDatetime.Now()
+    with test_lib.FakeTime(current_time):
       self.db.WriteFlowResults([sample_result])
 
     results = self.db.ReadFlowResults(client_id, flow_id, 0, 100)
     self.assertLen(results, 1)
     self.assertEqual(results[0].payload, sample_result.payload)
     self.assertEqual(
         rdfvalue.RDFDatetime.FromMicrosecondsSinceEpoch(results[0].timestamp),
-        rdfvalue.RDFDatetime.FromSecondsSinceEpoch(42),
+        current_time,
     )
 
   def testWritesAndReadsRDFStringFlowResult(self):
@@ -2335,7 +2336,8 @@ def testReadResultsRestoresAllFlowResultsFields(self):
     )
     sample_result.payload.Pack(jobs_pb2.ClientSummary(client_id=client_id))
 
-    with test_lib.FakeTime(42):
+    current_time = rdfvalue.RDFDatetime.Now()
+    with test_lib.FakeTime(current_time):
       self.db.WriteFlowResults([sample_result])
 
     results = self.db.ReadFlowResults(client_id, flow_id, 0, 100)
@@ -2346,7 +2348,7 @@ def testReadResultsRestoresAllFlowResultsFields(self):
     self.assertEqual(results[0].payload, sample_result.payload)
     self.assertEqual(
         rdfvalue.RDFDatetime.FromMicrosecondsSinceEpoch(results[0].timestamp),
-        rdfvalue.RDFDatetime.FromSecondsSinceEpoch(42),
+        current_time,
     )
 
   def testWritesAndReadsMultipleFlowResultsOfSingleType(self):
@@ -2743,17 +2745,16 @@ def testWritesAndReadsSingleFlowErrorOfSingleType(self):
     sample_error = flows_pb2.FlowError(client_id=client_id, flow_id=flow_id)
     sample_error.payload.Pack(jobs_pb2.ClientSummary(client_id=client_id))
 
-    with test_lib.FakeTime(42):
+    current_time = rdfvalue.RDFDatetime.Now()
+    with test_lib.FakeTime(current_time):
       self.db.WriteFlowErrors([sample_error])
 
     errors = self.db.ReadFlowErrors(client_id, flow_id, 0, 100)
     self.assertLen(errors, 1)
     self.assertEqual(errors[0].payload, sample_error.payload)
     self.assertEqual(
         errors[0].timestamp,
-        rdfvalue.RDFDatetime.FromSecondsSinceEpoch(
-            42
-        ).AsMicrosecondsSinceEpoch(),
+        current_time.AsMicrosecondsSinceEpoch(),
     )
 
   def testWritesAndReadsMultipleFlowErrorsOfSingleType(self):

grr/server/grr_response_server/databases/db_message_handler_test.py

@@ -93,6 +93,7 @@ def testMessageHandlerRequestLeasing(self):
         m.ClearField("timestamp")
       got += l
     self.db.DeleteMessageHandlerRequests(got)
+    self.db.UnregisterMessageHandler()
 
     got.sort(key=lambda req: req.request_id)
     self.assertEqual(requests, got)

grr/server/grr_response_server/databases/registry_init.py

@@ -3,9 +3,11 @@
 
 from grr_response_server.databases import mem
 from grr_response_server.databases import mysql
+from grr_response_server.databases import spanner
 
 # All available databases go into this registry.
 REGISTRY = {
     "InMemoryDB": mem.InMemoryDB,
     "MysqlDB": mysql.MysqlDB,
+    "SpannerDB": spanner.SpannerDB.FromConfig,
 }

grr/server/grr_response_server/databases/spanner.py

@@ -0,0 +1,82 @@
+
+from google.cloud.spanner import Client
+
+from grr_response_core import config
+from grr_response_core.lib import rdfvalue
+from grr_response_server import threadpool
+from grr_response_server.databases import db as db_module
+from grr_response_server.databases import spanner_artifacts
+from grr_response_server.databases import spanner_blob_keys
+from grr_response_server.databases import spanner_blob_references
+from grr_response_server.databases import spanner_clients
+from grr_response_server.databases import spanner_cron_jobs
+from grr_response_server.databases import spanner_events
+from grr_response_server.databases import spanner_flows
+from grr_response_server.databases import spanner_foreman_rules
+from grr_response_server.databases import spanner_hunts
+from grr_response_server.databases import spanner_paths
+from grr_response_server.databases import spanner_signed_binaries
+from grr_response_server.databases import spanner_signed_commands
+from grr_response_server.databases import spanner_users
+from grr_response_server.databases import spanner_utils
+from grr_response_server.databases import spanner_yara
+from grr_response_server.models import blobs as models_blobs
+from grr_response_server.rdfvalues import objects as rdf_objects
+
+class SpannerDB(
+    spanner_artifacts.ArtifactsMixin,
+    spanner_blob_keys.BlobKeysMixin,
+    spanner_blob_references.BlobReferencesMixin,
+    spanner_clients.ClientsMixin,
+    spanner_cron_jobs.CronJobsMixin,
+    spanner_events.EventsMixin,
+    spanner_flows.FlowsMixin,
+    spanner_foreman_rules.ForemanRulesMixin,
+    spanner_hunts.HuntsMixin,
+    spanner_paths.PathsMixin,
+    spanner_signed_binaries.SignedBinariesMixin,
+    spanner_signed_commands.SignedCommandsMixin,
+    spanner_users.UsersMixin,
+    spanner_yara.YaraMixin,
+    db_module.Database,
+):
+  """A Spanner implementation of the GRR database."""
+
+  def __init__(self, db: spanner_utils.Database) -> None:
+    """Initializes the database."""
+    self.db = db
+    self._write_rows_batch_size = 10000
+
+    self.handler_thread = None
+    self.handler_stop = True
+
+    self.flow_processing_request_handler_thread = None
+    self.flow_processing_request_handler_stop = None
+    self.flow_processing_request_handler_pool = threadpool.ThreadPool.Factory(
+        "spanner_flow_processing_pool",
+        min_threads=config.CONFIG["Spanner.flow_processing_threads_min"],
+        max_threads=config.CONFIG["Spanner.flow_processing_threads_max"],
+    )
+
+  @classmethod
+  def FromConfig(cls) -> "Database":
+    """Creates a GRR database instance for Spanner path specified in the config.
+
+    Returns:
+      A GRR database instance.
+    """
+    project_id = config.CONFIG["Spanner.project"]
+    spanner_client = Client(project_id)
+    spanner_instance = spanner_client.instance(config.CONFIG["Spanner.instance"])
+    spanner_database = spanner_instance.database(config.CONFIG["Spanner.database"])
+
+    return cls(spanner_utils.Database(spanner_database, project_id))
+
+  def Now(self) -> rdfvalue.RDFDatetime:
+    """Retrieves current time as reported by the database."""
+    (timestamp,) = self.db.QuerySingle("SELECT CURRENT_TIMESTAMP()")
+    return rdfvalue.RDFDatetime.FromDatetime(timestamp)
+
+  def MinTimestamp(self) -> rdfvalue.RDFDatetime:
+    """Returns minimal timestamp allowed by the DB."""
+    return rdfvalue.RDFDatetime.FromSecondsSinceEpoch(0)