vtoyexpand.c

/******************************************************************************
 * vtoyexpand.c  ---- ventoy auto install script variable expansion
 *
 * Copyright (c) 2022, longpanda <admin@ventoy.net>
 *
 * This program is free software; you can redistribute it and/or
 * modify it under the terms of the GNU General Public License as
 * published by the Free Software Foundation; either version 3 of the
 * License, or (at your option) any later version.
 * 
 * This program is distributed in the hope that it will be useful, but
 * WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
 * General Public License for more details.
 * 
 * You should have received a copy of the GNU General Public License
 * along with this program; if not, see <http://www.gnu.org/licenses/>.
 *
 */

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <stdarg.h>
#include <errno.h>
#include <unistd.h>
#include <fcntl.h>
#include <sys/types.h>
#include <sys/mman.h>
#include <sys/ioctl.h>
#include <sys/stat.h>
#include <sys/types.h>
#include <linux/fs.h>
#include <dirent.h>
#include "vtoytool.h"

#ifndef O_BINARY
#define O_BINARY 0
#endif

#define TMP_FILE    "/ventoy/tmp_var_expansion"

#define SIZE_1MB    (1024 * 1024)
#define ulong unsigned long
#define ulonglong  unsigned long long

typedef struct disk_info
{
    char name[128];
    ulonglong size;
    int isUSB;
    int isSDX;
}disk_info;

static disk_info *g_disk_list = NULL;
static int g_disk_num = 0;
static const char *g_vtoy_disk_name = NULL;

static void vlog(const char *fmt, ...)
{
    int n = 0;
    va_list arg;
    FILE *fp = NULL;
    char log[1024];

    fp = fopen("/ventoy/autoinstall.log", "a+");
    if (fp)
    {
        va_start(arg, fmt);
        n += vsnprintf(log, sizeof(log) - 1, fmt, arg);
        va_end(arg);

        fwrite(log, 1, n, fp);        
        fclose(fp);
    }
}

static int copy_file(const char *file1, const char *file2)
{
    int n;
    int size;
    int ret = 1;
    FILE *fp1 = NULL;
    FILE *fp2 = NULL;
    char *buf = NULL;

    fp1 = fopen(file1, "rb");
    if (!fp1)
    {
        vlog("Failed to read file <%s>\n", file1);
        goto end;
    }
    
    fp2 = fopen(file2, "wb+");
    if (!fp2)
    {
        vlog("Failed to create file <%s>\n", file2);
        goto end;
    }

    fseek(fp1, 0, SEEK_END);
    size = (int)ftell(fp1);
    fseek(fp1, 0, SEEK_SET);

    buf = malloc(size);
    if (!buf)
    {
        vlog("Failed to malloc buf\n");
        goto end;
    }

    n = fread(buf, 1, size, fp1);
    if (n != size)
    {
        vlog("Failed to read <%s> %d %d\n", file1, n, size);
        goto end;
    }
    
    n = fwrite(buf, 1, size, fp2);
    if (n != size)
    {
        vlog("Failed to write <%s> %d %d\n", file2, n, size);
        goto end;
    }

    ret = 0;

end:

    if (fp1)
        fclose(fp1);
    if (fp2)
        fclose(fp2);
    if (buf)
        free(buf);

    return ret;
}

static int vtoy_is_possible_blkdev(const char *name)
{
    if (name[0] == '.')
    {
        return 0;
    }

    /* /dev/ramX */
    if (name[0] == 'r' && name[1] == 'a' && name[2] == 'm')
    {
        return 0;
    }

    /* /dev/loopX */
    if (name[0] == 'l' && name[1] == 'o' && name[2] == 'o' && name[3] == 'p')
    {
        return 0;
    }

    /* /dev/dm-X */
    if (name[0] == 'd' && name[1] == 'm' && name[2] == '-' && IS_DIGIT(name[3]))
    {
        return 0;
    }

    /* /dev/srX */
    if (name[0] == 's' && name[1] == 'r' && (name[2] >= '0' && name[2] <= '9'))
    {
        return 0;
    }
    
    return 1;
}

