Repository: edgedb/edgedb

Issue:
Add support for SCRAM auth to HTTP-tunneled binary protocol



Discussion:
cc @jaclarke 
please

Resolving patch:
diff --git a/edb/server/protocol/auth/__init__.py b/edb/server/protocol/auth/__init__.py
new file mode 100644
--- /dev/null
+++ b/edb/server/protocol/auth/__init__.py
@@ -0,0 +1,80 @@
+#
+# This source file is part of the EdgeDB open source project.
+#
+# Copyright 2022-present MagicStack Inc. and the EdgeDB authors.
+#
+# 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.
+#
+
+
+import http
+import json
+
+from edb import errors
+from edb.common import debug
+from edb.common import markup
+
+from . import scram
+
+
+async def handle_request(request, response, path_parts, server):
+    try:
+        if path_parts == ["token"]:
+            if not request.authorization:
+                response.status = http.HTTPStatus.UNAUTHORIZED
+                response.custom_headers["WWW-Authenticate"] = "SCRAM-SHA-256"
+                return
+
+            scheme, _, auth_str = request.authorization.decode(
+                "ascii"
+            ).partition(" ")
+
+            if scheme.lower().startswith("scram"):
+                scram.handle_request(scheme, auth_str, response, server)
+            else:
+                response.body = b"Unsupported authentication scheme"
+                response.status = http.HTTPStatus.UNAUTHORIZED
+                response.custom_headers["WWW-Authenticate"] = "SCRAM-SHA-256"
+                response.close_connection = True
+        else:
+            response.body = b"Unknown path"
+            response.status = http.HTTPStatus.NOT_FOUND
+            response.close_connection = True
+    except errors.EdgeDBError as ex:
+        if debug.flags.server:
+            markup.dump(ex)
+        _response_error(
+            response, http.HTTPStatus.INTERNAL_SERVER_ERROR, str(ex), type(ex)
+        )
+    except Exception as ex:
+        if debug.flags.server:
+            markup.dump(ex)
+
+        # XXX Fix this when LSP "location" objects are implemented
+        ex_type = errors.InternalServerError
+
+        _response_error(
+            response, http.HTTPStatus.INTERNAL_SERVER_ERROR, str(ex), ex_type
+        )
+
+
+def _response_error(response, status, message, ex_type):
+    err_dct = {
+        "message": message,
+        "type": str(ex_type.__name__),
+        "code": ex_type.get_code(),
+    }
+
+    response.body = json.dumps({"error": err_dct}).encode()
+    response.status = status
+    response.close_connection = True
diff --git a/edb/server/protocol/auth/scram.py b/edb/server/protocol/auth/scram.py
new file mode 100644
--- /dev/null
+++ b/edb/server/protocol/auth/scram.py
@@ -0,0 +1,310 @@
+#
+# This source file is part of the EdgeDB open source project.
+#
+# Copyright 2022-present MagicStack Inc. and the EdgeDB authors.
+#
+# 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.
+#
+
+import base64
+import collections
+import hashlib
+import http
+import os
+import time
+from jwcrypto import jwt
+
+from edgedb import scram
+
+from edb.common import debug
+from edb.common import markup
+
+
+SESSION_TIMEOUT = 30
+SESSION_HIGH_WATER_MARK = SESSION_TIMEOUT * 10
+sessions: collections.OrderedDict[str, tuple] = collections.OrderedDict()
+
+
+def handle_request(scheme, auth_str, response, server):
+    if scheme != "SCRAM-SHA-256":
+        response.body = (
+            b"Client selected an invalid SASL authentication mechanism"
+        )
+        response.status = http.HTTPStatus.UNAUTHORIZED
+        response.custom_headers["WWW-Authenticate"] = "SCRAM-SHA-256"
+        return
+
+    data = None
+    sid = None
+    try:
+        for kv_str in auth_str.split():
+            key, _, value = kv_str.rstrip(",").partition("=")
+            if key == "data":
+                data = base64.b64decode(value.strip('"')).strip()
+            elif key == "sid":
+                sid = value.strip('"')
+        if data is None:
+            raise ValueError("Malformed SCRAM message: data is missing")
+    except Exception as ex:
+        if debug.flags.server:
+            markup.dump(ex)
+        response.body = str(ex).encode("ascii")
+        response.status = http.HTTPStatus.BAD_REQUEST
+        response.close_connection = True
+        return
+
+    if (
+        not server.get_jwe_key().has_public
+        or not server.get_jws_key().has_private
+    ):
+        response.body = b"Server doesn't support HTTP SCRAM authentication"
+        response.status = http.HTTPStatus.FORBIDDEN
+        response.close_connection = True
+        return
+
+    if sid is None:
+        try:
+            (
+                bare_offset,
+                cb_flag,
+                authzid,
+                username,
+                client_nonce,
+            ) = scram.parse_client_first_message(data)
+        except ValueError as ex:
+            if debug.flags.server:
+                markup.dump(ex)
+            response.body = f"Bad client first message: {ex!s}".encode("ascii")
+            response.status = http.HTTPStatus.BAD_REQUEST
+            response.close_connection = True
+            return
+
+        username = username.decode("utf-8")
+        client_first_bare = data[bare_offset:]
+
+        if isinstance(cb_flag, str):
+            response.body = (
+                b"Malformed SCRAM message: "
+                b"The client selected SCRAM-SHA-256 without "
+                b"channel binding, but the SCRAM message "
+                b"includes channel binding data."
+            )
+            response.status = http.HTTPStatus.BAD_REQUEST
+            response.close_connection = True
+            return
+
+        if authzid:
+            response.body = (
+                b"Client uses SASL authorization identity, "
+                b"which is not supported"
+            )
+            response.status = http.HTTPStatus.BAD_REQUEST
+            response.close_connection = True
+            return
+
+        try:
+            verifier, mock_auth = get_scram_verifier(username, server)
+        except ValueError as ex:
+            if debug.flags.server:
+                markup.dump(ex)
+            response.body = b"Authentication failed"
+            response.status = http.HTTPStatus.UNAUTHORIZED
+            response.custom_headers["WWW-Authenticate"] = "SCRAM-SHA-256"
+            return
+
+        server_nonce = scram.generate_nonce()
+        server_first = scram.build_server_first_message(
+            server_nonce, client_nonce, verifier.salt, verifier.iterations
+        ).encode("utf-8")
+
+        if len(sessions) > SESSION_HIGH_WATER_MARK:
+            while sessions:
+                key, value = sessions.popitem(last=False)
+                if value[0] + SESSION_TIMEOUT > time.monotonic():
+                    sessions[key] = value
+                    sessions.move_to_end(key, last=False)
+                    break
+
+        sid = (
+            base64.urlsafe_b64encode(os.urandom(16))
+            .decode("ascii")
+            .rstrip("=")
+        )
+        assert sid not in sessions
+        sessions[sid] = (
+            time.monotonic(),
+            client_nonce,
+            server_nonce,
+            client_first_bare,
+            cb_flag,
+            server_first,
+            verifier,
+            mock_auth,
+            username,
+        )
+
+        server_first = base64.b64encode(server_first).decode("ascii")
+        response.status = http.HTTPStatus.UNAUTHORIZED
+        response.custom_headers[
+            "WWW-Authenticate"
+        ] = f"SCRAM-SHA-256 sid={sid}, data={server_first}"
+
+    else:
+        session = sessions.pop(sid)
+        if session is None:
+            response.body = b"Bad session ID"
+            response.status = http.HTTPStatus.UNAUTHORIZED
+            response.custom_headers["WWW-Authenticate"] = "SCRAM-SHA-256"
+            return
+
+        (
+            ts,
+            client_nonce,
+            server_nonce,
+            client_first_bare,
+            cb_flag,
+            server_first,
+            verifier,
+            mock_auth,
+            username,
+        ) = session
+        if ts + SESSION_TIMEOUT < time.monotonic():
+            response.body = b"Session timed out"
+            response.status = http.HTTPStatus.UNAUTHORIZED
+            response.custom_headers["WWW-Authenticate"] = "SCRAM-SHA-256"
+            return
+
+        try:
+            (
+                cb_data,
+                client_proof,
+                proof_len,
+            ) = scram.parse_client_final_message(
+                data, client_nonce, server_nonce
+            )
+        except ValueError as ex:
+            if debug.flags.server:
+                markup.dump(ex)
+            response.body = f"Bad client final message: {ex!s}".encode("ascii")
+            response.status = http.HTTPStatus.BAD_REQUEST
+            response.close_connection = True
+            return
+
+        client_final_without_proof = data[:-proof_len]
+
+        cb_data_ok = (cb_flag is False and cb_data == b"biws") or (
+            cb_flag is True and cb_data == b"eSws"
+        )
+        if not cb_data_ok:
+            response.body = (
+                b"Malformed SCRAM message: "
+                b"Unexpected SCRAM channel-binding attribute "
+                b"in client-final-message."
+            )
+            response.status = http.HTTPStatus.BAD_REQUEST
+            response.close_connection = True
+            return
+
+        if (
+            not scram.verify_client_proof(
+                client_first_bare,
+                server_first,
+                client_final_without_proof,
+                verifier.stored_key,
+                client_proof,
+            )
+            or mock_auth
+        ):
+            response.body = b"Authentication failed"
+            response.status = http.HTTPStatus.UNAUTHORIZED
+            response.custom_headers["WWW-Authenticate"] = "SCRAM-SHA-256"
+            return
+
+        server_final = base64.b64encode(
+            scram.build_server_final_message(
+                client_first_bare,
+                server_first,
+                client_final_without_proof,
+                verifier.server_key,
+            ).encode("utf-8")
+        ).decode("ascii")
+
+        try:
+            response.body = generate_jwt_token(username, server)
+        except ValueError as ex:
+            if debug.flags.server:
+                markup.dump(ex)
+            response.body = b"Authentication failed"
+            response.status = http.HTTPStatus.UNAUTHORIZED
+            response.custom_headers["WWW-Authenticate"] = "SCRAM-SHA-256"
+            return
+
+        response.custom_headers[
+            "Authentication-Info"
+        ] = f"sid={sid}, data={server_final}"
+
+
+def get_scram_verifier(user, server):
+    roles = server.get_roles()
+
+    rolerec = roles.get(user)
+    if rolerec is not None:
+        verifier_string = rolerec["password"]
+        if verifier_string is not None:
+            verifier = scram.parse_verifier(verifier_string)
+            is_mock = False
+            return verifier, is_mock
+
+    # To avoid revealing the validity of the submitted user name,
+    # generate a mock verifier using a salt derived from the
+    # received user name and the cluster mock auth nonce.
+    # The same approach is taken by Postgres.
+    nonce = server.get_instance_data("mock_auth_nonce")
+    salt = hashlib.sha256(nonce.encode() + user.encode()).digest()
+
+    verifier = scram.SCRAMVerifier(
+        mechanism="SCRAM-SHA-256",
+        iterations=scram.DEFAULT_ITERATIONS,
+        salt=salt[: scram.DEFAULT_SALT_LENGTH],
+        stored_key=b"",
+        server_key=b"",
+    )
+    is_mock = True
+    return verifier, is_mock
+
+
+def generate_jwt_token(user, server):
+    rolerec = server.get_roles().get(user)
+    if rolerec is None or rolerec["password"] is None:
+        raise ValueError(f"invalid user: {user!r}")
+    ekey = server.get_jwe_key()
+    skey = server.get_jws_key()
+
+    namespace = "edgedb.server"
+    token = jwt.JWT(
+        header={"alg": "ES256" if skey["kty"] == "EC" else "RS256"},
+        claims={
+            f"{namespace}.roles": {user: rolerec["password"]},
+            "iat": int(time.time()),
+        },
+    )
+    token.make_signed_token(skey)
+    encrypted_token = jwt.JWT(
+        header={
+            "alg": "ECDH-ES" if ekey["kty"] == "EC" else "RSA-OAEP-256",
+            "enc": "A256GCM",
+        },
+        claims=token.serialize(),
+    )
+    encrypted_token.make_encrypted_token(ekey)
+    return encrypted_token.serialize().encode("ascii")


