// Copyright 2020 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
//
//    https://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";

option optimize_for = LITE_RUNTIME;

package private_membership.rlwe;

option go_package = "github.com/google/private-membership";

import "private_membership.proto";
import "serialization.proto";

// Request and Response protos
//
// The protocol consists of two roundtrips:
//   First Round: OprfRequest and OprfResponse
//   Second Round: QueryRequest and QueryResponse

// First request sent by the client for checking membership.
message PrivateMembershipRlweOprfRequest {
  // The identifier encrypted using the elliptic curve commutative
  // cipher under an ephemeral key generated by the client. Before encryption,
  // the identifier is hashed using either SHA256 or a more computationally
  // expensive hash such as Argon2.
  repeated bytes encrypted_ids = 1;

  // The name of the use case of this query.
  RlweUseCase use_case = 2;
}

// First response sent by the server for checking membership.
message PrivateMembershipRlweOprfResponse {
  // The requested encrypted_id entries that are re-encrypted using the elliptic
  // curve commutative cipher using the server’s key.
  //
  // The client will decrypt to simply get the identifier encrypted under the
  // server’s key. The client will use this to match with the associated
  // EncryptedBucket to see if there is a match or not. Additionally, the client
  // will use it to decrypt the AES encrypted value.
  repeated DoublyEncryptedId doubly_encrypted_ids = 1;

  // Unset if the server doesn't support non-sensitive identifier bucketing for
  // the requested use case.
  HashedBucketsParameters hashed_buckets_parameters = 2;

  // Parameters for the encrypted buckets (second level).
  EncryptedBucketsParameters encrypted_buckets_parameters = 3;

  // Parameters describing the Ring LWE group as well as the security parameter.
  RlweParameters rlwe_parameters = 4;

  // Version of key that encrypted the identifier-value pair.
  //
  // The client should provide this value in the second round of the protocol
  // to ensure consistency.
  int64 key_version = 5;
}

// Second request sent by the client for checking membership.
message PrivateMembershipRlweQueryRequest {
  // Upload a retrieval request consisting of RLWE ciphertexts.
  repeated PrivateMembershipRlweQuery queries = 1;

  // The name of the use case of this query.
  RlweUseCase use_case = 2;

  // Version of key that encrypted the identifier-value pair.
  int64 key_version = 3;
}

// Second response sent by the server for checking membership.
message PrivateMembershipRlweQueryResponse {
  repeated PrivateMembershipRlwePirResponse pir_responses = 1;
}

// Proto containing plaintext IDs supplied to the client library.
message RlwePlaintextId {
  // Non-sensitive portion of the identifier.
  string non_sensitive_id = 1;

  // Sensitive portion of the identifier.
  string sensitive_id = 2;
}

// Different use cases of the RLWE-based PSM protocol.
//
// NEXT ID: 25
enum RlweUseCase {
  RLWE_USE_CASE_UNDEFINED = 0;
  TEST_USE_CASE = 1;
  TEST_USE_CASE2 = 2;
  TEST_USE_CASE3 = 3;

  // Use case with empty database.
  EMPTY_USE_CASE = 22;

  // The device state of ChromeOS device, used for Zero Touch, License
  // Packaging. This is zero-disclosure database, i.e. client does not reveal
  // any bits to the server.
  CROS_DEVICE_STATE = 5;

  // The device state of ChromeOS device, used for all managed devices and
  // License Packaging. This is partial-disclosure database, i.e. client will be
  // asked to reveal some bits of the hash of non-sensitive identifier.
  CROS_DEVICE_STATE_UNIFIED = 23;

  // The device state of ChromeOS device.
  CROS_DEVICE_SECONDARY_STATE = 12 [deprecated = true];
  CROS_DEVICE_STATE_BACKUP = 21 [deprecated = true];

