Skip to content

Instantly share code, notes, and snippets.

@bwbroersma

bwbroersma/totp.pg.sql

Last active October 19, 2022 01:00
Show Gist options
  • Save bwbroersma/676d0de32263ed554584ab132434ebd9 to your computer and use it in GitHub Desktop.
Save bwbroersma/676d0de32263ed554584ab132434ebd9 to your computer and use it in GitHub Desktop.
Download ZIP
Raw
totp.pg.sql
CREATE EXTENSION pgcrypto;
CREATE OR REPLACE FUNCTION totp(key BYTEA, clock_offset INT DEFAULT 0) RETURNS INT AS $$
DECLARE
c BYTEA := '\x000000000' || TO_HEX(FLOOR(EXTRACT(EPOCH FROM NOW()) / 30)::INT + clock_offset);
mac BYTEA := HMAC(c, key, 'sha1');
trunc_offset INT := GET_BYTE(mac, 19) % 16;
result INT := SUBSTRING(SET_BIT(SUBSTRING(mac FROM 1 + trunc_offset FOR 4), 7, 0)::TEXT, 2)::BIT(32)::INT % 1000000;
BEGIN
RETURN result;
END;
$$ LANGUAGE plpgsql;
@steve-chavez
Copy link

This is really interesting @bwbroersma. Thanks for sharing!

If you'd like you could add this to our docs in the form of a how-to guide(more info about that here).

@bwbroersma
Copy link
Author

There was a small 'bug': prefix zero's since it returned an integer.

Here is a full htop/totp implementation with the support for other hash functions.

CREATE EXTENSION pgcrypto;
CREATE OR REPLACE FUNCTION hotp(key BYTEA, c INT, digits INT DEFAULT 6, hash TEXT DEFAULT 'sha1') RETURNS TEXT AS $$
DECLARE
    c BYTEA := '\x' || LPAD(TO_HEX(c), 16, '0');
    mac BYTEA := HMAC(c, key, hash);
    trunc_offset INT := GET_BYTE(mac, length(mac) - 1) % 16;
    result TEXT := SUBSTRING(SET_BIT(SUBSTRING(mac FROM 1 + trunc_offset FOR 4), 7, 0)::TEXT, 2)::BIT(32)::INT % (10 ^ digits)::INT;
BEGIN
    RETURN LPAD(result, digits, '0');
END;
$$ LANGUAGE plpgsql IMMUTABLE LEAKPROOF STRICT PARALLEL SAFE;


CREATE OR REPLACE FUNCTION totp(key BYTEA, clock_offset INT DEFAULT 0, digits INT DEFAULT 6, hash TEXT DEFAULT 'sha1') RETURNS TEXT AS $$
DECLARE
    c INT := FLOOR(EXTRACT(EPOCH FROM NOW()) / 30)::INT + clock_offset;
BEGIN
    RETURN hotp(key, c, digits, hash);
END;
$$ LANGUAGE plpgsql STABLE LEAKPROOF STRICT PARALLEL SAFE;


CREATE OR REPLACE FUNCTION otp_generate_secret(hash TEXT DEFAULT 'sha1') RETURNS BYTEA AS $$
BEGIN
    -- See https://tools.ietf.org/html/rfc4868#section-2.1.2
    -- The optimal key length for HMAC is the block size of the algorithm
    CASE
          WHEN hash = 'sha1'   THEN RETURN gen_random_bytes(20); -- = 160 bits
          WHEN hash = 'sha256' THEN RETURN gen_random_bytes(32); -- = 256 bits
          WHEN hash = 'sha512' THEN RETURN gen_random_bytes(64); -- = 512 bits
          ELSE
            RAISE EXCEPTION 'Unsupported hash algorithm for OTP (see RFC6238/4226).';
            RETURN NULL;
    END CASE;
END;
$$ LANGUAGE plpgsql VOLATILE LEAKPROOF STRICT PARALLEL SAFE;

@pyramation
Copy link

this is pretty cool! since authy and google authenticator use base32, do you by chance know a simple way to do the base32 decode in plpgsql?

@pyramation
Copy link

I modified this to expose a few parameters, and also the base32 encode https://gist.github.com/pyramation/15e2b531ea973fccd0011ce334030adc

CREATE FUNCTION totp.generate ( secret text, period int DEFAULT 30, digits int DEFAULT 6, time_from timestamptz DEFAULT now(), hash text DEFAULT 'sha1', encoding text DEFAULT 'base32', clock_offset int DEFAULT 0 ) RETURNS text AS $EOFCODE$
DECLARE
    c INT := FLOOR(EXTRACT(EPOCH FROM time_from) / period)::INT + clock_offset;
    key bytea;
BEGIN

  IF (encoding = 'base32') THEN 
    key = ( '\x' || totp.base32_to_hex(secret) )::bytea;
  ELSE 
    key = secret::bytea;
  END IF;

  RETURN totp.hotp(key, c, digits, hash);
END;
$EOFCODE$ LANGUAGE plpgsql STABLE;

@bwbroersma as you can see, my base32 is probably not as efficient as it could be, which is why I was wondering if there is a better way to implement base32. Currently the base32 does work if you'd like to use it! But it's more functional string parsing instead of bit shifting which would be better.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment