compooter-enhensa/3-arithmetic/main.c

#include "../common.h"

static void
print_inst(int instruction_offset, struct instruction *inst, int *labels, int label_count)
{
    if (inst->opcode > 0xF) {
        for (int i = 0; i < label_count; ++i) {
            if (instruction_offset + inst->size + inst->jump_offset == labels[i]) {
                printf("%s label_%d\n", OP_NAMES[inst->opcode], i);
                return;
            }
        }
        
        CRASH_SHOULDNOT();
    }
    
    const char **regs = inst->wide ? REGS_W1 : REGS_W0;
    struct rm dest = inst->dst;
    struct rm src = inst->src;
    
    if (inst->swap) {
        struct rm buf = dest;
        dest = src;
        src = buf;
    }
    
    printf("%s ", OP_NAMES[inst->opcode]);
    
    /* Dest can not be an immediate, so dest.imm indicates if there's a displacement */
    
    if (dest.reg1 == NO_REGISTER && dest.reg2 == NO_REGISTER && dest.disp != 0) {
        printf("[%d]", dest.disp);
    } else if (dest.reg2 == NO_REGISTER && dest.imm == 0) {
        printf("%s", regs[dest.reg1]);
    } else if (dest.reg2 != NO_REGISTER && dest.imm == 0) {
        printf("[%s + %s]", REGS_W1[dest.reg1], REGS_W1[dest.reg2]);
    } else if (dest.reg2 == NO_REGISTER && dest.imm != 0) {
        printf("[%s + %d]", REGS_W1[dest.reg1], dest.disp);
    } else if (dest.reg2 != NO_REGISTER && dest.imm != 0) {
        printf("[%s + %s + %d]", REGS_W1[dest.reg1], REGS_W1[dest.reg2], dest.disp);
    }
    
    printf(", ");
    
    if (src.reg1 == NO_REGISTER) {
        if (src.disp == 0) {
            printf("%s %d", inst->wide ? "word" : "byte", src.imm); /* only the src can be an immediate */
        } else {
            printf("[%d]", src.disp);
        }
    } else {
        if (src.reg2 == NO_REGISTER && src.imm == 0) {
            printf("%s", regs[src.reg1]);
        } else if (src.reg2 != NO_REGISTER && src.imm == 0) {
            printf("[%s + %s]", REGS_W1[src.reg1], REGS_W1[src.reg2]);
        } else if (src.reg2 == NO_REGISTER && src.imm != 0) {
            printf("[%s + %d]", REGS_W1[src.reg1], src.disp);
        } else if (src.reg2 != NO_REGISTER && src.imm != 0) {
            printf("[%s + %s + %d]", REGS_W1[src.reg1], REGS_W1[src.reg2], src.disp);
        }
    }
    
    printf("\n");
}

static int
decode_rm(u8 *data, int offset, struct rm *location)
{
    int advance = 0;
    
    u8 b2       = data[offset + 1];
    u8 mod      = (b2 & 0xC0) >> 6;
    u8 rm_field = b2 & 0x7;
    
    if (mod == MOD_REGISTER) {
        location->reg1 = rm_field;
        location->reg2 = -1;
        location->disp = 0;
        location->imm = 0;
        
        advance = 2;
    } else if (mod == MOD_MEMORY) {
        s16 b3 = data[offset + 2];
        s16 b4 = data[offset + 3];
        
        if (rm_field == 0x6) {
            location->reg1 = NO_REGISTER;
            location->reg2 = NO_REGISTER;
            location->disp = (b4 << 8) | b3;
            location->imm = 0;
            advance = 4;
        } else {
            location->reg1 = MOD_PAIRS_FIRST[rm_field];
            location->reg2 = MOD_PAIRS_SECOND[rm_field];
            location->disp = 0;
            location->imm = 1;
            advance = 2;
        }
    } else if (mod == MOD_MEMORY_8) {
        s8 b3 = data[offset + 2];
        
        location->reg1 = MOD_PAIRS_FIRST[rm_field];
        location->reg2 = MOD_PAIRS_SECOND[rm_field];
        location->disp = b3;
        location->imm = 1;
        
        advance = 3;
    } else if (mod == MOD_MEMORY_16) {
        s16 b3 = data[offset + 2];
        s16 b4 = data[offset + 3];
        
        location->reg1 = MOD_PAIRS_FIRST[rm_field];
        location->reg2 = MOD_PAIRS_SECOND[rm_field];
        location->disp = (b4 << 8) | b3;
        location->imm = 1;
        
        advance = 4;
    }
    
    return(advance);
}