  // Use-cases for ChromeOS Device Active Reporting to Fresnel.
  CROS_FRESNEL_DAILY = 13;
  CROS_FRESNEL_MONTHLY = 14;
  CROS_FRESNEL_FIRST_ACTIVE = 15;
  CROS_FRESNEL_7DAY_ACTIVE = 16;
  CROS_FRESNEL_28DAY_ACTIVE = 17;
  CROS_FRESNEL_CHURN_MONTHLY_COHORT = 19;
  CROS_FRESNEL_CHURN_MONTHLY_OBSERVATION = 20;

  // Use-cases for Cellular Carrier Lock (SimLock).
  CROS_SIM_LOCK = 18;
  CROS_SIM_LOCK_DEVMODE = 24;

  reserved 4, 6, 7, 8, 9, 10, 11;
}

// Parameters for the hashed buckets (first level).
//
// Hashed bucket id is derived by taking the prefix of the hash of the non-
// sensitive identifier.
message HashedBucketsParameters {
  // The bit length of the hashed bucket id.
  int32 hashed_bucket_id_length = 1;

  // The type of hash function to be used for the non-sensitive identifier
  // part.
  private_membership.HashType non_sensitive_id_hash_type = 2;
}

// Parameters for the encrypted buckets (second level).
message EncryptedBucketsParameters {
  // The bit length of the encrypted bucket ids.
  //
  // The total number of encrypted buckets can be derived from this field,
  // i.e. 2^encrypted_bucket_id_length.
  int32 encrypted_bucket_id_length = 1;

  // The type of hash function to be used to generate the encrypted bucket ID.
  private_membership.EncryptedBucketHashType sensitive_id_hash_type = 2;
}

// Public Ring-LWE encryption scheme parameters.
message RlweParameters {
  // Sequence of one modulus or two decreasing moduli.
  repeated Uint128 modulus = 1;

  // Log (base 2) of the polynomial degree N (dimension of the vectorspace).
  int32 log_degree = 2;

  // Plaintext modulus in bits log t.
  int32 log_t = 3;

  // Variance of Gaussian error in bits.
  int32 variance = 4;

  // Levels of recursion.
  int32 levels_of_recursion = 5;
}

// Represents a unsigned 128-bit integer.
message Uint128 {
  // The lower 64 bits of the modulus.
  uint64 lo = 1;

  // The higher 64 bits of the modulus.
  uint64 hi = 2;
}

// Single unit of RLWE query.
message PrivateMembershipRlweQuery {
  // Corresponds to the doubly_encrypted_id.queried_encrypted_id in the
  // PrivateMembershipRlweOprfResponse.
  //
  // Should be set by the client.
  bytes queried_encrypted_id = 1;

  // Upload a retrieval request consisting of RLWE ciphertexts.
  PirRequest pir_request = 2;

  // Hashed bucket id at bit level granularity.
  message HashedBucketId {
    // Hashed bucket id, in bytes.
    //
    // All bits after bit_length's bit should be set to 0.
    bytes hashed_bucket_id = 1;

    // Bit length of the hashed bucket id.
    //
    // This value should be in
    // ((hashed_bucket_id.size() - 1) * 8, hashed_bucket_id.size() * 8] range.
    int32 bit_length = 2;
  }
  // Hashed bucket id.
  //
  // The hashed bucket id is computed by taking the prefix of the hashed
  // non-sensitive id at the bit level.
  // e.g. if the hashed non-sensitive id is 11010111 and the bucket id length is
  // 6, this will be populated as 11010100.
  //
  // Should be unset if the server doesn't support non-sensitive identifier
  // bucketing for the requested use case.
  HashedBucketId hashed_bucket_id = 3;
}

// Response consisting of a unique query id and the associated PIR response.
message PrivateMembershipRlwePirResponse {
  // This corresponds to the queried_encrypted_id set by the client in the query
  // request.
  bytes queried_encrypted_id = 1;

  // PIR response corresponding to the PIR request.
  PirResponse pir_response = 2;
}

