Skip to content

Commit

Permalink
feat: Implement encryption in the client SDK's login function (#2054)
Browse files Browse the repository at this point in the history
<!--
Please precisely, concisely, and concretely describe what this PR changes, the rationale behind codes,
and how it affects the users and other developers.
-->

**Checklist:** (if applicable)

- [ ] Milestone metadata specifying the target backport version
- [ ] Mention to the original issue
- [ ] Installer updates including:
  - Fixtures for db schema changes
  - New mandatory config options
- [ ] Update of end-to-end CLI integration tests in `ai.backend.test`
- [ ] API server-client counterparts (e.g., manager API -> client SDK)
- [ ] Test case(s) to:
  - Demonstrate the difference of before/after
  - Demonstrate the flow of abstract/conceptual models with a concrete implementation
- [ ] Documentation
  - Contents in the `docs` directory
  - docstrings in public interfaces and type annotations
  • Loading branch information
achimnol committed Apr 23, 2024
1 parent bd83825 commit d605461
Showing 1 changed file with 23 additions and 8 deletions.
31 changes: 23 additions & 8 deletions src/ai/backend/client/func/auth.py
Original file line number Diff line number Diff line change
@@ -1,11 +1,23 @@
from typing import Optional
import json
from typing import Any, Optional

from ..auth import encrypt_payload
from ..request import Request
from .base import BaseFunction, api_function

__all__ = ("Auth",)


def _put_secure_body(rqst: Request, data: Any) -> None:
if rqst.config.endpoint.scheme == "https":
rqst.set_json(data)
else:
rqst.headers["X-BackendAI-Encoded"] = "true"
raw_body = json.dumps(data).encode()
encoded_body = encrypt_payload(str(rqst.config.endpoint), raw_body)
rqst.set_content(encoded_body)


class Auth(BaseFunction):
"""
Provides the function interface for login session management and authorization.
Expand All @@ -27,7 +39,7 @@ async def login(cls, user_id: str, password: str, otp: Optional[str] = None) ->
}
if otp:
body["otp"] = otp
rqst.set_json(body)
_put_secure_body(rqst, body)
async with rqst.fetch(anonymous=True) as resp:
data = await resp.json()
data["cookies"] = resp.raw_response.cookies
Expand Down Expand Up @@ -56,11 +68,12 @@ async def update_password(
Update user's password. This API works only for account owner.
"""
rqst = Request("POST", "/auth/update-password")
rqst.set_json({
body = {
"old_password": old_password,
"new_password": new_password,
"new_password2": new_password2,
})
}
_put_secure_body(rqst, body)
async with rqst.fetch() as resp:
return await resp.json()

Expand All @@ -75,12 +88,13 @@ async def update_password_no_auth(
"""

rqst = Request("POST", "/auth/update-password-no-auth")
rqst.set_json({
body = {
"domain": domain,
"username": user_id,
"current_password": current_password,
"new_password": new_password,
})
}
_put_secure_body(rqst, body)
async with rqst.fetch(anonymous=True) as resp:
return await resp.json()

Expand All @@ -95,10 +109,11 @@ async def update_password_no_auth_in_session(
"""

rqst = Request("POST", "/server/update-password-no-auth")
rqst.set_json({
body = {
"username": user_id,
"current_password": current_password,
"new_password": new_password,
})
}
_put_secure_body(rqst, body)
async with rqst.fetch(anonymous=True) as resp:
return await resp.json()

0 comments on commit d605461

Please sign in to comment.