static int
decode_and_print_add_sub_cmp_jmp(u8 *data, u64 size, struct instruction *stream)
{
    printf("bits 16\n\n");
    
    /* We still do not handle corrupt binaries */
    
    u64 offset = 0;
    int stream_at = 0;
    
    while (offset < size) {
        int saved_offset = offset;
        
        u8 b1 = data[offset + 0];
        u8 b2 = data[offset + 1];
        
        u8 antiop = b1 & ANTIOP_MASK;
        u8 op = 0;
        
        if (antiop == 0x00 || antiop == 0x4) {
            op = (b1 & 0x38) >> 3;
        } else {
            op = (b2 & 0x38) >> 3;
        }
        
        if (antiop == 0x00) {
            u8 W_flag = b1 & 0x1;
            u8 D_flag = (b1 & 0x2) >> 1;
            u8 reg_field = (b2 & 0x38) >> 3;
            
            struct rm dest = { 0 };
            struct rm src = { 0 };
            
            dest.reg1 = reg_field;
            dest.reg2 = -1;
            dest.disp = 0;
            dest.imm = 0;
            
            offset += decode_rm(data, offset, &src);
            
            struct instruction inst = { 
                .size = offset - saved_offset,
                .opcode = op,
                .dst = dest,
                .src = src,
                .wide = W_flag,
                .swap = !D_flag
            };
            
            stream[stream_at++] = inst;
            
            continue;
        } else if (antiop == 0x80) {
            u8 W_flag = b1 & 0x1;
            u8 S_flag = (b1 & 0x2) >> 1;
            u8 D_flag = 1; /* Immediate to register-memory always has mod/rm as destination */
            
            struct rm dest = { 0 };
            struct rm src = { 0 };
            
            offset += decode_rm(data, offset, &dest);
            
            s16 immediate = 0;
            
            if (S_flag && !W_flag) {
                s16 b5 = data[offset];
                s16 b6 = data[offset + 1];
                immediate = (b6 << 8) | b5;
                offset += 2;
            } else {
                s16 b5 = data[offset];
                immediate = b5;
                offset += 1;
            }
            
            src.reg1 = NO_REGISTER;
            src.imm = immediate;
            
            struct instruction inst = { 
                .size = offset - saved_offset,
                .opcode = op,
                .dst = dest,
                .src = src,
                .wide = W_flag,
                .swap = !D_flag
            };
            
            stream[stream_at++] = inst;
            
            continue;
        } else if (antiop == 0x4) {
            u8 W_flag = b1 & 0x1;
            s16 immediate = 0;
            
            struct rm dest = { 0 };
            struct rm src = { 0 };
            
            dest.reg1 = AX;
            dest.reg2 = NO_REGISTER;
            dest.disp = 0;
            dest.imm = 0;
            
            if (W_flag) {
                s16 b3 = data[offset + 2];
                immediate= (b3 << 8) | b2;
                offset += 3;
            } else {
                immediate = b2;
                offset += 2;
            }
            
            src.reg1 = NO_REGISTER;
            src.reg2 = NO_REGISTER;
            src.disp = 0;
            src.imm = immediate;
            
            struct instruction inst = { 
                .size = offset - saved_offset,
                .opcode = op,
                .dst = dest,
                .src = src,
                .wide = W_flag,
                .swap = 0,
            };
            
            stream[stream_at++] = inst;
            
            continue;
        }
        
        offset += 2;
        
        struct instruction inst = { 
            .size = offset - saved_offset,
            .opcode = b1,
            .jump_offset = (s8) b2
        };
        
        stream[stream_at++] = inst;
    }
    
    return(stream_at);
}

static int
insert_labels(struct instruction *in_stream, int instruction_count, int *labels)
{
    int instruction_offset = 0;
    int label_count = 0;
    
    for (int i = 0; i < instruction_count; ++i) {
        struct instruction *inst = in_stream + i;
        
        if (inst->opcode > 0xF) {
            /* jmp-like */
            int target = instruction_offset + inst->size + inst->jump_offset;
            int label_exists = 0;
            
            for (int j = 0; j < label_count; ++j) {
                if (labels[j] == target) {
                    label_exists = 1;
                    break;
                }
            }
            
            if (!label_exists) {
                labels[label_count++] = target;
            }
        }
        
        instruction_offset += inst->size;
    }
    
    return(label_count);
}

static void
print_instruction_stream(struct instruction *stream, int instruction_count, 
                         int *labels, int label_count)
{
    int instruction_offset = 0;
    
    for (int i = 0; i < instruction_count; ++i) {
        struct instruction *inst = stream + i;
        
        for (int j = 0; j < label_count; ++j) {
            if (labels[j] == instruction_offset) {
                printf("label_%d:\n", j);
                break;
            }
        }
        
        print_inst(instruction_offset, stream + i, labels, label_count);
        
        instruction_offset += inst->size;
    }
}

int
main(int argc, char **argv)
{
    if (argc != 2) {
        fprintf(stderr, "Usage: %s binary_file\n", argv[0]);
        return(1);
    }
    
    char *filename = argv[1];
    int fd = open(filename, O_RDONLY);
    
    if (fd == -1) {
        perror("open");
        return(1);
    }
    
    u64 size = file_size(fd);
    u8 *data = read_file(fd, size);
    
    struct instruction stream[MAX_INSTRUCTIONS] = { 0 };
    int labels[MAX_LABELS] = { 0 };
    
    int instruction_count = decode_and_print_add_sub_cmp_jmp(data, size, stream);
    int label_count = insert_labels(stream, instruction_count, labels);
    
    print_instruction_stream(stream, instruction_count, labels, label_count);
    
    return(0);
}