static ulonglong vtoy_get_disk_size_in_byte(const char *disk)
{
    int fd;
    int rc;
    unsigned long long size = 0;
    char diskpath[256] = {0};
    char sizebuf[64] = {0};

    // Try 1: get size from sysfs
    snprintf(diskpath, sizeof(diskpath) - 1, "/sys/block/%s/size", disk);
    if (access(diskpath, F_OK) >= 0)
    {
        vlog("get disk size from sysfs for %s\n", disk);
        
        fd = open(diskpath, O_RDONLY | O_BINARY);
        if (fd >= 0)
        {
            read(fd, sizebuf, sizeof(sizebuf));
            size = strtoull(sizebuf, NULL, 10);
            close(fd);
            return (size * 512);
        }
    }
    else
    {
        vlog("%s not exist \n", diskpath);
    }

    // Try 2: get size from ioctl
    snprintf(diskpath, sizeof(diskpath) - 1, "/dev/%s", disk);
    fd = open(diskpath, O_RDONLY);
    if (fd >= 0)
    {
        vlog("get disk size from ioctl for %s\n", disk);
        rc = ioctl(fd, BLKGETSIZE64, &size);
        if (rc == -1)
        {
            size = 0;
            vlog("failed to ioctl %d\n", rc);
        }
        close(fd);
    }
    else
    {
        vlog("failed to open %s %d\n", diskpath, errno);
    }

    vlog("disk %s size %llu bytes\n", disk, (ulonglong)size);
    return size;
}

static int get_disk_num(void)
{
    int n = 0;
    DIR* dir = NULL;
    struct dirent* p = NULL;

    dir = opendir("/sys/block");
    if (!dir)
    {
        return 0;
    }
    
    while ((p = readdir(dir)) != NULL)
    {
        n++;
    }
    closedir(dir);

    return n;
}

static int is_usb_disk(const char *diskname)
{
    int rc;
    char dstpath[1024] = { 0 };
    char syspath[1024] = { 0 };
    
    snprintf(syspath, sizeof(syspath), "/sys/block/%s", diskname);
    rc = readlink(syspath, dstpath, sizeof(dstpath) - 1);
    if (rc > 0 && strstr(dstpath, "/usb"))
    {
        return 1;
    }

    return 0;
}

static int get_all_disk(void)
{
    int i = 0;
    int j = 0;
    int num = 0;
    ulonglong cursize = 0;
    DIR* dir = NULL;
    struct dirent* p = NULL;
    disk_info *node = NULL;
    disk_info tmpnode;

    num = get_disk_num();
    if (num <= 0)
    {
        return 1;
    }

    g_disk_list = malloc(num * sizeof(disk_info));
    if (!g_disk_list)
    {
        return 1;
    }
    memset(g_disk_list, 0, num * sizeof(disk_info));

    dir = opendir("/sys/block");
    if (!dir)
    {
        return 0;
    }

    while (((p = readdir(dir)) != NULL) && g_disk_num < num)
    {
        if (!vtoy_is_possible_blkdev(p->d_name))
        {
            vlog("disk %s is filted by name\n", p->d_name);
            continue;
        }
    
        cursize = vtoy_get_disk_size_in_byte(p->d_name);

        node = g_disk_list + g_disk_num;
        g_disk_num++;

        
        snprintf(node->name, sizeof(node->name), p->d_name);
        node->size = cursize;
        node->isUSB = is_usb_disk(p->d_name);
        if (strncmp(node->name, "sd", 2) == 0)
        {
            node->isSDX = 1;
        }
    }
    closedir(dir);

    /* sort disks */
    for (i = 0; i < g_disk_num; i++)
    {
        for (j = i + 1; j < g_disk_num; j++)
        {
            if (g_disk_list[i].isSDX && g_disk_list[j].isSDX)
            {
                if (strcmp(g_disk_list[i].name, g_disk_list[j].name) > 0)
                {
                    memcpy(&tmpnode, g_disk_list + i, sizeof(tmpnode));
                    memcpy(g_disk_list + i, g_disk_list + j, sizeof(tmpnode));
                    memcpy(g_disk_list + j, &tmpnode, sizeof(tmpnode));
                }
            }
        }
    }

    vlog("============ DISK DUMP BEGIN ===========\n");
    for (i = 0; i < g_disk_num; i++)
    {
        node = g_disk_list + i;
        vlog("[%d] %s %dGB(%llu) USB:%d\n", i, node->name,
            node->size / 1024 / 1024 / 1024, node->size, node->isUSB);
    }
    vlog("============ DISK DUMP END ===========\n");
    
    return 0;
}