// Request for private bucket retrieval.
message PirRequest {
  // The request is an encoding of a flattened table of encrypted {1} and {0}
  // messages as necessary for PIR. If the database has D items and we are using
  // L levels of recursion, this table will contain L*(D^(1/L)) messages,
  // representing a table with L rows and D^(1/L) items per row. Each row should
  // contain exactly one {1} ciphertext, with the rest of the items containing
  // {0} messages.
  //
  // The flattened table is laid out in the following order:
  //
  //     level  |  ciphertexts
  //     -----------------------------
  //         1  | {0} ... {1} ... {0}
  //         2  | {0} ... {1} ... {0}
  //       ...  | ...
  //
  // The ciphertexts are uploaded in a compressed format as NTT Polynomials
  // using a PRNG seed to expand them into ciphertexts.
  // If this field is set, none of the sharded_request_type should be set.
  repeated .rlwe.SerializedNttPolynomial request = 1;

  // ExpandedRequest can be useful if the querier would like to manipulate the
  // PirRequest after it's been created by the client library. To modify it, the
  // ciphertext must be expanded first; hence the name ExpandedRequest.
  message ExpandedRequest {
    // Wrapper around .rlwe.SerializedSymmetricRlweCiphertext so we can
    // represent "empty" ciphertext without ambiguity. If `ciphertext` is not
    // set, it is consiered "not requested" and the corresponding index will be
    // ignored and excluded from the response. When used multiple levels of
    // recursion, it is the querier's responsibility to leave the correct
    // positions unset so the final query is valid.
    message SerializedSymmetricRlweCiphertext {
      optional .rlwe.SerializedSymmetricRlweCiphertext ciphertext = 1;
    }
    repeated SerializedSymmetricRlweCiphertext ciphertexts = 1;
  }

  // CompactRequest should be the same as the request field above, except it
  // indicates sharding.
  message CompactRequest {
    repeated .rlwe.SerializedNttPolynomial ciphertexts = 1;
  }

  // If one of the sharded_request_type is set, the above `request` field should
  // be set.
  oneof sharded_request_type {
    // Will be used by hashed bucket sharded request because we don't need to
    // modify the ciphertext.
    CompactRequest compact_request = 3;

    // Will be used by encrypted bucket sharded request because we need to
    // modify the ciphertext.
    ExpandedRequest expanded_request = 4;
  }

  // The PRNG seed used to expand the compressed RLWE ciphertexts.
  bytes prng_seed = 2;
}

// Response for private bucket retrieval.
message PirResponse {
  // Sequence of ciphertexts consisting of encrypted buckets.
  repeated .rlwe.SerializedSymmetricRlweCiphertext response = 1;

  // Size of the plaintext entry in bytes.
  int32 plaintext_entry_size = 2;
}

// Bucket consisting of encrypted id and values.
message EncryptedBucket {
  // All encrypted pairs in the bucket.
  repeated EncryptedIdValuePair encrypted_id_value_pairs = 1;

  // Proto storing encrypted id-value pairs.
  message EncryptedIdValuePair {
    // Identifier that is encrypted using the Elliptic curve commutative cipher.
    // Furthermore, the serialization of the encryption is hashed and only a
    // prefix is used.
    bytes encrypted_id = 1;

    // Value that is encrypted using AES where the key is a serialization
    // of the identifier encrypted using the Elliptic curve commutative cipher.
    bytes encrypted_value = 2;

    // Proto containing the sensitive and non-sensitive portion of the
    // identifier in plaintext.
    RlwePlaintextId id = 3;
  }
}

// A collection of plaintext ids and associated membership responses.
message RlweMembershipResponses {
  // A pair of queried plaintext id and its associated membership response.
  message MembershipResponseEntry {
    // Plaintext id that was queried.
    RlwePlaintextId plaintext_id = 1;

    // Membership response of the queried plaintext id.
    MembershipResponse membership_response = 2;
  }
  repeated MembershipResponseEntry membership_responses = 1;
}