Image Steganography — Zig source

//! steganography — least-significant-bit (LSB) encoding on RGBA pixel data.
//!
//! Language: Zig 0.14 (standard library only)
//! Ported from: src/lib/steganography.ts (the canonical TypeScript implementation).
//! display source — part of CosmoDev's polyglot tool pages.
//!
//! Pure logic - no Canvas, no DOM. Per the polyglot task spec this port covers
//! the LSB encode/decode math: capacity, payload-header packing, and the
//! bit-level embed/extract over R/G/B channels. The optional AES-256-GCM
//! body encryption from the TS reference is included too (Zig's std.crypto
//! provides PBKDF2 + GCM natively); only the browser Canvas API parts are
//! out of scope.
//!
//! Wire format (the "payload" hidden in the pixels):
//!   4-byte big-endian header, then the body. The header's top bit is an
//!   encryption flag (1 = body is salt+IV+AES-GCM ciphertext, 0 = body is raw
//!   UTF-8); the low 31 bits are the body length in bytes. The flag makes the
//!   "password required" / "not password-protected" errors deterministic.
//!
//! Payload bits are written MSB-first, one per R/G/B channel in raster order
//! (Alpha is never touched): bit i lands in pixel floor(i/3), channel i%3.
//! Capacity = floor(width * height * 3 / 8) payload bytes.

const std = @import("std");

const pbkdf2_iterations: u32 = 100_000;
const salt_bytes: usize = 16;
const iv_bytes: usize = 12;
const gcm_tag_bytes: usize = 16;

/// Salt + IV + GCM tag overhead added to the body when a password is used.
pub const encryption_overhead_bytes: usize = salt_bytes + iv_bytes + gcm_tag_bytes;
/// The 4-byte length header is also stored in the pixels, so it consumes capacity.
pub const header_bytes: usize = 4;

const Aes256Gcm = std.crypto.aead.aes_gcm.Aes256Gcm;
const HmacSha256 = std.crypto.auth.hmac.sha2.HmacSha256;

/// RGBA pixel buffer with its dimensions (what the browser's ImageData holds).
pub const StegoImageData = struct {
    width: usize,
    height: usize,
    /// RGBA, 4 bytes per pixel; exactly width * height * 4 bytes.
    data: []u8,
};

pub const Error = error{
    InvalidDimensions,
    PasswordEmpty,
    MessageTooLong,
    NoHiddenMessage,
    NotPasswordProtected,
    PasswordRequired,
    DecryptionFailed,
    OutOfMemory,
};

/// Max payload bytes (header + body) an image of this size can carry.
pub fn calculateCapacity(width: usize, height: usize) Error!usize {
    if (width == 0 or height == 0) return Error.InvalidDimensions;
    return (width * height * 3) / 8;
}

/// PBKDF2-SHA256 (100k iterations) -> AES-256-GCM key.
fn deriveKey(password: []const u8, salt: [salt_bytes]u8) [32]u8 {
    var key: [32]u8 = undefined;
    std.crypto.pwhash.pbkdf2(&key, password, &salt, pbkdf2_iterations, HmacSha256) catch
        unreachable; // fixed-size output can only fail on weak parameters
    return key;
}

/// AES-256-GCM encrypt bytes -> packed salt + IV + ciphertext (+ tag).
fn encryptBytes(allocator: std.mem.Allocator, plain: []const u8, password: []const u8) ![]u8 {
    var salt: [salt_bytes]u8 = undefined;
    var iv: [iv_bytes]u8 = undefined;
    std.crypto.random.bytes(&salt);
    std.crypto.random.bytes(&iv);
    const key = deriveKey(password, salt);

    const packed_len = salt_bytes + iv_bytes + plain.len + Aes256Gcm.tag_length;
    const packed = try allocator.alloc(u8, packed_len);
    errdefer allocator.free(packed);
    @memcpy(packed[0..salt_bytes], &salt);
    @memcpy(packed[salt_bytes .. salt_bytes + iv_bytes], &iv);

    var tag: [Aes256Gcm.tag_length]u8 = undefined;
    Aes256Gcm.encrypt(
        packed[salt_bytes + iv_bytes ..][0..plain.len],
        &tag,
        plain,
        "", // no associated data
        iv,
        key,
    ) catch return Error.DecryptionFailed;
    @memcpy(packed[packed_len - Aes256Gcm.tag_length ..], &tag);
    return packed;
}

/// Unpack and AES-256-GCM decrypt a salt + IV + ciphertext payload.
fn decryptBytes(allocator: std.mem.Allocator, packed: []const u8, password: []const u8) Error![]u8 {
    var salt: [salt_bytes]u8 = undefined;
    @memcpy(&salt, packed[0..salt_bytes]);
    var iv: [iv_bytes]u8 = undefined;
    @memcpy(&iv, packed[salt_bytes .. salt_bytes + iv_bytes]);

    const ct = packed[salt_bytes + iv_bytes .. packed.len - Aes256Gcm.tag_length];
    var tag: [Aes256Gcm.tag_length]u8 = undefined;
    @memcpy(&tag, packed[packed.len - Aes256Gcm.tag_length ..]);

    const key = deriveKey(password, salt);
    const plain = allocator.alloc(u8, ct.len) catch return Error.OutOfMemory;
    errdefer allocator.free(plain);
    Aes256Gcm.decrypt(plain, ct, tag, "", iv, key) catch return Error.DecryptionFailed;
    return plain;
}

