SET search_path TO documentdb_api_catalog;
SET documentdb.next_collection_id TO 1800;
SET documentdb.next_collection_index_id TO 1800;

CREATE OR REPLACE FUNCTION generate_auth_message_client_proof(
    p_user_name text,
    p_password text)
RETURNS documentdb_core.bson LANGUAGE C AS 'pg_documentdb', $$command_generate_auth_message_client_proof_for_test$$;

CREATE OR REPLACE FUNCTION generate_server_signature(
    p_user_name text,
    p_password text,
    p_auth_message text)
RETURNS documentdb_core.bson LANGUAGE C AS 'pg_documentdb', $$command_generate_server_signature_for_test$$;

-- TEST SCRAM SHA-256 authentication support functions from extension
--   Test file to test the functions:
--     1. documentdb_api_internal.scram_sha256_get_salt_and_iterations() and 
--     2. documentdb_api_internal.authenticate_with_scram_sha256()
CREATE OR REPLACE FUNCTION test_documentdb_scram_sha256_dual_api(p_user_name text,
															  p_password text)
RETURNS text 
AS $$
DECLARE
    rol_pas            text; -- Role Password from pg_authid table
    iter_ext           text; -- iterations count extracted from shadow password
    salt_ext           text; -- salt extracted from shadow password
    salt_n_iter        text; -- SALT and iterations from scram_sha256_get_salt_and_iterations()
    salt_n_iter_formed text; -- json format text with SALT and iterations formed in this test.
    auth_result        documentdb_core.bson; -- Authentication result.
    cp_n_authmsg       text; -- Client Proof and Auth message from test helper function.
    auth_message       text; -- Auth Message generated by test suite
    client_proof       text;
    result             text;
    serv_sign          text; -- Server signature sent by extension
    serv_sign_gen      text; -- Generated server signature by Test suite.
    auth_result_t      text; -- Scram Authentication result in text type
