Skip to content

Commit

Permalink
feat: implement key expiry notifications
Browse files Browse the repository at this point in the history
  • Loading branch information
david-r-cox committed Dec 4, 2024
1 parent 12c5c75 commit df18b50
Show file tree
Hide file tree
Showing 2 changed files with 309 additions and 0 deletions.
189 changes: 189 additions & 0 deletions extension/keyhippo--1.0.1.sql
Original file line number Diff line number Diff line change
Expand Up @@ -43,6 +43,8 @@ CREATE EXTENSION IF NOT EXISTS pgcrypto;

CREATE EXTENSION IF NOT EXISTS pg_net;

CREATE EXTENSION IF NOT EXISTS pg_cron;

-- Create custom types
CREATE TYPE keyhippo.app_permission AS ENUM (
'manage_groups',
Expand Down Expand Up @@ -307,6 +309,186 @@ BEGIN
END;
$$;

CREATE OR REPLACE FUNCTION keyhippo.notify_expiring_key ()
RETURNS TRIGGER
LANGUAGE plpgsql
SECURITY DEFINER
AS $$
DECLARE
v_endpoint text;
v_installation_uuid uuid;
v_payload jsonb;
v_enable_notifications boolean;
v_notification_hours int;
v_request_id bigint;
v_response_status int;
v_response_body text;
BEGIN
RAISE LOG 'notify_expiring_key triggered for key: %', NEW.id;
-- Check if notifications are enabled
SELECT
value::boolean INTO v_enable_notifications
FROM
keyhippo_internal.config
WHERE
key = 'enable_key_expiry_notifications';
IF NOT v_enable_notifications THEN
RAISE LOG 'Key expiry notifications are disabled';
RETURN NEW;
END IF;
-- Get the notification hours from config
SELECT
value::int INTO v_notification_hours
FROM
keyhippo_internal.config
WHERE
key = 'key_expiry_notification_hours';
RAISE LOG 'Notification hours: %, Key expires at: %, Current time: %', v_notification_hours, NEW.expires_at, NOW();
-- If the key is about to expire within the specified hours
IF NEW.expires_at <= NOW() + (v_notification_hours || ' hours')::interval AND (OLD.expires_at IS NULL OR OLD.expires_at > NOW() + (v_notification_hours || ' hours')::interval) THEN
RAISE LOG 'Key % is expiring soon, preparing notification', NEW.id;
-- Get the endpoint and installation UUID from the config table
SELECT
value INTO v_endpoint
FROM
keyhippo_internal.config
WHERE
key = 'audit_log_endpoint';
SELECT
value::uuid INTO v_installation_uuid
FROM
keyhippo_internal.config
WHERE
key = 'installation_uuid';
-- Prepare the payload
v_payload := jsonb_build_object('event', 'expiring_key', 'installation_uuid', v_installation_uuid, 'timestamp', now(), 'expiring_key', jsonb_build_object('id', NEW.id, 'user_id', NEW.user_id, 'description', NEW.description, 'expires_at', NEW.expires_at));
RAISE LOG 'Sending notification to endpoint: % with payload: %', v_endpoint, v_payload;
-- Send the notification with error handling
BEGIN
SELECT
(result).status,
(result).content::text,
(result).id INTO v_response_status,
v_response_body,
v_request_id
FROM
net.http_post (url := v_endpoint, body := v_payload, headers := jsonb_build_object('Content-Type', 'application/json')) AS result;
RAISE LOG 'Notification sent. Request ID: %, Status: %, Response: %', v_request_id, v_response_status, v_response_body;
IF v_response_status < 200 OR v_response_status >= 300 THEN
RAISE EXCEPTION 'HTTP request failed with status %', v_response_status;
END IF;
EXCEPTION
WHEN OTHERS THEN
RAISE LOG 'Error sending notification: %', SQLERRM;
END;
-- Insert into audit_log
INSERT INTO keyhippo.audit_log (action, table_name, data)
VALUES ('expiring_key', TG_TABLE_NAME, v_payload);
RAISE LOG 'Audit log entry created for expiring key %', NEW.id;
ELSE
RAISE LOG 'Key % is not expiring soon, no notification sent', NEW.id;
END IF;
RETURN NEW;
END;

$$;

CREATE OR REPLACE FUNCTION keyhippo.notify_expiring_key ()
RETURNS TRIGGER
LANGUAGE plpgsql
SECURITY DEFINER
AS $$
DECLARE
v_endpoint text;
v_installation_uuid uuid;
v_payload jsonb;
v_enable_notifications boolean;
v_notification_hours int;
v_request_id bigint;
BEGIN
RAISE LOG 'notify_expiring_key triggered for key: %', NEW.id;
-- Check if notifications are enabled
SELECT
value::boolean INTO v_enable_notifications
FROM
keyhippo_internal.config
WHERE
key = 'enable_key_expiry_notifications';
IF NOT v_enable_notifications THEN
RAISE LOG 'Key expiry notifications are disabled';
RETURN NEW;
END IF;
-- Get the notification hours from config
SELECT
value::int INTO v_notification_hours
FROM
keyhippo_internal.config
WHERE
key = 'key_expiry_notification_hours';
RAISE LOG 'Notification hours: %, Key expires at: %, Current time: %', v_notification_hours, NEW.expires_at, NOW();
-- If the key is about to expire within the specified hours
IF NEW.expires_at <= NOW() + (v_notification_hours || ' hours')::interval AND (OLD.expires_at IS NULL OR OLD.expires_at > NOW() + (v_notification_hours || ' hours')::interval) THEN
RAISE LOG 'Key % is expiring soon, preparing notification', NEW.id;
-- Get the endpoint and installation UUID from the config table
SELECT
value INTO v_endpoint
FROM
keyhippo_internal.config
WHERE
key = 'audit_log_endpoint';
SELECT
value::uuid INTO v_installation_uuid
FROM
keyhippo_internal.config
WHERE
key = 'installation_uuid';
-- Prepare the payload
v_payload := jsonb_build_object('event', 'expiring_key', 'installation_uuid', v_installation_uuid, 'timestamp', now(), 'expiring_key', jsonb_build_object('id', NEW.id, 'user_id', NEW.user_id, 'description', NEW.description, 'expires_at', NEW.expires_at));
RAISE LOG 'Sending notification to endpoint: % with payload: %', v_endpoint, v_payload;
-- Send the notification
SELECT
net.http_post (url := v_endpoint, body := v_payload, headers := jsonb_build_object('Content-Type', 'application/json')) INTO v_request_id;
RAISE LOG 'Notification sent. Request ID: %', v_request_id;
-- Insert into audit_log
INSERT INTO keyhippo.audit_log (action, table_name, data)
VALUES ('expiring_key', TG_TABLE_NAME, v_payload);
RAISE LOG 'Audit log entry created for expiring key %', NEW.id;
ELSE
RAISE LOG 'Key % is not expiring soon, no notification sent', NEW.id;
END IF;
RETURN NEW;
END;
$$;

CREATE TRIGGER keyhippo_notify_expiring_key_trigger
BEFORE UPDATE OF expires_at ON keyhippo.api_key_metadata
FOR EACH ROW
EXECUTE FUNCTION keyhippo.notify_expiring_key ();

CREATE OR REPLACE FUNCTION keyhippo.update_expiring_keys ()
RETURNS VOID
LANGUAGE plpgsql
SECURITY DEFINER
AS $$
BEGIN
-- Update expires_at for keys that are about to expire
UPDATE
keyhippo.api_key_metadata
SET
expires_at = expires_at
WHERE
expires_at <= NOW() + INTERVAL '1 day'
AND is_revoked = FALSE;
END;
$$;

SELECT
cron.schedule ('0 * * * *', $$
SELECT
keyhippo.update_expiring_keys ();

$$);

CREATE OR REPLACE FUNCTION keyhippo.is_authorized (target_resource regclass, required_permission text)
RETURNS boolean
AS $$
Expand Down Expand Up @@ -1071,6 +1253,13 @@ DECLARE
user_group_id uuid;
user_role_id uuid;
BEGIN
-- Add configurations for key expiry notifications
INSERT INTO keyhippo_internal.config (key, value, description)
VALUES ('key_expiry_notification_hours', '72', 'Number of hours before key expiry to send notification'),
('enable_key_expiry_notifications', 'true', 'Flag to enable/disable key expiry notifications')
ON CONFLICT (key)
DO UPDATE SET
value = EXCLUDED.value, description = EXCLUDED.description;
-- Upsert the default value for the audit log endpoint
INSERT INTO keyhippo_internal.config (key, value, description)
VALUES ('audit_log_endpoint', 'https://app.keyhippo.com/api/ingest', 'Endpoint for sending audit log notifications')
Expand Down
120 changes: 120 additions & 0 deletions tests/tests.sql
Original file line number Diff line number Diff line change
Expand Up @@ -285,5 +285,125 @@ BEGIN
'assign_permission_to_role should assign the permission to the role';
END
$$;
-- Test key expiry notification
SET ROLE postgres;
DO $$
DECLARE
v_api_key_id uuid;
v_api_key text;
v_user_id uuid;
v_notification_sent boolean := FALSE;
v_audit_log_entry jsonb;
BEGIN
-- Set expiry notification time to 2 hours for testing purposes
UPDATE
keyhippo_internal.config
SET
value = '2'
WHERE
key = 'key_expiry_notification_hours';
-- Ensure notifications are enabled
UPDATE
keyhippo_internal.config
SET
value = 'true'
WHERE
key = 'enable_key_expiry_notifications';
-- Create a test user
INSERT INTO auth.users (id, email)
VALUES (gen_random_uuid (), 'testuser@example.com')
RETURNING
id INTO v_user_id;
-- Login as the test user
PERFORM
set_config('request.jwt.claim.sub', v_user_id::text, TRUE);
PERFORM
set_config('request.jwt.claims', json_build_object('sub', v_user_id, 'role', 'authenticated')::text, TRUE);
-- Create a test API key using the real create_api_key function
SELECT
api_key,
api_key_id INTO v_api_key,
v_api_key_id
FROM
keyhippo.create_api_key ('Test Expiring Key');
RAISE NOTICE 'Created API key with ID: %', v_api_key_id;
-- Logout
PERFORM
set_config('request.jwt.claim.sub', '', TRUE);
PERFORM
set_config('request.jwt.claims', '', TRUE);
-- Update the expiry to trigger the notification
UPDATE
keyhippo.api_key_metadata
SET
expires_at = NOW() + INTERVAL '1 hour'
WHERE
id = v_api_key_id;
RAISE NOTICE 'Updated API key expiry';
-- Check if a notification was logged
SELECT
EXISTS (
SELECT
1
FROM
keyhippo.audit_log
WHERE
action = 'expiring_key'
AND (data ->> 'expiring_key')::jsonb ->> 'id' = v_api_key_id::text
AND timestamp > NOW() - INTERVAL '1 minute') INTO v_notification_sent;
IF v_notification_sent THEN
RAISE NOTICE 'Notification sent for key %', v_api_key_id;
-- Fetch and display the audit log entry
SELECT
data INTO v_audit_log_entry
FROM
keyhippo.audit_log
WHERE
action = 'expiring_key'
AND (data ->> 'expiring_key')::jsonb ->> 'id' = v_api_key_id::text
ORDER BY
timestamp DESC
LIMIT 1;
RAISE NOTICE 'Audit log entry: %', v_audit_log_entry;
ELSE
RAISE NOTICE 'No notification found for key %', v_api_key_id;
-- Display recent audit log entries for debugging
RAISE NOTICE 'Recent audit log entries:';
FOR v_audit_log_entry IN (
SELECT
data
FROM
keyhippo.audit_log
WHERE
timestamp > NOW() - INTERVAL '5 minutes'
ORDER BY
timestamp DESC
LIMIT 5)
LOOP
RAISE NOTICE '%', v_audit_log_entry;
END LOOP;
END IF;
-- Cleanup
DELETE FROM keyhippo.api_key_metadata
WHERE id = v_api_key_id;
DELETE FROM auth.users
WHERE id = v_user_id;
-- Assert the result
ASSERT v_notification_sent,
'An expiry notification should have been sent';
EXCEPTION
WHEN OTHERS THEN
-- Cleanup in case of error
IF v_api_key_id IS NOT NULL THEN
DELETE FROM keyhippo.api_key_metadata
WHERE id = v_api_key_id;
END IF;
IF v_user_id IS NOT NULL THEN
DELETE FROM auth.users
WHERE id = v_user_id;
END IF;
RAISE;
END
$$;
-- Clean up
ROLLBACK;

0 comments on commit df18b50

Please sign in to comment.