static int expand_var(const char *var, char *value, int len)
{
    int i;
    int index = -1;
    ulonglong uiDst = 0;
    ulonglong delta = 0;
    ulonglong maxsize = 0;
    ulonglong maxdelta = 0xFFFFFFFFFFFFFFFFULL;
    disk_info *node = NULL;
    value[0] = 0;

    if (strcmp(var, "VT_LINUX_DISK_SDX_1ST_NONVTOY") == 0)
    {
        for (i = 0; i < g_disk_num; i++)
        {
            node = g_disk_list + i;
            if (node->size > 0 && node->isSDX && strcmp(node->name, g_vtoy_disk_name) != 0)
            {
                vlog("%s=<%s>\n", var, node->name);
                snprintf(value, len, "%s", node->name);
                return 0;
            }
        }

        vlog("[Error] %s not found\n", var);
    }
    else if (strcmp(var, "VT_LINUX_DISK_SDX_1ST_NONUSB") == 0)
    {
        for (i = 0; i < g_disk_num; i++)
        {
            node = g_disk_list + i;
            if (node->size > 0 && node->isSDX && node->isUSB == 0)
            {
                vlog("%s=<%s>\n", var, node->name);
                snprintf(value, len, "%s", node->name);
                return 0;
            }
        }

        vlog("[Error] %s not found\n", var);
    }
    else if (strcmp(var, "VT_LINUX_DISK_MAX_SIZE") == 0)
    {
        for (i = 0; i < g_disk_num; i++)
        {
            node = g_disk_list + i;
            if (node->size > 0 && node->size > maxsize)
            {
                index = i;
                maxsize = node->size;
            }
        }

        if (index >= 0)
        {
            vlog("%s=<%s>\n", var, g_disk_list[index].name);
            snprintf(value, len, "%s", g_disk_list[index].name);
            return 0;
        }
        else
        {
            vlog("[Error] %s not found\n", var);
        }
    }
    else if (strncmp(var, "VT_LINUX_DISK_CLOSEST_", 22) == 0)
    {
        uiDst = strtoul(var + 22, NULL, 10);
        uiDst = uiDst * (1024ULL * 1024ULL * 1024ULL);
    
        for (i = 0; i < g_disk_num; i++)
        {
            node = g_disk_list + i;
            if (node->size == 0)
            {
                continue;
            }
        
            if (node->size > uiDst)
            {
                delta = node->size - uiDst;
            }
            else
            {
                delta = uiDst - node->size;
            }
            
            if (delta < maxdelta)
            {
                index = i;
                maxdelta = delta;
            }
        }

        if (index >= 0)
        {
            vlog("%s=<%s>\n", var, g_disk_list[index].name);
            snprintf(value, len, "%s", g_disk_list[index].name);
            return 0;
        }
        else
        {
            vlog("[Error] %s not found\n", var);
        }
    }
    else
    {
        vlog("Invalid var name <%s>\n", var);
        snprintf(value, len, "$$%s$$", var);
    }

    if (value[0] == 0)
    {
        snprintf(value, len, "$$%s$$", var);
    }

    return 0;
}

int vtoyexpand_main(int argc, char **argv)
{
    FILE *fp = NULL;
    FILE *fout = NULL;
    char *start = NULL;
    char *end = NULL;
    char line[4096];
    char value[256];

    vlog("========= vtoyexpand_main %d ========\n", argc);

    if (argc != 3)
    {
        return 1;
    }

    g_vtoy_disk_name = argv[2];
    if (strncmp(g_vtoy_disk_name, "/dev/", 5) == 0)
    {
        g_vtoy_disk_name += 5;
    }
    vlog("<%s> <%s> <%s>\n", argv[1], argv[2], g_vtoy_disk_name);    

    get_all_disk();
    
    fp = fopen(argv[1], "r");
    if (!fp)
    {
        vlog("Failed to open file <%s>\n", argv[1]);
        return 1;
    }

    fout = fopen(TMP_FILE, "w+");
    if (!fout)
    {
        vlog("Failed to create file <%s>\n", TMP_FILE);
        fclose(fp);
        return 1;
    }

    memset(line, 0, sizeof(line));
    memset(value, 0, sizeof(value));
    
    while (fgets(line, sizeof(line), fp))
    {
        start = strstr(line, "$$VT_");
        if (start)
        {
            end = strstr(start + 5, "$$");
        }

        if (start && end)
        {
            *start = 0;
            fprintf(fout, "%s", line);

            *end = 0;
            expand_var(start + 2, value, sizeof(value));
            fprintf(fout, "%s", value);
            
            fprintf(fout, "%s", end + 2);

            memset(value, 0, sizeof(value));
        }
        else
        {
            fprintf(fout, "%s", line);
        }

        line[0] = line[4095] = 0;
    }
    
    fclose(fp);
    fclose(fout);

    vlog("delete file <%s>\n", argv[1]);
    remove(argv[1]);

    vlog("Copy file <%s> --> <%s>\n", TMP_FILE, argv[1]);
    copy_file(TMP_FILE, argv[1]);
    
    return 0;
}

// wrapper main
#ifndef BUILD_VTOY_TOOL
int main(int argc, char **argv)
{
    return vtoyexpand_main(argc, argv);
}
#endif