Supabase Auth Hook for MFA/TOTP/AAL2 account lock

migration.sql

-- create a table to track failed MFA verification attempts
create table if not exists hub_admin.mfa_failed_verification_attempts(
  user_id uuid not null references auth.users(id) on delete cascade
  , factor_id uuid not null
  , failed_at timestamp with time zone not null default now()
  , deleted_at timestamp with time zone
);
create index if not exists idx_mfa_failed_verification_attempts_user_id on hub_admin.mfa_failed_verification_attempts(user_id);

-- a database function that takes a jsonb event and returns a jsonb response
-- if the mfa attempt is valid, it soft-deletes previous failed attempts for this user & returns a decision to continue
-- if it's not valid & there's not too many attempts, it records the attempt and rejects with the count of remaining attempts
-- if it's too many, it locks the user account & returns an error

create or replace function hub_admin.hook_mfa_verification_attempt(event jsonb)
  returns jsonb language plpgsql as $_$
declare
  max_attempts int = 10;
  failed_attempts int;
begin
  if (event -> 'valid')::boolean is true then
    -- soft-delete previous failed attempts for this user
    update
      hub_admin.mfa_failed_verification_attempts
    set
      deleted_at = now()
    where
      user_id = (event ->> 'user_id')::uuid;

    return jsonb_build_object('decision', 'continue');
  end if;

  -- count the number of failed attempts for this user
  select
    count(*) + 1 into failed_attempts -- existing + this one
  from
    hub_admin.mfa_failed_verification_attempts
  where
    user_id = (event ->> 'user_id')::uuid
    and deleted_at is null;

  -- if it's not too many, record the attempt and error with the count of remaining attempts
  if failed_attempts < max_attempts then
    insert into hub_admin.mfa_failed_verification_attempts(user_id, factor_id)
      values((event ->> 'user_id')::uuid, (event ->> 'factor_id')::uuid);

    return jsonb_build_object(
      'error', jsonb_build_object(
        'http_code', 400,
        'message' , 'You have ' || (max_attempts - failed_attempts) || ' attempts remaining.'
      )
    );
  end if;

  -- it's too many; lock the user account & reject
  update
    auth.users
  set
    banned_until = '2100-01-01'
  where
    id = (event ->> 'user_id')::uuid;

  return jsonb_build_object(
    'decision', 'reject',
    'message' , 'The account has been locked.'
  );
end;
$_$;


-- soft delete all failed attempts for this user so they aren't instantly
-- re-banned if they fail their very first attempt after unlock
create or replace function hub_admin.delete_mfa_failed_verification_attempts()
  returns trigger language plpgsql as $_$
begin
  update
    hub_admin.mfa_failed_verification_attempts
  set
    deleted_at = now()
  where
    user_id = new.id
    and deleted_at is null;

  return new;
end;
$_$;

-- a trigger that runs the above function when the user's banned_until column is updated
create or replace trigger delete_mfa_failed_verification_attempts
  after update of banned_until
  on auth.users
  for each row
  when (old.banned_until is distinct from new.banned_until and new.banned_until is null)
  execute procedure hub_admin.delete_mfa_failed_verification_attempts();



-- permissions for the role that runs the hook:
grant all on schema hub_admin to supabase_auth_admin;
grant all on table hub_admin.mfa_failed_verification_attempts to supabase_auth_admin;
grant execute on function hub_admin.hook_mfa_verification_attempt to supabase_auth_admin;

-- if you need to revert for testing:
-- revoke all on schema hub_admin from supabase_auth_admin;
-- revoke all on table hub_admin.mfa_failed_verification_attempts from supabase_auth_admin;
-- revoke execute on function hub_admin.hook_mfa_verification_attempt from supabase_auth_admin;
-- drop function if exists hub_admin.hook_mfa_verification_attempt(jsonb) cascade;
-- drop function if exists hub_admin.delete_mfa_failed_verification_attempts() cascade;
-- drop table if exists hub_admin.mfa_failed_verification_attempts cascade;

migration.test.sql

begin;
select plan (1);

-- table to store test results b/c variables can't be used in pgtap 😭
create temporary table test_results (
  name text primary key
  , result boolean
  , debug text
);

-- function to record test result when run inside a block with vars
create or replace function tests.record_result(name text, result boolean, debug text DEFAULT NULL)
  returns void language plpgsql as $$
begin
  insert into test_results(name, result, debug) values (name, result, debug);
end;
$$;

