flwr-nightly 1.9.0.dev20240422__tar.gz → 1.9.0.dev20240424__tar.gz

{flwr_nightly-1.9.0.dev20240422 → flwr_nightly-1.9.0.dev20240424}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: flwr-nightly
-Version: 1.9.0.dev20240422
+Version: 1.9.0.dev20240424
 Summary: Flower: A Friendly Federated Learning Framework
 Home-page: https://flower.ai
 License: Apache-2.0

{flwr_nightly-1.9.0.dev20240422 → flwr_nightly-1.9.0.dev20240424}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "poetry.core.masonry.api"
 
 [tool.poetry]
 name = "flwr-nightly"
-version = "1.9.0.dev20240422"
+version = "1.9.0.dev20240424"
 description = "Flower: A Friendly Federated Learning Framework"
 license = "Apache-2.0"
 authors = ["The Flower Authors <hello@flower.ai>"]

{flwr_nightly-1.9.0.dev20240422 → flwr_nightly-1.9.0.dev20240424}/src/py/flwr/cli/new/new.py

@@ -36,6 +36,7 @@ class MlFramework(str, Enum):
     NUMPY = "NumPy"
     PYTORCH = "PyTorch"
     TENSORFLOW = "TensorFlow"
+    SKLEARN = "sklearn"
 
 
 class TemplateNotFound(Exception):

flwr_nightly-1.9.0.dev20240424/src/py/flwr/cli/new/templates/app/code/client.sklearn.py.tpl

@@ -0,0 +1,94 @@
+"""$project_name: A Flower / Scikit-Learn app."""
+
+import warnings
+
+import numpy as np
+from flwr.client import NumPyClient, ClientApp
+from flwr_datasets import FederatedDataset
+from sklearn.linear_model import LogisticRegression
+from sklearn.metrics import log_loss
+
+
+def get_model_parameters(model):
+    if model.fit_intercept:
+        params = [
+            model.coef_,
+            model.intercept_,
+        ]
+    else:
+        params = [model.coef_]
+    return params
+
+
+def set_model_params(model, params):
+    model.coef_ = params[0]
+    if model.fit_intercept:
+        model.intercept_ = params[1]
+    return model
+
+
+def set_initial_params(model):
+    n_classes = 10  # MNIST has 10 classes
+    n_features = 784  # Number of features in dataset
+    model.classes_ = np.array([i for i in range(10)])
+
+    model.coef_ = np.zeros((n_classes, n_features))
+    if model.fit_intercept:
+        model.intercept_ = np.zeros((n_classes,))
+
+
+class FlowerClient(NumPyClient):
+    def __init__(self, model, X_train, X_test, y_train, y_test):
+        self.model = model
+        self.X_train = X_train
+        self.X_test = X_test
+        self.y_train = y_train
+        self.y_test = y_test
+
+    def get_parameters(self, config):
+        return get_model_parameters(self.model)
+
+    def fit(self, parameters, config):
+        set_model_params(self.model, parameters)
+
+        # Ignore convergence failure due to low local epochs
+        with warnings.catch_warnings():
+            warnings.simplefilter("ignore")
+            self.model.fit(self.X_train, self.y_train)
+
+        return get_model_parameters(self.model), len(self.X_train), {}
+
+    def evaluate(self, parameters, config):
+        set_model_params(self.model, parameters)
+
+        loss = log_loss(self.y_test, self.model.predict_proba(self.X_test))
+        accuracy = self.model.score(self.X_test, self.y_test)
+
+        return loss, len(self.X_test), {"accuracy": accuracy}
+
+fds = FederatedDataset(dataset="mnist", partitioners={"train": 2})
+
+def client_fn(cid: str):
+    dataset = fds.load_partition(int(cid), "train").with_format("numpy")
+
+    X, y = dataset["image"].reshape((len(dataset), -1)), dataset["label"]
+
+    # Split the on edge data: 80% train, 20% test
+    X_train, X_test = X[: int(0.8 * len(X))], X[int(0.8 * len(X)) :]
+    y_train, y_test = y[: int(0.8 * len(y))], y[int(0.8 * len(y)) :]
+
+    # Create LogisticRegression Model
+    model = LogisticRegression(
+        penalty="l2",
+        max_iter=1,  # local epoch
+        warm_start=True,  # prevent refreshing weights when fitting
+    )
+
+    # Setting initial parameters, akin to model.compile for keras models
+    set_initial_params(model)
+
+    return FlowerClient(model, X_train, X_test, y_train, y_test).to_client()
+
+
+# Flower ClientApp
+app = ClientApp(client_fn=client_fn)

flwr_nightly-1.9.0.dev20240424/src/py/flwr/cli/new/templates/app/code/server.sklearn.py.tpl

@@ -0,0 +1,17 @@
+"""$project_name: A Flower / Scikit-Learn app."""
+
+from flwr.server import ServerApp, ServerConfig
+from flwr.server.strategy import FedAvg
+
+
+strategy = FedAvg(
+    fraction_fit=1.0,
+    fraction_evaluate=1.0,
+    min_available_clients=2,
+)
+
+# Create ServerApp
+app = ServerApp(
+    config=ServerConfig(num_rounds=3),
+    strategy=strategy,
+)

flwr_nightly-1.9.0.dev20240424/src/py/flwr/cli/new/templates/app/pyproject.sklearn.toml.tpl

@@ -0,0 +1,24 @@
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
+[project]
+name = "$project_name"
+version = "1.0.0"
+description = ""
+authors = [
+    { name = "The Flower Authors", email = "hello@flower.ai" },
+]
+license = {text = "Apache License (2.0)"}
+dependencies = [
+    "flwr[simulation]>=1.8.0,<2.0",
+    "flwr-datasets[vision]>=0.0.2,<1.0.0",
+    "scikit-learn>=1.1.1",
+]
+
+[tool.hatch.build.targets.wheel]
+packages = ["."]
+
+[flower.components]
+serverapp = "$project_name.server:app"
+clientapp = "$project_name.client:app"

{flwr_nightly-1.9.0.dev20240422 → flwr_nightly-1.9.0.dev20240424}/src/py/flwr/client/grpc_client/connection.py

@@ -22,6 +22,8 @@ from pathlib import Path
 from queue import Queue
 from typing import Callable, Iterator, Optional, Tuple, Union, cast
 