/// Write `payload` into the LSBs of the R/G/B channels of a COPY of `data`
/// (the input is never mutated). Caller owns the returned buffer.
fn embedBits(allocator: std.mem.Allocator, data: []const u8, payload: []const u8) Error![]u8 {
    const out = allocator.dupe(u8, data) catch return Error.OutOfMemory;
    errdefer allocator.free(out);
    const total_bits = payload.len * 8;
    var i: usize = 0;
    while (i < total_bits) : (i += 1) {
        const byte = payload[i >> 3];
        const bit = (byte >> @intCast(7 - (i & 7))) & 1;
        const px = i / 3;
        const channel = i % 3;
        const idx = px * 4 + channel;
        out[idx] = (out[idx] & 0xfe) | bit;
    }
    return out;
}

/// Read `count` payload bytes back out of the R/G/B LSBs, starting at
/// `offset_bytes` into the logical payload stream. `out` must be `count` long.
fn extractBits(data: []const u8, out: []u8, offset_bytes: usize, count: usize) void {
    const start_bit = offset_bytes * 8;
    var i: usize = 0;
    while (i < count * 8) : (i += 1) {
        const bit_index = start_bit + i;
        const px = bit_index / 3;
        const channel = bit_index % 3;
        const bit = data[px * 4 + channel] & 1;
        out[i >> 3] |= bit << @intCast(7 - (i & 7));
    }
}

/// Hide `message` inside a copy of `image`'s pixels (LSB of R/G/B) and
/// return the modified pixel data. With `password`, the message body is
/// AES-256-GCM encrypted first. Fails if the message (including header and
/// encryption overhead) exceeds the image capacity, or on an empty password.
pub fn hideMessage(
    allocator: std.mem.Allocator,
    image: StegoImageData,
    message: []const u8,
    password: ?[]const u8,
) Error!StegoImageData {
    if (password) |p| {
        if (p.len == 0) return Error.PasswordEmpty;
    }
    const capacity = try calculateCapacity(image.width, image.height);

    const body = if (password) |p|
        try encryptBytes(allocator, message, p)
    else
        message;
    var body_owned = password != null;
    defer if (body_owned) allocator.free(@constCast(body));

    const payload_len = header_bytes + body.len;
    const payload = allocator.alloc(u8, payload_len) catch return Error.OutOfMemory;
    defer allocator.free(payload);

    // 4-byte big-endian header: top bit = encrypted flag, low 31 bits = length.
    const header: u32 = @as(u32, @intCast(body.len)) |
        (if (password != null) @as(u32, 0x8000_0000) else 0);
    std.mem.writeInt(u32, payload[0..4], header, .big);
    @memcpy(payload[header_bytes..], body);
    body_owned = false;

    if (payload.len > capacity) return Error.MessageTooLong;

    return .{
        .width = image.width,
        .height = image.height,
        .data = try embedBits(allocator, image.data, payload),
    };
}

/// Read the hidden message out of `image`'s pixels. Fails when the pixels
/// carry no valid payload, when the payload is encrypted but no password is
/// given, when a password is given but the payload is plaintext, and on a
/// wrong password (GCM authentication failure). Caller owns the result.
pub fn extractMessage(
    allocator: std.mem.Allocator,
    image: StegoImageData,
    password: ?[]const u8,
) Error![]u8 {
    if (password) |p| {
        if (p.len == 0) return Error.PasswordEmpty;
    }
    const capacity = try calculateCapacity(image.width, image.height);

    var header_buf: [header_bytes]u8 = undefined;
    extractBits(image.data, &header_buf, 0, header_bytes);
    const header = std.mem.readInt(u32, &header_buf, .big);
    const encrypted = (header & 0x8000_0000) != 0;
    const length: usize = header & 0x7fff_ffff;
    if (length == 0 and !encrypted) {
        return allocator.dupe(u8, "") catch Error.OutOfMemory;
    }
    if (header_bytes + length > capacity or
        length < (if (encrypted) salt_bytes + iv_bytes + gcm_tag_bytes else 1))
    {
        return Error.NoHiddenMessage;
    }

    const body = allocator.alloc(u8, length) catch return Error.OutOfMemory;
    defer allocator.free(body);
    extractBits(image.data, body, header_bytes, length);

    if (!encrypted) {
        if (password != null) return Error.NotPasswordProtected;
        // Validate UTF-8 the way the TS TextDecoder(fatal: true) does.
        if (!std.unicode.utf8ValidateSlice(body)) return Error.NoHiddenMessage;
        return allocator.dupe(u8, body) catch Error.OutOfMemory;
    }
    const pass = password orelse return Error.PasswordRequired;
    const plain = try decryptBytes(allocator, body, pass);
    if (!std.unicode.utf8ValidateSlice(plain)) return Error.NoHiddenMessage;
    return plain;
}