memeater.c

/*
 * memeater utility (SivanagBalla@gmail.com)
 */

#include <stdio.h>
#include <stdlib.h> 
#include <string.h> 
#include <errno.h>
#include <unistd.h>
#include <sys/types.h>
#include <sys/stat.h>
#include <fcntl.h>
#include <sys/mman.h>
#include <inttypes.h>

#include "util.h"

typedef struct {
    uint64_t alloc_total;
    uint64_t free_total;
    uint64_t alloc_fail;
    uint64_t alloc_pass;
    uint64_t sbrk_t;
    uint64_t mmap_t;
    uint64_t file_open;
    uint64_t file_close;
    uint64_t shm_fail;
    uint64_t shm_pass;
    uint64_t shm_total;
} _stats;

typedef struct  {
    uint64_t ptr;
    uint64_t size;
}malloc_info;

typedef struct  {
    char name[16];
    int fd;
    uint64_t ptr;
    uint64_t size;
}shm_info;
#define MAX_SHM 100

int memeater_main(int argc, char* argv[]) {
    int quit = FALSE;
    char buf[1024];

    int sbrk_c = 0;
    int mmap_c = 0;
    pid_t pid = 0;
    int fd = -1;
    _stats st = {0};

    malloc_info mali[50000] = {0};
    uint32_t mind = 0;

    shm_info shi[MAX_SHM] = {0};
    uint32_t shind = 0;

    uint64_t len;

    void* brk_pos;
    void* cur_brk_pos;

    printf("Memory Eater\n");
    brk_pos = sbrk(0);
    pid = getpid();
    printf("sbrk: %p pid: %d\n", brk_pos, pid);

    while (quit == FALSE) {
        buf[0] = '\0';
        printf(">> ");
        fflush(stdout);
        len = getInput(buf, 1024);
        tokenize(buf, len);

        if (tokCnt == 0) {
            continue;
        } else if (!strcmp("h", tokens[0]) || 
                   !strcmp("?", tokens[0]) ||
                   !strcmp("help", tokens[0])
                   ) {
            printf(" help/?/h - help\n");
            printf(" eat      - Eat memory\n");
            printf(" vomit    - Vomit memory\n");
            printf(" eatshm   - Eat shared memory\n");
            printf(" system   - Execute shell command\n");
            printf(" stat     - Stats\n");
            printf(" open     - open a file\n");
            printf(" close    - close a file\n");

        } else if (!strcmp("q", tokens[0])) {
            quit = TRUE;
        } else if (!strcmp("eat", tokens[0])) {
            uint64_t size, no_acc, cnt, echo;
            void *ptr;
            int i;
            if (tokCnt != 5) {
                printf("Usage: eat <size> <cnt> <echo> <no_access>\n");
                continue;
            }
            if(parseNumericArg(tokens[1], &size) == FALSE ||
                    parseNumericArg(tokens[2], &cnt) == FALSE ||
                    parseNumericArg(tokens[3], &echo) == FALSE ||
                    parseNumericArg(tokens[4], &no_acc) == FALSE) {
                printf("Parsing error: Usage: eat <size> <cnt> <echo> <no_access>\n");
                continue;
            }
            printf("Eating %"PRId64" byte chunks %"PRId64" times%s\n", size, cnt,
                    no_acc ? " without accessing" : "");
            sbrk_c = 0;
            mmap_c = 0;
            for(i=0; i < cnt; i++) {
                ptr = malloc(size);

                if (!ptr) {
                    printf("Alloc failed\n");
                    st.alloc_fail++;
                    break;
                }

                st.alloc_total += size;
                st.alloc_pass ++;

                if (mind <= 50000) {
                    mali[mind].ptr = (uint64_t)ptr;
                    mali[mind].size = size;
                    mind++;
                }

                if(!no_acc)
                    memset(ptr, 123, size); 

                cur_brk_pos = sbrk(0);
                if ( ptr < cur_brk_pos) {
                    brk_pos = cur_brk_pos;
                    sbrk_c++;
                    st.sbrk_t++;
                    if (echo)
                        printf("%09d: SBRK ptr=%p sbrk=%p \n",i, ptr, sbrk(0));
                } else {
                    mmap_c++;
                    st.mmap_t++;
                    if (echo)
                        printf("%09d: MMAP ptr=%p sbrk=%p \n", i, ptr, sbrk(0));

                }
            }
            printf("%09d: last allocation ptr=%p sbrk=%p \n", i, ptr, sbrk(0));
            printf("================\n");
            printf("sbrk_c:%08d st.sbrk_t:%08lu mmap_c:%08d st.mmap_t:%08lu\n",
                    sbrk_c, st.sbrk_t, mmap_c, st.mmap_t);
            printf("================\n");
        } else if (!strcmp("vomit", tokens[0])) {
            uint64_t cnt = 0;
            uint64_t vomit_t = 0;
            uint64_t ptr = 0;
            uint64_t size = 0;

            if (tokCnt != 2) {
                printf("Usage: vomit <cnt>\n");
                continue;
            }
            if (parseNumericArg(tokens[1], &cnt) == FALSE) {
                printf("Parsing error - Usage: vomit <cnt>\n");
                continue;
            }

            for(; cnt > 0; cnt--) {
                if (mind == 0 && mali[mind].size == 0) {
                    printf("Nothing in stomach to vomit\n");
                    break;
                }
                mind--;
                ptr = mali[mind].ptr;
                size = mali[mind].size;
                if (size == 0 ) {
                    printf("Error: Zero size allocation");
                    continue;
                }
                if (ptr == 0 ) {
                    printf("Error: Null");
                    continue;
                }
                printf("Vomiting %ld bytes\n", size);
                vomit_t += size;
                st.free_total +=  size;

                free((void*)ptr);

                mali[mind].ptr = 0;
                mali[mind].size = 0;
            }
            printf("==================\n");
            printf("Total vomit: %ld bytes\n", vomit_t);
            printf("==================\n");

        } else if (!strcmp("eatshm", tokens[0])) {
            printf("NOP \n");
            /* uint64_t size,  no_acc; */
            /* void *ptr; */
            /* int sfd; */
            /* char shmpath[16]; */

            /* if (tokCnt < 3) { */
            /*     printf("Usage: eatshm <size> <echo> <no_access>\n"); */
            /*     continue; */
            /* } */

            /* if(shind >= MAX_SHM) { */
            /*     printf("Error: Shm limit %d\n", MAX_SHM); */
            /*     continue; */
            /* } */


            /* if(parseNumericArg(tokens[1], &size) == FALSE || */
            /*         parseNumericArg(tokens[2], &no_acc) == FALSE) { */
            /*     printf("Parsing error: Usage: eatshm <size> <no_access>\n"); */
            /*     continue; */
            /* } */
            /* snprintf(shmpath, 16, "memeaterSHM_%d", shind); */
            /* printf("Eating %"PRId64" byte via shm:%s %s\n", size,  shmpath, */
            /*         no_acc ? " without accessing" : ""); */

            /* sfd = shm_open(shmpath, O_CREAT | O_EXCL | O_RDWR, S_IRUSR | S_IWUSR); */ 
            /* if (sfd == -1) { */
            /*     printf("Error doing shm_open: %d - %s\n", errno, strerror(errno)); */
            /*     st.shm_fail++; */
            /*     shm_unlink(shmpath); */
            /*     continue; */
            /* } */
            /* if (ftruncate(sfd, size) == -1) { */
            /*     printf("Error doing ftruncate: %d - %s\n", errno, strerror(errno)); */
            /*     st.shm_fail++; */
            /*     continue; */
            /* } */
            /* ptr = mmap(NULL, size, PROT_READ | PROT_WRITE, MAP_SHARED, sfd, 0); */
            /* if (ptr == MAP_FAILED) { */
            /*     printf("Error doing mmap: %d - %s\n", errno, strerror(errno)); */
            /*     st.shm_fail++; */
            /*     shm_unlink(shmpath); */
            /*     continue; */
            /* } */

            /* strcpy(shi[shind].name, shmpath); */
            /* shi[shind].fd = sfd; */
            /* shi[shind].size = size; */
            /* shi[shind].ptr = (uint64_t)ptr; */
            /* shind ++; */

            /* if(!no_acc) */
            /*     memset(ptr, 123, size); */ 


        } else if (!strcmp("system", tokens[0])) {
            int retCode = 0;
            if (tokCnt < 2 ) {
                printf("Usage: system <cmd>\n");
                continue;
            }
            printf("%s\n", tokens[1]);
            retCode = system(tokens[1]);
            printf("retCode = %d\n", retCode);
        } else if (!strcmp("open", tokens[0])) {
            int flags;

            if (fd != -1) {
                printf("Error: fd already open\n");
                continue;
            }
            if (tokCnt != 3){
                printf("Usage: open <file> <mode>\n");
                continue;
            }

            if (!strcmp(tokens[2], "r")) {
                flags = O_RDONLY;
            } else if (!strcmp(tokens[2], "w")) {
                flags = O_WRONLY;
            } else if (!strcmp(tokens[2], "rw")) {
                flags = O_RDWR;
            } else {
                printf("Unknown mode: %s\n", tokens[2]);
                continue;
            }

            fd = open(tokens[1], flags);
            if ( fd == -1 ) {
                printf("Failed to  open %s: %d: %s\n", tokens[1],
                        errno, strerror(errno));
            }
            st.file_open ++;
        } else if (!strcmp("close", tokens[0])) {
            if (fd == -1) {
                printf("Nothing open\n");
                continue;
            }
            close(fd);
            st.file_close ++;
            printf("fd closed\n");
        } else if (!strcmp("stat", tokens[0])) {
#define PRN(arg)   printf("%15s : %ld\n", #arg, arg);
            PRN(st.alloc_total);
            PRN(st.free_total);
            printf("%15s : %ld\n", "diff", st.alloc_total - st.free_total);
            PRN(st.alloc_fail);
            PRN(st.alloc_pass);
            PRN(st.sbrk_t);
            PRN(st.mmap_t);
        } else {
            printf("Unknown Cmd: %s\n", tokens[0]);
        }
            
    }
    return EXIT_SUCCESS;
}