:- module(poly1305, [poly1305/3, test_poly1305/0]).
:- use_module(binchunking).
:- use_module(types).

% docs: https://www.rfc-editor.org/rfc/rfc7539.html#section-2.5.2

poly1305(Key, Text, Tag) :-
    assertion(is_poly1305_key(Key)),
    assertion(is_u8s(Text)),
    parse_key(Key, R, Pad),

    poly1305_impl(R, Pad, 0, TagInt, Text),
    as_int_le(16, Tag, TagInt),
    assertion(is_poly1305_tag(Tag)).

parse_key(Key, R, Pad) :-
    length(RBytes, 16),
    length(PadBytes, 16),
    append(RBytes, PadBytes, Key),
    as_int_le(16, RBytes, R0),
    as_int_le(16, PadBytes, Pad),
    R is R0 /\ 0x0FFFFFFC0FFFFFFC0FFFFFFC0FFFFFFF.

prime(P) :- P is 2^130-5.
poly1305_impl(_, Pad, Acc, Result, []) :- 
    Result is (Acc + Pad) /\ (1<<128 - 1).
poly1305_impl(R, Pad, Acc, Result, Text) :-
    prime(Prime),
    append_chunk(16, BytesTrunc, Remaining, Text),
    append(BytesTrunc, [1], Bytes),
    length(Bytes, W),
    as_int_le(W, Bytes, Block),
    Acc2 is ((Acc + Block) * R) mod Prime, !,
    poly1305_impl(R, Pad, Acc2, Result, Remaining).

test_poly1305 :-
    test_parse_key,
    test_main.

test_parse_key :-
    parse_key(
        [
            0x85, 0xd6, 0xbe, 0x78, 0x57, 0x55, 0x6d, 0x33, 
            0x7f, 0x44, 0x52, 0xfe, 0x42, 0xd5, 0x06, 0xa8,
            0x01, 0x03, 0x80, 0x8a, 0xfb, 0x0d, 0xb2, 0xfd,
            0x4a, 0xbf, 0xf6, 0xaf, 0x41, 0x49, 0xf5, 0x1b
        ],
        R, 
        Pad
    ),
    assertion(Pad = 0x1bf54941aff6bf4afdb20dfb8a800301),
    assertion(R = 0x806d5400e52447c036d555408bed685).

test_main :- 
    string_bytes("Cryptographic Forum Research Group", Bytes, utf8),
    poly1305(
        [
            0x85, 0xd6, 0xbe, 0x78, 0x57, 0x55, 0x6d, 0x33, 
            0x7f, 0x44, 0x52, 0xfe, 0x42, 0xd5, 0x06, 0xa8,
            0x01, 0x03, 0x80, 0x8a, 0xfb, 0x0d, 0xb2, 0xfd,
            0x4a, 0xbf, 0xf6, 0xaf, 0x41, 0x49, 0xf5, 0x1b
        ],
        Bytes,
        Tag
    ),
    assertion(Tag = [
        0xa8, 0x06, 0x1d, 0xc1, 0x30, 0x51, 0x36, 0xc6,
        0xc2, 0x2b, 0x8b, 0xaf, 0x0c, 0x01, 0x27, 0xa9
    ]).