+from cryptography.hazmat.primitives.asymmetric import ec
+
 from flwr.common import (
     DEFAULT_TTL,
     GRPC_MAX_MESSAGE_LENGTH,
@@ -56,12 +58,15 @@ def on_channel_state_change(channel_connectivity: str) -> None:
 
 
 @contextmanager
-def grpc_connection(  # pylint: disable=R0915
+def grpc_connection(  # pylint: disable=R0913, R0915
     server_address: str,
     insecure: bool,
     retry_invoker: RetryInvoker,  # pylint: disable=unused-argument
     max_message_length: int = GRPC_MAX_MESSAGE_LENGTH,
     root_certificates: Optional[Union[bytes, str]] = None,
+    authentication_keys: Optional[  # pylint: disable=unused-argument
+        Tuple[ec.EllipticCurvePrivateKey, ec.EllipticCurvePublicKey]
+    ] = None,
 ) -> Iterator[
     Tuple[
         Callable[[], Optional[Message]],

flwr_nightly-1.9.0.dev20240424/src/py/flwr/client/grpc_rere_client/client_interceptor.py

@@ -0,0 +1,150 @@
+# Copyright 2024 Flower Labs GmbH. All Rights Reserved.
+#
+# 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.
+# ==============================================================================
+"""Flower client interceptor."""
+
+
+import base64
+import collections
+from typing import Any, Callable, Optional, Sequence, Tuple, Union
+
+import grpc
+from cryptography.hazmat.primitives.asymmetric import ec
+
+from flwr.common.secure_aggregation.crypto.symmetric_encryption import (
+    bytes_to_public_key,
+    compute_hmac,
+    generate_shared_key,
+    public_key_to_bytes,
+)
+from flwr.proto.fleet_pb2 import (  # pylint: disable=E0611
+    CreateNodeRequest,
+    DeleteNodeRequest,
+    GetRunRequest,
+    PullTaskInsRequest,
+    PushTaskResRequest,
+)
+
+_PUBLIC_KEY_HEADER = "public-key"
+_AUTH_TOKEN_HEADER = "auth-token"
+
+Request = Union[
+    CreateNodeRequest,
+    DeleteNodeRequest,
+    PullTaskInsRequest,
+    PushTaskResRequest,
+    GetRunRequest,
+]
+
+
+def _get_value_from_tuples(
+    key_string: str, tuples: Sequence[Tuple[str, Union[str, bytes]]]
+) -> bytes:
+    value = next((value for key, value in tuples if key == key_string), "")
+    if isinstance(value, str):
+        return value.encode()
+
+    return value
+
+
+class _ClientCallDetails(
+    collections.namedtuple(
+        "_ClientCallDetails", ("method", "timeout", "metadata", "credentials")
+    ),
+    grpc.ClientCallDetails,  # type: ignore
+):
+    """Details for each client call.
+
+    The class will be passed on as the first argument in continuation function.
+    In our case, `AuthenticateClientInterceptor` adds new metadata to the construct.
+    """
+
+
+class AuthenticateClientInterceptor(grpc.UnaryUnaryClientInterceptor):  # type: ignore
+    """Client interceptor for client authentication."""
+
+    def __init__(
+        self,
+        private_key: ec.EllipticCurvePrivateKey,
+        public_key: ec.EllipticCurvePublicKey,
+    ):
+        self.private_key = private_key
+        self.public_key = public_key
+        self.shared_secret: Optional[bytes] = None
+        self.server_public_key: Optional[ec.EllipticCurvePublicKey] = None
+        self.encoded_public_key = base64.urlsafe_b64encode(
+            public_key_to_bytes(self.public_key)
+        )
+
+    def intercept_unary_unary(
+        self,
+        continuation: Callable[[Any, Any], Any],
+        client_call_details: grpc.ClientCallDetails,
+        request: Request,
+    ) -> grpc.Call:
+        """Flower client interceptor.
+
+        Intercept unary call from client and add necessary authentication header in the
+        RPC metadata.
+        """
+        metadata = []
+        postprocess = False
+        if client_call_details.metadata is not None:
+            metadata = list(client_call_details.metadata)
+
+        # Always add the public key header
+        metadata.append(
+            (
+                _PUBLIC_KEY_HEADER,
+                self.encoded_public_key,
+            )
+        )
+
+        if isinstance(request, CreateNodeRequest):
+            postprocess = True
+        elif isinstance(
+            request,
+            (DeleteNodeRequest, PullTaskInsRequest, PushTaskResRequest, GetRunRequest),
+        ):
+            if self.shared_secret is None:
+                raise RuntimeError("Failure to compute hmac")
+
+            metadata.append(
+                (
+                    _AUTH_TOKEN_HEADER,
+                    base64.urlsafe_b64encode(
+                        compute_hmac(
+                            self.shared_secret, request.SerializeToString(True)
+                        )
+                    ),
+                )
+            )
+
+        client_call_details = _ClientCallDetails(
+            client_call_details.method,
+            client_call_details.timeout,
+            metadata,
+            client_call_details.credentials,
+        )
+
+        response = continuation(client_call_details, request)
+        if postprocess:
+            server_public_key_bytes = base64.urlsafe_b64decode(
+                _get_value_from_tuples(_PUBLIC_KEY_HEADER, response.initial_metadata())
+            )
+            self.server_public_key = bytes_to_public_key(server_public_key_bytes)
+            self.shared_secret = generate_shared_key(
+                self.private_key, self.server_public_key
+            )
+        return response

{flwr_nightly-1.9.0.dev20240422 → flwr_nightly-1.9.0.dev20240424}/src/py/flwr/client/grpc_rere_client/connection.py

@@ -21,7 +21,10 @@ from contextlib import contextmanager
 from copy import copy
 from logging import DEBUG, ERROR
 from pathlib import Path
-from typing import Callable, Iterator, Optional, Tuple, Union, cast
+from typing import Callable, Iterator, Optional, Sequence, Tuple, Union, cast
+
+import grpc
+from cryptography.hazmat.primitives.asymmetric import ec
 
 from flwr.client.heartbeat import start_ping_loop
 from flwr.client.message_handler.message_handler import validate_out_message
@@ -52,6 +55,8 @@ from flwr.proto.fleet_pb2_grpc import FleetStub  # pylint: disable=E0611
 from flwr.proto.node_pb2 import Node  # pylint: disable=E0611
 from flwr.proto.task_pb2 import TaskIns  # pylint: disable=E0611
 
+from .client_interceptor import AuthenticateClientInterceptor
+
 
 def on_channel_state_change(channel_connectivity: str) -> None:
     """Log channel connectivity."""
@@ -59,12 +64,15 @@ def on_channel_state_change(channel_connectivity: str) -> None:
 
 
 @contextmanager
-def grpc_request_response(  # pylint: disable=R0914, R0915
+def grpc_request_response(  # pylint: disable=R0913, R0914, R0915
     server_address: str,
     insecure: bool,
     retry_invoker: RetryInvoker,
     max_message_length: int = GRPC_MAX_MESSAGE_LENGTH,  # pylint: disable=W0613
     root_certificates: Optional[Union[bytes, str]] = None,
+    authentication_keys: Optional[
+        Tuple[ec.EllipticCurvePrivateKey, ec.EllipticCurvePublicKey]
+    ] = None,
 ) -> Iterator[
     Tuple[
         Callable[[], Optional[Message]],
@@ -109,11 +117,18 @@ def grpc_request_response(  # pylint: disable=R0914, R0915
     if isinstance(root_certificates, str):
         root_certificates = Path(root_certificates).read_bytes()
 
+    interceptors: Optional[Sequence[grpc.UnaryUnaryClientInterceptor]] = None
+    if authentication_keys is not None:
+        interceptors = AuthenticateClientInterceptor(
+            authentication_keys[0], authentication_keys[1]
+        )
+
     channel = create_channel(
         server_address=server_address,
         insecure=insecure,
         root_certificates=root_certificates,
         max_message_length=max_message_length,
+        interceptors=interceptors,
     )
     channel.subscribe(on_channel_state_change)
 

{flwr_nightly-1.9.0.dev20240422 → flwr_nightly-1.9.0.dev20240424}/src/py/flwr/client/rest_client/connection.py

@@ -23,6 +23,7 @@ from copy import copy
 from logging import ERROR, INFO, WARN
 from typing import Callable, Iterator, Optional, Tuple, Type, TypeVar, Union
 
+from cryptography.hazmat.primitives.asymmetric import ec
 from google.protobuf.message import Message as GrpcMessage
 
 from flwr.client.heartbeat import start_ping_loop
@@ -74,7 +75,7 @@ T = TypeVar("T", bound=GrpcMessage)
 
 
 @contextmanager
-def http_request_response(  # pylint: disable=R0914, R0915
+def http_request_response(  # pylint: disable=,R0913, R0914, R0915
     server_address: str,
     insecure: bool,  # pylint: disable=unused-argument
     retry_invoker: RetryInvoker,
@@ -82,6 +83,9 @@ def http_request_response(  # pylint: disable=R0914, R0915
     root_certificates: Optional[
         Union[bytes, str]
     ] = None,  # pylint: disable=unused-argument
+    authentication_keys: Optional[  # pylint: disable=unused-argument
+        Tuple[ec.EllipticCurvePrivateKey, ec.EllipticCurvePublicKey]
+    ] = None,
 ) -> Iterator[
     Tuple[
         Callable[[], Optional[Message]],

{flwr_nightly-1.9.0.dev20240422 → flwr_nightly-1.9.0.dev20240424}/src/py/flwr/common/grpc.py

@@ -16,7 +16,7 @@
 
 
 from logging import DEBUG
-from typing import Optional
+from typing import Optional, Sequence
 
 import grpc
 
@@ -30,6 +30,7 @@ def create_channel(
     insecure: bool,
     root_certificates: Optional[bytes] = None,
     max_message_length: int = GRPC_MAX_MESSAGE_LENGTH,
+    interceptors: Optional[Sequence[grpc.UnaryUnaryClientInterceptor]] = None,
 ) -> grpc.Channel:
     """Create a gRPC channel, either secure or insecure."""
     # Check for conflicting parameters
@@ -57,4 +58,7 @@ def create_channel(
         )
         log(DEBUG, "Opened secure gRPC connection using certificates")
 
+    if interceptors is not None:
+        channel = grpc.intercept_channel(channel, interceptors)
+
     return channel

{flwr_nightly-1.9.0.dev20240422 → flwr_nightly-1.9.0.dev20240424}/src/py/flwr/common/secure_aggregation/crypto/symmetric_encryption.py

@@ -18,8 +18,9 @@
 import base64
 from typing import Tuple, cast
 
+from cryptography.exceptions import InvalidSignature
 from cryptography.fernet import Fernet
-from cryptography.hazmat.primitives import hashes, serialization
+from cryptography.hazmat.primitives import hashes, hmac, serialization
 from cryptography.hazmat.primitives.asymmetric import ec
 from cryptography.hazmat.primitives.kdf.hkdf import HKDF
 
@@ -98,3 +99,21 @@ def decrypt(key: bytes, ciphertext: bytes) -> bytes:
     # The input key must be url safe
     fernet = Fernet(key)
     return fernet.decrypt(ciphertext)
+
+
+def compute_hmac(key: bytes, message: bytes) -> bytes:
+    """Compute hmac of a message using key as hash."""
+    computed_hmac = hmac.HMAC(key, hashes.SHA256())
+    computed_hmac.update(message)
+    return computed_hmac.finalize()
+
+
+def verify_hmac(key: bytes, message: bytes, hmac_value: bytes) -> bool:
+    """Verify hmac of a message using key as hash."""
+    computed_hmac = hmac.HMAC(key, hashes.SHA256())
+    computed_hmac.update(message)
+    try:
+        computed_hmac.verify(hmac_value)
+        return True
+    except InvalidSignature:
+        return False

{flwr_nightly-1.9.0.dev20240422 → flwr_nightly-1.9.0.dev20240424}/src/py/flwr/server/superlink/driver/driver_servicer.py

@@ -64,7 +64,7 @@ class DriverServicer(driver_pb2_grpc.DriverServicer):
         """Create run ID."""
         log(INFO, "DriverServicer.CreateRun")
         state: State = self.state_factory.state()
-        run_id = state.create_run("None/None", "None")
+        run_id = state.create_run(request.fab_id, request.fab_version)
         return CreateRunResponse(run_id=run_id)
 
     def PushTaskIns(

{flwr_nightly-1.9.0.dev20240422 → flwr_nightly-1.9.0.dev20240424}/src/py/flwr/server/superlink/state/in_memory_state.py

@@ -30,7 +30,7 @@ from flwr.server.utils import validate_task_ins_or_res
 from .utils import make_node_unavailable_taskres
 
 
-class InMemoryState(State):
+class InMemoryState(State):  # pylint: disable=R0902
     """In-memory State implementation."""
 
     def __init__(self) -> None:
@@ -40,6 +40,9 @@ class InMemoryState(State):
         self.run_ids: Dict[int, Tuple[str, str]] = {}
         self.task_ins_store: Dict[UUID, TaskIns] = {}
         self.task_res_store: Dict[UUID, TaskRes] = {}
+        self.client_public_keys: Set[bytes] = set()
+        self.server_public_key: Optional[bytes] = None
+        self.server_private_key: Optional[bytes] = None
         self.lock = threading.Lock()
 
     def store_task_ins(self, task_ins: TaskIns) -> Optional[UUID]:
@@ -251,6 +254,39 @@ class InMemoryState(State):
         log(ERROR, "Unexpected run creation failure.")
         return 0
 
+    def store_server_public_private_key(
+        self, public_key: bytes, private_key: bytes
+    ) -> None:
+        """Store `server_public_key` and `server_private_key` in state."""
+        with self.lock:
+            if self.server_private_key is None and self.server_public_key is None:
+                self.server_private_key = private_key
+                self.server_public_key = public_key
+            else:
+                raise RuntimeError("Server public and private key already set")
+
+    def get_server_private_key(self) -> Optional[bytes]:
+        """Retrieve `server_private_key` in urlsafe bytes."""
+        return self.server_private_key
+
+    def get_server_public_key(self) -> Optional[bytes]:
+        """Retrieve `server_public_key` in urlsafe bytes."""
+        return self.server_public_key
+
+    def store_client_public_keys(self, public_keys: Set[bytes]) -> None:
+        """Store a set of `client_public_keys` in state."""
+        with self.lock:
+            self.client_public_keys = public_keys
+
+    def store_client_public_key(self, public_key: bytes) -> None:
+        """Store a `client_public_key` in state."""
+        with self.lock:
+            self.client_public_keys.add(public_key)
+
+    def get_client_public_keys(self) -> Set[bytes]:
+        """Retrieve all currently stored `client_public_keys` as a set."""
+        return self.client_public_keys
+
     def get_run(self, run_id: int) -> Tuple[int, str, str]:
         """Retrieve information about the run with the specified `run_id`."""
         with self.lock:

{flwr_nightly-1.9.0.dev20240422 → flwr_nightly-1.9.0.dev20240424}/src/py/flwr/server/superlink/state/sqlite_state.py

@@ -20,7 +20,7 @@ import re
 import sqlite3
 import time
 from logging import DEBUG, ERROR
-from typing import Any, Dict, List, Optional, Set, Tuple, Union, cast
+from typing import Any, Dict, List, Optional, Sequence, Set, Tuple, Union, cast
 from uuid import UUID, uuid4
 
 from flwr.common import log, now
@@ -40,6 +40,19 @@ CREATE TABLE IF NOT EXISTS node(
 );
 """
 
+SQL_CREATE_TABLE_CREDENTIAL = """
+CREATE TABLE IF NOT EXISTS credential(
+    public_key BLOB PRIMARY KEY,
+    private_key BLOB
+);
+"""
+
+SQL_CREATE_TABLE_PUBLIC_KEY = """
+CREATE TABLE IF NOT EXISTS public_key(
+    public_key BLOB UNIQUE
+);
+"""
+
 SQL_CREATE_INDEX_ONLINE_UNTIL = """
 CREATE INDEX IF NOT EXISTS idx_online_until ON node (online_until);
 """
@@ -72,7 +85,6 @@ CREATE TABLE IF NOT EXISTS task_ins(
 );
 """
 
-
 SQL_CREATE_TABLE_TASK_RES = """
 CREATE TABLE IF NOT EXISTS task_res(
     task_id                 TEXT UNIQUE,
@@ -96,7 +108,7 @@ CREATE TABLE IF NOT EXISTS task_res(
 DictOrTuple = Union[Tuple[Any, ...], Dict[str, Any]]
 
 
-class SqliteState(State):
+class SqliteState(State):  # pylint: disable=R0904
     """SQLite-based state implementation."""
 
     def __init__(
@@ -134,6 +146,8 @@ class SqliteState(State):
         cur.execute(SQL_CREATE_TABLE_TASK_INS)
         cur.execute(SQL_CREATE_TABLE_TASK_RES)
         cur.execute(SQL_CREATE_TABLE_NODE)
+        cur.execute(SQL_CREATE_TABLE_CREDENTIAL)
+        cur.execute(SQL_CREATE_TABLE_PUBLIC_KEY)
         cur.execute(SQL_CREATE_INDEX_ONLINE_UNTIL)
         res = cur.execute("SELECT name FROM sqlite_schema;")
 
@@ -142,7 +156,7 @@ class SqliteState(State):
     def query(
         self,
         query: str,
-        data: Optional[Union[List[DictOrTuple], DictOrTuple]] = None,
+        data: Optional[Union[Sequence[DictOrTuple], DictOrTuple]] = None,
     ) -> List[Dict[str, Any]]:
         """Execute a SQL query."""
         if self.conn is None:
@@ -575,6 +589,59 @@ class SqliteState(State):
         log(ERROR, "Unexpected run creation failure.")
         return 0
 
+    def store_server_public_private_key(
+        self, public_key: bytes, private_key: bytes
+    ) -> None:
+        """Store `server_public_key` and `server_private_key` in state."""
+        query = "SELECT COUNT(*) FROM credential"
+        count = self.query(query)[0]["COUNT(*)"]
+        if count < 1:
+            query = (
+                "INSERT OR REPLACE INTO credential (public_key, private_key) "
+                "VALUES (:public_key, :private_key)"
+            )
+            self.query(query, {"public_key": public_key, "private_key": private_key})
+        else:
+            raise RuntimeError("Server public and private key already set")
+
+    def get_server_private_key(self) -> Optional[bytes]:
+        """Retrieve `server_private_key` in urlsafe bytes."""
+        query = "SELECT private_key FROM credential"
+        rows = self.query(query)
+        try:
+            private_key: Optional[bytes] = rows[0]["private_key"]
+        except IndexError:
+            private_key = None
+        return private_key
+
+    def get_server_public_key(self) -> Optional[bytes]:
+        """Retrieve `server_public_key` in urlsafe bytes."""
+        query = "SELECT public_key FROM credential"
+        rows = self.query(query)
+        try:
+            public_key: Optional[bytes] = rows[0]["public_key"]
+        except IndexError:
+            public_key = None
+        return public_key
+
+    def store_client_public_keys(self, public_keys: Set[bytes]) -> None:
+        """Store a set of `client_public_keys` in state."""
+        query = "INSERT INTO public_key (public_key) VALUES (?)"
+        data = [(key,) for key in public_keys]
+        self.query(query, data)
+
+    def store_client_public_key(self, public_key: bytes) -> None:
+        """Store a `client_public_key` in state."""
+        query = "INSERT INTO public_key (public_key) VALUES (:public_key)"
+        self.query(query, {"public_key": public_key})
+
+    def get_client_public_keys(self) -> Set[bytes]:
+        """Retrieve all currently stored `client_public_keys` as a set."""
+        query = "SELECT public_key FROM public_key"
+        rows = self.query(query)
+        result: Set[bytes] = {row["public_key"] for row in rows}
+        return result
+
     def get_run(self, run_id: int) -> Tuple[int, str, str]:
         """Retrieve information about the run with the specified `run_id`."""
         query = "SELECT * FROM run WHERE run_id = ?;"