Test patch:
diff --git a/tests/test_http_auth.py b/tests/test_http_auth.py
new file mode 100644
--- /dev/null
+++ b/tests/test_http_auth.py
@@ -0,0 +1,214 @@
+#
+# This source file is part of the EdgeDB open source project.
+#
+# Copyright 2022-present MagicStack Inc. and the EdgeDB authors.
+#
+# 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.
+#
+
+
+import base64
+import struct
+
+import edgedb
+from edgedb import scram
+
+from edb import protocol
+from edb.testbase import http as tb
+from edb.testbase import server as tb_server
+
+
+class BaseTestHttpAuth(tb.BaseHttpTest, tb_server.ConnectedTestCase):
+    @classmethod
+    def get_api_path(cls) -> str:
+        return "/auth"
+
+    def _scram_auth(self, user, password):
+        with self.http_con() as con:
+            _, headers, status = self.http_con_request(con, {}, path="token")
+            self.assertEqual(status, 401)
+            self.assertEqual(
+                headers, headers | {"www-authenticate": "scram-sha-256"}
+            )
+
+        client_nonce = scram.generate_nonce()
+        client_first, client_first_bare = scram.build_client_first_message(
+            client_nonce, user
+        )
+        client_first_b64 = base64.b64encode(
+            client_first.encode("ascii")
+        ).decode("ascii")
+
+        with self.http_con() as con:
+            con.request(
+                "GET",
+                "/auth/token",
+                headers={
+                    "Authorization": f"SCRAM-SHA-256 data={client_first_b64}"
+                },
+            )
+            resp = con.getresponse()
+            headers = {k.lower(): v for k, v in resp.getheaders()}
+            self.assertEqual(resp.status, 401)
+
+        scheme, _, data = headers["www-authenticate"].partition(" ")
+        self.assertEqual(scheme, "SCRAM-SHA-256")
+
+        values = {}
+        for kv_str in data.split():
+            key, _, value = kv_str.rstrip(",").partition("=")
+            values[key] = value
+
+        self.assertIn("sid", values)
+        self.assertIn("data", values)
+
+        sid = values["sid"]
+        server_first = base64.b64decode(values["data"])
+
+        server_nonce, salt, itercount = scram.parse_server_first_message(
+            server_first
+        )
+
+        client_final, expected_server_sig = scram.build_client_final_message(
+            password,
+            salt,
+            itercount,
+            client_first_bare.encode("utf-8"),
+            server_first,
+            server_nonce,
+        )
+        client_final_b64 = base64.b64encode(
+            client_final.encode("ascii")
+        ).decode("ascii")
+
+        with self.http_con() as con:
+            con.request(
+                "GET",
+                "/auth/token",
+                headers={
+                    "Authorization": f"SCRAM-SHA-256 sid={sid} "
+                    f"data={client_final_b64}"
+                },
+            )
+            resp = con.getresponse()
+            content = resp.read()
+            headers = {k.lower(): v for k, v in resp.getheaders()}
+            return content, headers, resp.status, sid, expected_server_sig
+
+    def _scram_auth_expect_failure(self, user, password):
+        (
+            content,
+            headers,
+            status,
+            sid,
+            expected_server_sig,
+        ) = self._scram_auth(user, password)
+        self.assertEqual(status, 401)
+        self.assertEqual(content, b"Authentication failed")
+        self.assertEqual(
+            headers, headers | {"www-authenticate": "SCRAM-SHA-256"}
+        )
+
+
+class TestHttpAuth(BaseTestHttpAuth):
+    def test_http_auth_scram(self):
+        args = self.get_connect_args()
+        (token, headers, status, sid, expected_server_sig) = self._scram_auth(
+            args["user"], args["password"]
+        )
+        self.assertEqual(status, 200)
+        values = {}
+        for kv_str in headers["authentication-info"].split():
+            key, _, value = kv_str.rstrip(",").partition("=")
+            values[key] = value
+
+        self.assertEqual(values["sid"], sid)
+        self.assertIn("data", values)
+        server_final = base64.b64decode(values["data"])
+        server_sig = scram.parse_server_final_message(server_final)
+        self.assertEqual(server_sig, expected_server_sig)
+
+        with self.http_con() as con:
+            con.request(
+                "POST",
+                f"/db/{args['database']}",
+                body=protocol.Execute(
+                    annotations=[],
+                    allowed_capabilities=protocol.Capability.ALL,
+                    compilation_flags=protocol.CompilationFlag(0),
+                    implicit_limit=0,
+                    command_text="SELECT 42",
+                    output_format=protocol.OutputFormat.JSON,
+                    expected_cardinality=protocol.Cardinality.AT_MOST_ONE,
+                    input_typedesc_id=b"\0" * 16,
+                    output_typedesc_id=b"\0" * 16,
+                    state_typedesc_id=b"\0" * 16,
+                    arguments=b"",
+                    state_data=b"",
+                ).dump()
+                + protocol.Sync().dump(),
+                headers={
+                    "Content-Type": "application/x.edgedb.v_1_0.binary",
+                    "Authorization": f"Bearer {token.decode('ascii')}",
+                    "X-EdgeDB-User": args["user"],
+                },
+            )
+            content, headers, status = self.http_con_read_response(con)
+        self.assertEqual(status, 200)
+        self.assertEqual(
+            headers,
+            headers | {"content-type": "application/x.edgedb.v_1_0.binary"},
+        )
+        uint32_unpack = struct.Struct("!L").unpack
+        msgs = []
+        while content:
+            mtype = content[0]
+            (msize,) = uint32_unpack(content[1:5])
+            msg = protocol.ServerMessage.parse(mtype, content[5 : msize + 1])
+            msgs.append(msg)
+            content = content[msize + 1 :]
+        self.assertIsInstance(msgs[0], protocol.CommandDataDescription)
+        self.assertIsInstance(msgs[1], protocol.Data)
+        self.assertEqual(bytes(msgs[1].data[0].data), b"42")
+        self.assertIsInstance(msgs[2], protocol.CommandComplete)
+        self.assertEqual(msgs[2].status, "SELECT")
+        self.assertIsInstance(msgs[3], protocol.ReadyForCommand)
+        self.assertEqual(
+            msgs[3].transaction_state,
+            protocol.TransactionState.NOT_IN_TRANSACTION,
+        )
+
+    def test_http_auth_scram_invalid_password(self):
+        args = self.get_connect_args()
+        self._scram_auth_expect_failure(args["user"], "bad-password")
+
+    def test_http_auth_scram_no_user(self):
+        self._scram_auth_expect_failure("scram_no_user", "bad-password")
+
+
+class TestHttpAuthSystem(BaseTestHttpAuth):
+
+    PARALLELISM_GRANULARITY = "system"
+    TRANSACTION_ISOLATION = False
+
+    async def test_http_auth_scram_no_password(self):
+        if not self.has_create_role:
+            self.skipTest("create role is not supported by the backend")
+        await self.con.execute("create superuser role scram_no_pass")
+        roles = await self.con.query("SELECT sys::Role.name")
+        self.assertIn("scram_no_pass", roles)
+        with self.assertRaisesRegex(
+            edgedb.AuthenticationError, "authentication failed"
+        ):
+            await self.connect(user="scram_no_pass", password="password")
+        self._scram_auth_expect_failure("scram_no_pass", "password")