// Copyright 2024 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

syntax = "proto3";

package fcp.confidentialcompute;

import "google/protobuf/any.proto";
import "fcp/protos/confidentialcompute/access_policy.proto";
import "fcp/protos/confidentialcompute/confidential_transform.proto";
import "fcp/protos/confidentialcompute/private_inference.proto";
import "fcp/protos/confidentialcompute/private_logger_uploads_config.proto";
import "tensorflow_federated/cc/core/impl/aggregation/protocol/configuration.proto";

// Configures the FedSQL container initialization. This configuration will last
// for the lifetime of the TEE and includes the privacy-relevant properties.
message FedSqlContainerInitializeConfiguration {
  tensorflow_federated.aggregation.Configuration agg_configuration = 1;

  // Configuration for running private inference in FedSQL container.
  InferenceInitializeConfiguration inference_init_config = 6;

  // Configuration for processing PrivateLogger uploads.
  PrivateLoggerUploadsConfig private_logger_uploads_config = 10;

  // Configuration for possibly encrypted auto-tuning parameters.
  FedSqlContainerAutoTuningConfig auto_tuning_config = 11;

  reserved 2, 3, 4, 5, 7, 8, 9;
}

// Configuration for serialized and possibly encrypted auto-tuning parameters.
// This parametes are passed in the form of Federate Compute checkpoint that
// contains individual auto-tuning parameters as scalar tensors.
message FedSqlContainerAutoTuningConfig {
  // Contains encryption and encoding metadata for the auto-tuning parameters.
  BlobMetadata auto_tuning_metadata = 1;
  // The blob containing serialized and possibly encrypted auto-tuning data.
  bytes auto_tuning_data = 2 [ctype = CORD];
}

// TODO: b/505099049 - Streamline the naming of these aggregation types.
enum AggCoreAggregationType {
  AGGREGATION_TYPE_UNSPECIFIED = 0;
  // Accumulates the results in memory which are applied during `Commit`.
  AGGREGATION_TYPE_ACCUMULATE = 1;
  // Merges the serialized checkpoint.
  AGGREGATION_TYPE_MERGE = 2;
  // Accumulates the private state produced as a result of
  // `FINALIZATION_TYPE_REPORT_PARTITION`.
  AGGREGATION_TYPE_ACCUMULATE_PRIVATE_STATE = 3;
  // Merges the serialized private state.
  AGGREGATION_TYPE_MERGE_PRIVATE_STATE = 4;
}

// Configures how a blob should be aggregated by the trusted AggCore container.
// This is specified through `WriteRequest.first_request_configuration`.
message FedSqlContainerWriteConfiguration {
  AggCoreAggregationType type = 1;
}

// Configures the range of blobs to commit. This is specified through
// `CommitRequest.configuration`.
message FedSqlContainerCommitConfiguration {
  message Range {
    uint64 start = 1;
    uint64 end = 2;
  }
  // The range of blobs to commit.
  Range range = 1;
}

// TODO: b/505099049 - Streamline the naming of these finalization types.
enum AggCoreFinalizationType {
  FINALIZATION_TYPE_UNSPECIFIED = 0;
  // Return noised results from the checkpoint aggregator.
  FINALIZATION_TYPE_REPORT = 1;
  // Return the serialized internal state of the checkpoint aggregator.
  FINALIZATION_TYPE_SERIALIZE = 2;
  // The below enums are used when output partitioning is enabled.
  // Return the partitioned internal state of the checkpoint aggregator
  // bundled with the budget tracking state.
  FINALIZATION_TYPE_PARTITION = 3;
  // Returns the noised results for the partitioned checkpoint aggregator.
  FINALIZATION_TYPE_REPORT_PARTITION = 4;
  // Returns the serialized private state.
  FINALIZATION_TYPE_SERIALIZE_PRIVATE_STATE = 5;
  // Returns the final private state generating a release token that can be
  // passed to KMS for releasing the results.
  FINALIZATION_TYPE_REPORT_PRIVATE_STATE = 6;
  // Returns the result of aggregation (e.g. with autotuning parameters)
  // bundled with the budget tracking state.
  FINALIZATION_TYPE_REPORT_AUTOTUNING_PARAMS = 7;
}

// Configures how a result should be output by the trusted AggCore container.
// This is specified through `FinalizeRequest.configuration`.
message FedSqlContainerFinalizeConfiguration {
  AggCoreFinalizationType type = 1;

  // The number of partitions to use for the partitioned output. This is only
  // used when `type` is `FINALIZATION_TYPE_PARTITION`. This value must be less
  // than or equal to `max_output_partitions` in
  // `FedSqlContainerConfigConstraints`.
  uint64 num_partitions = 3;

  reserved 2;
}

// The configuration for a partitioned output. This is specified through
// `ReadResponse.first_response_configuration` for each partition.
message FedSqlContainerPartitionedOutputConfiguration {
  // The index of the partition that this output corresponds to.
  uint64 partition_index = 1;
}

// The symmetric keys used to encrypt each output partition.
message FedSqlContainerPartitionKeys {
  repeated PartitionKey keys = 1;

  message PartitionKey {
    uint64 partition_index = 1;
    bytes symmetric_key = 2;
  }
}

// Config constraints for the FedSQL container that must be validated before any
// of the blobs written to the container are decrypted.
// This is required when using the KeyManagementService API, it is ignored when
// using the Ledger API.
message FedSqlContainerConfigConstraints {
  // DP parameters for the FedSQL container.
  // The epsilon and delta parameters represent an upper bound on the privacy
  // usage. It should be fine to define server workloads that consume less than
  // this budget, but not more.
  double epsilon = 1;
  double delta = 2;
  // The column names that are part of the DP unit.
  repeated string dp_column_names = 5;
  // The fcp.confidentialcompute.WindowingSchedule for the DP unit. This is an
  // `Any` because using a WindowingSchedule would have required using the
  // `tf_proto_library` build rule for
  // fed_sql_container_config.proto and all its dependencies.
  // TODO: b/425744048 - Change this to a WindowingSchedule once we remove the
  // dependency on plan.proto from sql_query.proto.
  google.protobuf.Any dp_windowing_schedule = 6;

  // The uri of the aggregation intrinsic to use with AggCore.
  // Use `intrinsic_uris` instead.
  string intrinsic_uri = 3 [deprecated = true];

  // Defines limits on the accesses to blobs.
  AccessBudget access_budget = 4;

  // The uris of the aggregation intrinsics to use with AggCore. If non-empty,
  // this field is used instead of `intrinsic_uri`.
  repeated string intrinsic_uris = 7;

  // This value represents an upper bound on the number of partitions that can
  // be used when performing partitioned aggregation. If unset or set to 0,
  // unpartitioned aggregation will be used i.e. all aggregates are
  // merged together in a single worker before release.
  optional uint64 max_output_partitions = 8;
}