BEGIN
    -- Get salt and Iterations from Postgres for the provided User name
    SELECT documentdb_api_internal.scram_sha256_get_salt_and_iterations(p_user_name) into salt_n_iter;
    
    select rolpassword into rol_pas from pg_catalog.pg_authid where rolname = p_user_name;

	SELECT substring(rol_pas similar '%SCRAM-SHA-256$#"_+#":_+$_+:%' escape '#') into iter_ext;
	SELECT substring(rol_pas similar '%SCRAM-SHA-256$_+:#"_+#"$_+:%' escape '#') into salt_ext;
	
    salt_n_iter_formed := '{ "ok" : { "$numberInt" : "1" }, "iterations" : { "$numberInt" : "' 
                          || iter_ext || '" }, "salt" : "' || salt_ext || '" }';
	
	-- IF the salt and iterations got from scram_sha256_get_salt_and_iterations is 
	-- not not matching with the extracted ones from the shadow password then return false.
	IF salt_n_iter <> salt_n_iter_formed THEN
        RETURN 'SALT request result mismatch';
    END IF;

    SELECT generate_auth_message_client_proof(p_user_name, p_password) into cp_n_authmsg;

    SELECT substring(cp_n_authmsg similar '%"AuthMessage" : "#"_+#"", "ClientProof"%' escape '#') into auth_message;
    SELECT substring(cp_n_authmsg similar '%"AuthMessage" : "_+", "ClientProof" : "#"_+#""%' escape '#') into client_proof;

    SELECT REPLACE(auth_message, '\"', '"') into auth_message;
    SELECT REPLACE(auth_message, '\\', '\') into auth_message;
    
    SELECT documentdb_api_internal.authenticate_with_scram_sha256(p_user_name, auth_message, client_proof) into auth_result;

    result := auth_result @= '{"ok":1}';

    IF result = 'true' THEN
        SELECT auth_result into auth_result_t;
        SELECT substring(auth_result_t similar '%"ServerSignature" : "#"_+#""%' escape '#') into serv_sign;

        -- Validate Server Signature in auth_result
        --     ServerKey = HMAC(SaltedPassword, "Server Key")
        --     ServerSignature = HMAC(ServerKey, AuthMessage)
        --   This ServerSignature has to be compared against serv_sign (which is received from extension).
        --       If matched, then authentication is success.
        SELECT generate_server_signature(p_user_name, p_password, auth_message) into serv_sign_gen;
        SELECT substring(serv_sign_gen similar '%"ServerSignature" : "#"_+#""%' escape '#') into serv_sign_gen;
              
        IF serv_sign IS DISTINCT FROM serv_sign_gen THEN
            RAISE NOTICE 'Server Signature mismatch or NULL, change result to false';
            RETURN 'false';
        END IF;
        
    END IF;

    RETURN result;
	
END;
$$ 
LANGUAGE plpgsql;

/* create users */
SET client_min_messages TO ERROR;
CREATE ROLE user1      WITH LOGIN PASSWORD '<password_placeholder1>';
CREATE ROLE uSeR2      WITH LOGIN PASSWORD '<password_placeholder2>';
CREATE ROLE "User3"    WITH LOGIN PASSWORD '<password_placeholder3>';
RESET client_min_messages;

/* negative test cases for scram_sha256_get_salt_and_iterations */
SET client_min_messages = NOTICE;
-- 1. CALL scram_sha256_get_salt_and_iterations with null input parameter
SELECT documentdb_api_internal.scram_sha256_get_salt_and_iterations(null);

-- 2. CALL scram_sha256_get_salt_and_iterations with empty input parameter
SELECT documentdb_api_internal.scram_sha256_get_salt_and_iterations('');

-- 3. CALL scram_sha256_get_salt_and_iterations with non existent user name
SELECT documentdb_api_internal.scram_sha256_get_salt_and_iterations('nonexistent');

-- 4. CALL scram_sha256_get_salt_and_iterations with non existent user name with a white space in it
SELECT documentdb_api_internal.scram_sha256_get_salt_and_iterations('white space');

-- 5. CALL scram_sha256_get_salt_and_iterations with non existent quoted user name with a white space in it. Note that quoted user name is not supported by this API
SELECT documentdb_api_internal.scram_sha256_get_salt_and_iterations('"white space"');

-- 6. CALL scram_sha256_get_salt_and_iterations with non existent user name which is of length more than 64
SELECT documentdb_api_internal.scram_sha256_get_salt_and_iterations('abcdefghijklmnopqrstuvwxyabcdefghijklmnopqrstuvwxyabcdefghijklmnopqrstuvwxy');

-- 7. FAIL CASE because username is given in lower case but was created with mixed case using double quotes.
SELECT documentdb_api_internal.scram_sha256_get_salt_and_iterations('user3');

-- 8. FAIL CASE because username is given in upper case but was created with mixed case using double quotes.
SELECT documentdb_api_internal.scram_sha256_get_salt_and_iterations('USER3');

-- 9. FAIL CASE because unquotated user name will be converted to lower case and will not match with the mixed case user name.
SELECT documentdb_api_internal.scram_sha256_get_salt_and_iterations('uSeR2');

/* negative test cases for authenticate_with_scram_sha256 */
-- 1. Call with User Name as NULL
SELECT documentdb_api_internal.authenticate_with_scram_sha256(null, 'authmsg', 'client proof');

-- 2. Call with AUTH MESSAGE as NULL
SELECT documentdb_api_internal.authenticate_with_scram_sha256('user1', null, 'client proof');

-- 3. Call with User Name, AUTH MESSAGE, CLIENT PROOF as NULL
SELECT documentdb_api_internal.authenticate_with_scram_sha256(null, null, null);

-- 4. Empty user name
SELECT documentdb_api_internal.authenticate_with_scram_sha256('', 'abc', 'defg');

-- 5. Non existent user name
SELECT documentdb_api_internal.authenticate_with_scram_sha256('abc', 'abc', 'defgh');

-- 6. User name > 64
SELECT documentdb_api_internal.authenticate_with_scram_sha256('abcdefghijklmnopqrstuvwxyzabcdefghijklmnopqrstuvwxyzabcdefghijklmnopqrstuvwxyzabcdefghijklmnopqrstuvwxyz', 'abcdefghijklmnopqrstuvwxyzabcdefghijklmnopqrstuvwxyzabcdefghijklmnopqrstuvwxyzabcdefghijklmnopqrstuvwxyz', 'client proof');

-- 7. Incorrect auth message for a valid user
select documentdb_api_internal.authenticate_with_scram_sha256('user1', 'authMsg1', 'clientProof123');

DROP ROLE user1;
DROP ROLE uSeR2;
DROP ROLE "User3";

/* tests with helper testing functions */
SET client_min_messages TO ERROR;
create role "testuser"    with LOGIN PASSWORD '<password_placeholder111>';
create role "test""user"    with LOGIN PASSWORD '<password_placeholder222>';
create role "test""user."   with LOGIN PASSWORD '<password_placeholder333>';
create role "test user"   with LOGIN PASSWORD '<password_placeholder444>';
create role "test\user"    with LOGIN PASSWORD '<password_placeholder555>';
create role TestUserCase    with LOGIN PASSWORD '<password_placeholder666>';
RESET client_min_messages;

SET client_min_messages = LOG;
-- CALL THE TEST FUNCTION with a valid user name
SELECT test_documentdb_scram_sha256_dual_api('testuser', '<password_placeholder111>');

-- with quotation marks in role name
SELECT test_documentdb_scram_sha256_dual_api('test"user', '<password_placeholder222>');

-- with quotation marks and dot in role name
SELECT test_documentdb_scram_sha256_dual_api('test"user.', '<password_placeholder333>');

-- with space in role name
SELECT test_documentdb_scram_sha256_dual_api('test user', '<password_placeholder444>');

-- with backslash in role name
SELECT test_documentdb_scram_sha256_dual_api('test\user', '<password_placeholder555>');

-- Failed test for incorrect password
SELECT test_documentdb_scram_sha256_dual_api('test\user', '<password_placeholder111>');

-- Test with isNativeAuthEnabled set to OFF
SET documentdb.isNativeAuthEnabled TO OFF;
SELECT test_documentdb_scram_sha256_dual_api('testuser', '<password_placeholder111>');
SET documentdb.isNativeAuthEnabled TO ON;

-- DROP THE USERS CREATED FOR THE TEST
DROP role "testuser";
DROP role "test""user";
DROP role "test""user.";
DROP role "test user";
DROP role "test\user";
DROP role TestUserCase;