do $$
  declare
    _user1 uuid = '00000000-0000-0000-0000-000000000001';
    _factor1 uuid = 'FFFFFFFF-FFFF-FFFF-FFFF-FFFFFFFFFFF1';
    _valid_user1_event jsonb = jsonb_build_object('user_id', _user1, 'factor_id', _factor1, 'valid', true);
    _invalid_user1_event jsonb = jsonb_build_object('user_id', _user1, 'factor_id', _factor1, 'valid', false);
    _user2 uuid = '00000000-0000-0000-0000-000000000002';
    _factor2 uuid = 'FFFFFFFF-FFFF-FFFF-FFFF-FFFFFFFFFFF2';
    _valid_user2_event jsonb = jsonb_build_object('user_id', _user2, 'factor_id', _factor2, 'valid', true);
    _invalid_user2_event jsonb = jsonb_build_object('user_id', _user2, 'factor_id', _factor2, 'valid', false);
    _result jsonb;
  begin
    set search_path to 'hub_admin';

    insert into auth.users(id) values (_user1), (_user2);
    -- ensure we don't touch _user2's failed attempts
    perform hook_mfa_verification_attempt(_invalid_user2_event);
    perform hook_mfa_verification_attempt(_invalid_user2_event);

    perform tests.record_result(
      'Invalid event result'
      , (select hook_mfa_verification_attempt(_invalid_user1_event)) -> 'error' ->> 'message' = '9'
    );

    perform tests.record_result(
      'Valid event result'
      , (select hook_mfa_verification_attempt(_valid_user1_event)) ->> 'decision' = 'continue'
    );

    perform tests.record_result(
      'User2''s failed attempts aren''t touched'
      , (select count(*) = 2 from mfa_failed_verification_attempts where user_id = _user2 and deleted_at is null)
    );

    -- 9 more failed attempts, the next will be the last
    for _ in 1..9 loop
      perform hook_mfa_verification_attempt(_invalid_user1_event);
    end loop;

    perform tests.record_result(
      '9 failed attempts'
      , (select count(*) = 9 from mfa_failed_verification_attempts where user_id = _user1 and deleted_at is null)
    );

    -- test account lock - 10th & final attempt
    select hook_mfa_verification_attempt(_invalid_user1_event) into _result;
    perform tests.record_result(
      'Account locked response'
      , _result ->> 'decision' = 'reject' and _result ->> 'message' = 'The account has been locked.'
      -- , _result -> 'error' ->> 'message'
    );

    perform tests.record_result(
      'User account is banned until 2100'
      , (select banned_until >= '2100-01-01' from auth.users where id = _user1)
      -- , (select banned_until is null from auth.users)
    );

    perform tests.record_result(
      'User2''s failed attempts still aren''t touched'
      , (select count(*) > 0 from mfa_failed_verification_attempts where user_id = _user2 and deleted_at is null)
    );

    select hook_mfa_verification_attempt(_invalid_user1_event) into _result;
    perform tests.record_result(
      'Account still locked'
      , _result ->> 'decision' = 'reject' and _result ->> 'message' = 'The account has been locked.'
      -- , _result
    );

    -- unban the user (failed attempt deletions taken care of by trigger)
    update auth.users set banned_until = null where id = _user1;

    -- unsuccessful attempt cause the user can't type
    perform tests.record_result(
      'User account unbanned; but user fail'
      , (select hook_mfa_verification_attempt(_invalid_user1_event)) -> 'error' ->> 'message' = '9'
    );

    -- successful attempt clears out previous failed attempts
    perform tests.record_result(
      'User account unbanned; user succeeds'
      , (select hook_mfa_verification_attempt(_valid_user1_event)) ->> 'decision' = 'continue'
    );

    -- cycle starts again
    perform tests.record_result(
      'User account unbanned; 9 new attempts'
      , (select hook_mfa_verification_attempt(_invalid_user1_event)) -> 'error' ->> 'message' = '9'
    );
  end;
$$;

set search_path to 'extensions';

select set_eq(
  $$ select name, result, debug from test_results $$,
  $$ select * from (values
      ('Invalid event result', true, null)
    , ('Valid event result', true, null)
    , ('User2''s failed attempts aren''t touched', true, null)
    , ('9 failed attempts', true, null)
    , ('Account locked response', true, null)
    , ('User account is banned until 2100', true, null)
    , ('User2''s failed attempts still aren''t touched', true, null)
    , ('Account still locked', true, null)
    , ('User account unbanned; but user fail', true, null)
    , ('User account unbanned; user succeeds', true, null)
    , ('User account unbanned; 9 new attempts', true, null)
  ) as t $$,
  'MFA Auth hook tests:'
);

select * from finish();
rollback;