Add support for user-provided SSH public keys

Part-of: #3644

Author: un-def

src/dstack/_internal/core/models/keys.py

@@ -0,0 +1,12 @@
+import datetime
+import uuid
+
+from dstack._internal.core.models.common import CoreModel
+
+
+class PublicKeyInfo(CoreModel):
+    id: uuid.UUID
+    added_at: datetime.datetime
+    name: str
+    type: str
+    fingerprint: str

src/dstack/_internal/server/app.py

@@ -42,6 +42,7 @@
     metrics,
     projects,
     prometheus,
+    public_keys,
     repos,
     runs,
     secrets,
@@ -259,6 +260,7 @@ def register_routes(app: FastAPI, ui: bool = True):
     app.include_router(exports.project_router)
     app.include_router(imports.project_router)
     app.include_router(sshproxy.router)
+    app.include_router(public_keys.router)
 
     @app.exception_handler(ForbiddenError)
     async def forbidden_error_handler(request: Request, exc: ForbiddenError):

src/dstack/_internal/server/migrations/versions/2026/03_24_1145_59e328ced74c_add_userpublickeymodel.py

@@ -0,0 +1,50 @@
+"""Add UserPublicKeyModel
+
+Revision ID: 59e328ced74c
+Revises: c1c2ecaee45c
+Create Date: 2026-03-24 11:45:13.560594+00:00
+
+"""
+
+import sqlalchemy as sa
+import sqlalchemy_utils
+from alembic import op
+
+import dstack._internal.server.models
+
+# revision identifiers, used by Alembic.
+revision = "59e328ced74c"
+down_revision = "c1c2ecaee45c"
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.create_table(
+        "user_public_keys",
+        sa.Column("id", sqlalchemy_utils.types.uuid.UUIDType(binary=False), nullable=False),
+        sa.Column("created_at", dstack._internal.server.models.NaiveDateTime(), nullable=False),
+        sa.Column("user_id", sqlalchemy_utils.types.uuid.UUIDType(binary=False), nullable=False),
+        sa.Column("name", sa.String(length=100), nullable=False),
+        sa.Column("type", sa.String(length=100), nullable=False),
+        sa.Column("fingerprint", sa.String(length=100), nullable=False),
+        sa.Column("key", sa.Text(), nullable=False),
+        sa.ForeignKeyConstraint(
+            ["user_id"],
+            ["users.id"],
+            name=op.f("fk_user_public_keys_user_id_users"),
+            ondelete="CASCADE",
+        ),
+        sa.PrimaryKeyConstraint("id", name=op.f("pk_user_public_keys")),
+        sa.UniqueConstraint(
+            "user_id", "fingerprint", name="uq_user_public_keys_user_id_fingerprint"
+        ),
+    )
+    # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.drop_table("user_public_keys")
+    # ### end Alembic commands ###

src/dstack/_internal/server/models.py

@@ -1115,3 +1115,24 @@ class ExportedFleetModel(BaseModel):
         ForeignKey("fleets.id", ondelete="CASCADE"), index=True
     )
     fleet: Mapped["FleetModel"] = relationship()
+
+
+class UserPublicKeyModel(BaseModel):
+    __tablename__ = "user_public_keys"
+    __table_args__ = (
+        UniqueConstraint("user_id", "fingerprint", name="uq_user_public_keys_user_id_fingerprint"),
+    )
+
+    id: Mapped[uuid.UUID] = mapped_column(
+        UUIDType(binary=False), primary_key=True, default=uuid.uuid4
+    )
+    created_at: Mapped[datetime] = mapped_column(NaiveDateTime, default=get_current_datetime)
+    user_id: Mapped[uuid.UUID] = mapped_column(ForeignKey("users.id", ondelete="CASCADE"))
+    user: Mapped["UserModel"] = relationship()
+    name: Mapped[str] = mapped_column(String(100))
+    type: Mapped[str] = mapped_column(String(100))
+    """`type` is a key type identifier used by OpenSSH, e.g., `ssh-rsa`, `ecdsa-sha2-nistp521`."""
+    fingerprint: Mapped[str] = mapped_column(String(100))
+    """`fingerprint` stores a key digest in the format used by OpenSSH: `SHA256:<base64>`."""
+    key: Mapped[str] = mapped_column(Text)
+    """`key` stores a public key in the OpenSSH disk (ASCII-armored) format."""

src/dstack/_internal/server/routers/public_keys.py

@@ -0,0 +1,54 @@
+from typing import Annotated
+
+from fastapi import APIRouter, Depends
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from dstack._internal.core.models.keys import PublicKeyInfo
+from dstack._internal.server.db import get_session
+from dstack._internal.server.models import UserModel
+from dstack._internal.server.schemas.public_keys import (
+    AddPublicKeyRequest,
+    DeletePublicKeysRequest,
+)
+from dstack._internal.server.security.permissions import Authenticated
+from dstack._internal.server.services import public_keys as public_keys_services
+from dstack._internal.server.utils.routers import (
+    CustomORJSONResponse,
+    get_base_api_additional_responses,
+)
+
+router = APIRouter(
+    prefix="/api/users/public_keys",
+    tags=["user public keys"],
+    responses=get_base_api_additional_responses(),
+)
+
+
+@router.post("/list", response_model=list[PublicKeyInfo])
+async def list_user_public_keys(
+    session: Annotated[AsyncSession, Depends(get_session)],
+    user: Annotated[UserModel, Depends(Authenticated())],
+):
+    public_keys = await public_keys_services.list_user_public_keys(session=session, user=user)
+    return CustomORJSONResponse(public_keys)
+
+
+@router.post("/add", response_model=PublicKeyInfo)
+async def add_user_public_key(
+    body: AddPublicKeyRequest,
+    session: Annotated[AsyncSession, Depends(get_session)],
+    user: Annotated[UserModel, Depends(Authenticated())],
+):
+    public_key = await public_keys_services.add_user_public_key(
+        session=session, user=user, key=body.key, name=body.name
+    )
+    return CustomORJSONResponse(public_key)
+
+
+@router.post("/delete")
+async def delete_user_public_keys(
+    body: DeletePublicKeysRequest,
+    session: Annotated[AsyncSession, Depends(get_session)],
+    user: Annotated[UserModel, Depends(Authenticated())],
+):
+    await public_keys_services.delete_user_public_keys(session=session, user=user, ids=body.ids)

src/dstack/_internal/server/schemas/public_keys.py

@@ -0,0 +1,13 @@
+import uuid
+from typing import Optional
+
+from dstack._internal.core.models.common import CoreModel
+
+
+class AddPublicKeyRequest(CoreModel):
+    key: str
+    name: Optional[str] = None
+
+
+class DeletePublicKeysRequest(CoreModel):
+    ids: list[uuid.UUID]