Multi-precision integer library for Lua 5.2/5.3+, with functions for RSA encryption

mpi.lua

---@class MultiPrecisionInteger
---@field sign boolean|nil
local mpi = {}
mpi.__mt = {__index = mpi, __name = "MultiPrecisionInteger"}

local floor = math.floor

mpi.zero = setmetatable({0}, mpi.__mt)
mpi.one = setmetatable({1}, mpi.__mt)

--- Creates a new multi-precision number.
---@param bytes string|number The bytes representing the number in big-endian form
---@return MultiPrecisionInteger num The new number
function mpi:new(bytes)
    if type(bytes) == "number" then
        bytes = floor(bytes)
        if bytes == 0 then return self.zero end
        local num = setmetatable({sign = bytes < 0}, self.__mt) ---@type MultiPrecisionInteger
        if bytes < 0 then bytes = -bytes end
        while bytes > 0 do
            num[#num+1] = bytes % 0x10000
            bytes = floor(bytes / 0x10000)
        end
        return num
    end
    if #bytes % 2 == 1 then bytes = "\0" .. bytes end
    local n = #bytes / 2
    local num = setmetatable({}, self.__mt) ---@type MultiPrecisionInteger
    local parts = {(">I2"):rep(#bytes / 2):unpack(bytes)}
    table.remove(parts)
    for i, p in ipairs(parts) do num[n - i + 1] = p end
    return num
end

--- Adds two numbers.
---@param b MultiPrecisionInteger
---@return MultiPrecisionInteger sum
function mpi:add(b)
    if self.sign and not b.sign then return b:sub(self:unm())
    elseif not self.sign and b.sign then return self:sub(b:unm()) end
    local n = math.max(#self, #b)
    local sum = setmetatable({sign = self.sign}, self.__mt) ---@type MultiPrecisionInteger
    local c = 0
    for i = 1, n do
        local s = (self[i] or 0) + (b[i] or 0) + c
        if s > 0xFFFF then sum[i], c = s - 0x10000, 1
        else sum[i], c = s, 0 end
    end
    if c ~= 0 then
        sum[n+1] = 1
    end
    return sum
end
mpi.__mt.__add = mpi.add

--- Subtracts two numbers.
---@param b MultiPrecisionInteger
---@return MultiPrecisionInteger diff
function mpi:sub(b)
    if b.sign then return self:add(b:unm()) end
    if self.sign then local n = b:add(self:unm()) n.sign = true return n end
    if self:lt(b) then local n = b:sub(self) n.sign = not n.sign return n end
    local n = math.max(#self, #b)
    local diff = setmetatable({}, self.__mt) ---@type MultiPrecisionInteger
    local c = 0
    for i = 1, n do
        local s = (self[i] or 0) - (b[i] or 0) - c
        if s < 0 then diff[i], c = s + 0x10000, 1
        else diff[i], c = s, 0 end
    end
    if c ~= 0 then
        error("Signed number not caught properly", 2)
    end
    return diff
end
mpi.__mt.__sub = mpi.sub

--- Multiplies two numbers.
---@param b MultiPrecisionInteger
---@return MultiPrecisionInteger product
function mpi:mul(b)
    local product = setmetatable({sign = (not self.sign) ~= (not b.sign)}, self.__mt)
    for i = 1, #self + #b do product[i] = 0 end
    for i = 1, #self do
        local c = 0
        for j = 1, #b do
            local p = self[i] * b[j] + c
            c, p = floor(p / 0x10000), p % 0x10000
            local s = product[i+j-1] + p
            if s > 0xFFFF then
                local k = i+j
                while product[k] == 0xFFFF do
                    product[k] = 0
                    k = k + 1
                end
                product[k] = product[k] + 1
                s = s - 0x10000
            end
            product[i+j-1] = s
        end
        product[i+#b] = product[i+#b] + c
    end
    return product
end
mpi.__mt.__mul = mpi.mul

--- Divides two numbers with remainder.
---@param b MultiPrecisionInteger
---@return MultiPrecisionInteger q
---@return MultiPrecisionInteger r
function mpi:rdiv(b)
    -- Algorithm from L. Huang, Y. Luo, H. Zhong, H. Shen: An efficient multiple-precision division algorithm;
    -- improved by Debapriyay Mukhopadhyay, Subhas C. Nandy: Efficient multiple-precision integer division algorithm
    local m, n = #self, #b
    while self[m] == 0 do m = m - 1 end
    while b[n] == 0 do n = n - 1 end
    if n == 1 then
        -- fall back to long division for small divisor
        local q = setmetatable({sign = (not self.sign) ~= (not b.sign)}, self.__mt)
        local c = 0
        b = b[1]
        for i = m, 1, -1 do
            q[i] = floor((c + self[i]) / b)
            c = ((c + self[i]) % b) * 0x10000
        end
        return q, setmetatable({c / 0x10000}, self.__mt)
    end
    local W = {[0] = 0}
    for i = 1, m do W[i] = self[m-i+1] end
    local D = b[n] * 0x10000 + b[n-1]
    for i = 0, m - n do
        local N = W[i] * 0x100000000 + W[i+1] * 0x10000 + W[i+2]
        local Q = floor(N / D)
        for j = 1, n do
            W[i+j] = W[i+j] - Q * b[n-j+1]
        end
        W[i+1] = W[i] * 0x10000 + W[i+1]
        W[i] = Q
    end
    local adjust, X, savec = false, {}, nil
    for i = 0, m do X[i] = W[i] end
    for i = m, 1, -1 do
        local c = 0
        if X[i] < 0 then c = floor((-X[i] - 1) / 0x10000) + 1
        elseif X[i] >= 0x10000 then c = -floor(X[i] / 0x10000) end
        X[i] = X[i] + c * 0x10000
        X[i-1] = X[i-1] - c
        if c ~= 0 and i == m - n + 1 then
            savec = c
            adjust = true
        end
    end
    if adjust then
        for i = m, m - n + 1, -1 do
            W[i] = W[i] + savec * b[n-i+1]
        end
        for i = m - n, 0, -1 do
            W[i] = X[i]
        end
        for i = m, 1, -1 do
            local c = 0
            if W[i] < 0 then c = floor((-W[i] - 1) / 0x10000) + 1
            elseif W[i] >= 0x10000 then c = -floor(W[i] / 0x10000) end
            W[i] = W[i] + c * 0x10000
            W[i-1] = W[i-1] - c
        end
    else
        for i = 0, m do
            W[i] = X[i]
        end
    end
    local q, r = setmetatable({sign = (not self.sign) ~= (not b.sign)}, self.__mt), setmetatable({}, self.__mt)
    local l = m - n + 1
    for i = 0, l - 1 do q[l-i] = W[i] end
    for i = l, m do r[m-i+1] = W[i] end
    return q, r
end

--- Divides two numbers.
---@param b MultiPrecisionInteger
---@return MultiPrecisionInteger quotient
function mpi:div(b)
    local q = self:rdiv(b)
    return q
end
mpi.__mt.__div = mpi.div

--- Returns the modulus of two numbers.
---@param b MultiPrecisionInteger
---@return MultiPrecisionInteger remainder
function mpi:mod(b)
    local _, r = self:rdiv(b)
    return r
end
mpi.__mt.__mod = mpi.mod

--- Returns the negative version of this number.
---@return MultiPrecisionInteger n
function mpi:unm()
    local n = setmetatable({sign = not self.sign}, self.__mt)
    for i = 1, #self do n[i] = self[i] end
    return n
end
mpi.__mt.__unm = mpi.unm

--- Returns whether self == b.
---@param b MultiPrecisionInteger
---@return boolean eq
function mpi:eq(b)
    if (not not self.sign) ~= (not not b.sign) then return false end
    for i = 1, math.max(#self, #b) do if (self[i] or 0) ~= (b[i] or 0) then return false end end
    return true
end
mpi.__mt.__eq = mpi.eq

--- Returns whether self < b.
---@param b MultiPrecisionInteger
---@return boolean lt
function mpi:lt(b)
    if self.sign and not b.sign then return true
    elseif not self.sign and b.sign then return false end
    for i = math.max(#self, #b), 1, -1 do
        if (self[i] or 0) < (b[i] or 0) then return not self.sign
        elseif (self[i] or 0) > (b[i] or 0) then return not not self.sign end
    end
    return false
end
mpi.__mt.__lt = mpi.lt

--- Returns whether self <= b.
---@param b MultiPrecisionInteger
---@return boolean le
function mpi:le(b)
    if self.sign and not b.sign then return true
    elseif not self.sign and b.sign then return false end
    for i = math.max(#self, #b), 1, -1 do
        if (self[i] or 0) < (b[i] or 0) then return not self.sign
        elseif (self[i] or 0) > (b[i] or 0) then return not not self.sign end
    end
    return true
end
mpi.__mt.__le = mpi.le

--- Trims leading zeros from the number.
---@return MultiPrecisionInteger n
function mpi:trim()
    local n = setmetatable({sign = self.sign}, self.__mt)
    for i = 1, #self do n[i] = self[i] end
    for i = #n, 1, -1 do if n[i] == 0 then n[i] = nil else break end end
    return n
end

--- Performs Montgomery reduction on a number, which is assumed to be in Montgomery form.
---@param r number The number of words in the original factors
---@param n MultiPrecisionInteger The original modulus to use
---@param np number The inverse of the modulus, such that (n * np) % 2^(#self*16) = 2^(#self*16) - 1, found with extended Euclidean algorithm
---@return MultiPrecisionInteger s The result of the reduction
function mpi:redc(r, n, np)
    local p = #n
    local t = {}
    for i = 1, #self do t[i-1] = self[i] end
    for i = #self, r + p do t[i] = 0 end
    for i = 0, r - 1 do
        local c = 0
        local m = (t[i] * np) % 0x10000
        for j = 0, p - 1 do
            local x = t[i+j] + m * n[j+1] + c
            t[i+j], c = x % 0x10000, floor(x / 0x10000)
        end
        for j = p, r + p - i do
            local x = t[i+j] + c
            t[i+j], c = x % 0x10000, floor(x / 0x10000)
        end
    end
    local s = setmetatable({}, self.__mt)
    for i = 0, p do s[i+1] = t[i+r] end
    if s >= n then return s - n
    else return s end
end

--- Performs the extended Euclidean algorithm on this number and (2^16)^(#self), returning the Bezout factors.
---@return MultiPrecisionInteger rp
---@return number np
function mpi:exteuclid()
    local b = setmetatable({}, self.__mt)
    for i = 1, #self do b[i] = 0 end
    b[#b+1] = 1
    local ra, rb, sa, sb, ta, tb = b, self, mpi.one, mpi.zero, mpi.zero, mpi.one
    while rb ~= mpi.zero do
        local q, r = ra:rdiv(rb)
        rb, sb, tb, ra, sa, ta = r, sa - (q * sb):trim(), ta - (q * tb):trim(), rb, sb, tb
    end
    assert(ra == mpi.one)
    return sa, ta[1]
end

--- Performs modular exponentiation.
---@param e MultiPrecisionInteger The exponent to raise to
---@param n MultiPrecisionInteger The modulus to use
---@return MultiPrecisionInteger res The result of the modular exponentiation
function mpi:modexp(e, n)
    local r2m = mpi.one
    for p = 1, #self * 32 do
        r2m = r2m + r2m
        if r2m >= n then r2m = r2m - n end
    end
    local _, np = n:exteuclid()
    local m = (self * r2m):redc(#self, n, np)
    local res = (mpi.one * r2m):redc(#self, n, np)
    for i = 1, #e do
        for j = 0, 15 do
            if bit32.btest(e[i], 2^j) then
                res = (res * m):redc(#self, n, np)
            end
            m = (m * m):redc(#self, n, np)
        end
    end
    return res:redc(#self, n, np)
end

--- Returns an estimation of the numerical representation of this number.
---@return number num
function mpi:tonumber()
    local num = 0
    for i = 1, #self do num = num + 2^(16*(i-1)) * self[i] end
    return num
end

function mpi.__mt:__tostring()
    return tostring(self:tonumber())
end

return mpi