File size: 6,874 Bytes
287a0bc
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
import importlib
import logging
from typing import Optional, cast, Dict, TypeVar, Any

import requests
from overrides import override
from pydantic import SecretStr
from chromadb.auth import (
    ServerAuthCredentialsProvider,
    AbstractCredentials,
    ClientAuthCredentialsProvider,
    AuthInfoType,
    ClientAuthProvider,
    ClientAuthProtocolAdapter,
    SimpleUserIdentity,
)
from chromadb.auth.registry import register_provider, resolve_provider
from chromadb.config import System
from chromadb.telemetry.opentelemetry import (
    OpenTelemetryGranularity,
    trace_method,
)

T = TypeVar("T")

logger = logging.getLogger(__name__)


class HtpasswdServerAuthCredentialsProvider(ServerAuthCredentialsProvider):
    _creds: Dict[str, SecretStr]

    def __init__(self, system: System) -> None:
        super().__init__(system)
        try:
            # Equivalent to import onnxruntime
            self.bc = importlib.import_module("bcrypt")
        except ImportError:
            raise ValueError(
                "The bcrypt python package is not installed. "
                "Please install it with `pip install bcrypt`"
            )

    @trace_method(
        "HtpasswdServerAuthCredentialsProvider.validate_credentials",
        OpenTelemetryGranularity.ALL,
    )
    @override
    def validate_credentials(self, credentials: AbstractCredentials[T]) -> bool:
        _creds = cast(Dict[str, SecretStr], credentials.get_credentials())
        if len(_creds) != 2:
            logger.error(
                "Returned credentials did match expected format: "
                "dict[username:SecretStr, password: SecretStr]"
            )
            return False
        if "username" not in _creds or "password" not in _creds:
            logger.error(
                "Returned credentials do not contain username or password")
            return False
        _usr_check = bool(
            _creds["username"].get_secret_value()
            == self._creds["username"].get_secret_value()
        )
        return _usr_check and self.bc.checkpw(
            _creds["password"].get_secret_value().encode("utf-8"),
            self._creds["password"].get_secret_value().encode("utf-8"),
        )

    @override
    def get_user_identity(
        self, credentials: AbstractCredentials[T]
    ) -> Optional[SimpleUserIdentity]:
        _creds = cast(Dict[str, SecretStr], credentials.get_credentials())
        return SimpleUserIdentity(_creds["username"].get_secret_value())


@register_provider("htpasswd_file")
class HtpasswdFileServerAuthCredentialsProvider(HtpasswdServerAuthCredentialsProvider):
    def __init__(self, system: System) -> None:
        super().__init__(system)
        system.settings.require("chroma_server_auth_credentials_file")
        _file = str(system.settings.chroma_server_auth_credentials_file)
        with open(_file, "r") as f:
            _raw_creds = [v for v in f.readline().strip().split(":")]
            self._creds = {
                "username": SecretStr(_raw_creds[0]),
                "password": SecretStr(_raw_creds[1]),
            }
        if (
            len(self._creds) != 2
            or "username" not in self._creds
            or "password" not in self._creds
        ):
            raise ValueError(
                "Invalid Htpasswd credentials found in "
                "[chroma_server_auth_credentials]. "
                "Must be <username>:<bcrypt passwd>."
            )


class HtpasswdConfigurationServerAuthCredentialsProvider(
    HtpasswdServerAuthCredentialsProvider
):
    def __init__(self, system: System) -> None:
        super().__init__(system)
        system.settings.require("chroma_server_auth_credentials")
        _raw_creds = (
            str(system.settings.chroma_server_auth_credentials).strip().split(":")
        )
        self._creds = {
            "username": SecretStr(_raw_creds[0]),
            "password": SecretStr(_raw_creds[1]),
        }
        if (
            len(self._creds) != 2
            or "username" not in self._creds
            or "password" not in self._creds
        ):
            raise ValueError(
                "Invalid Htpasswd credentials found in "
                "[chroma_server_auth_credentials]. "
                "Must be <username>:<bcrypt passwd>."
            )


class RequestsClientAuthProtocolAdapter(
    ClientAuthProtocolAdapter[requests.PreparedRequest]
):
    class _Session(requests.Session):
        _protocol_adapter: ClientAuthProtocolAdapter[requests.PreparedRequest]

        def __init__(
            self, protocol_adapter: ClientAuthProtocolAdapter[requests.PreparedRequest]
        ) -> None:
            super().__init__()
            self._protocol_adapter = protocol_adapter

        @override
        def send(
            self, request: requests.PreparedRequest, **kwargs: Any
        ) -> requests.Response:
            self._protocol_adapter.inject_credentials(request)
            return super().send(request, **kwargs)

    _session: _Session
    _auth_provider: ClientAuthProvider

    def __init__(self, system: System) -> None:
        super().__init__(system)
        system.settings.require("chroma_client_auth_provider")
        self._auth_provider = cast(
            ClientAuthProvider,
            system.require(
                resolve_provider(
                    str(system.settings.chroma_client_auth_provider), ClientAuthProvider
                ),
            ),
        )
        self._session = self._Session(self)
        self._auth_header = self._auth_provider.authenticate()

    @property
    def session(self) -> requests.Session:
        return self._session

    @override
    def inject_credentials(self, injection_context: requests.PreparedRequest) -> None:
        if self._auth_header.get_auth_info_type() == AuthInfoType.HEADER:
            _header_info = self._auth_header.get_auth_info()
            if isinstance(_header_info, tuple):
                injection_context.headers[_header_info[0]] = _header_info[
                    1
                ].get_secret_value()
            else:
                for header in _header_info:
                    injection_context.headers[header[0]
                                              ] = header[1].get_secret_value()
        else:
            raise ValueError(
                f"Unsupported auth type: {self._auth_header.get_auth_info_type()}"
            )


class ConfigurationClientAuthCredentialsProvider(
    ClientAuthCredentialsProvider[SecretStr]
):
    _creds: SecretStr

    def __init__(self, system: System) -> None:
        super().__init__(system)
        system.settings.require("chroma_client_auth_credentials")
        self._creds = SecretStr(
            str(system.settings.chroma_client_auth_credentials))

    @override
    def get_credentials(self) -> SecretStr:
        return self._creds