; Name:        shmdemo.asm
; Build:       nasm "-felf64" shmdemo.asm -l shmdemo.lst -o shmdemo.o
;              ld -s -melf_x86_64 -o shmdemo shmdemo.o 
; Description: Demonstration on shared memory segments based on an example from Beej's Guide to IPC.
; Remark:      The C code only use shmget to "get and create" a shared emory segment.  In assembly language
;              we must take care of this.  That's why the syscall shmget appears twice.  Once to get and
;              create a shared memory segment, the second time if the shared memory segment already exists
;              and we want to attach to it.
; Source:      http://beej.us/guide/bgipc/output/html/multipage/shm.html

bits 64

[list -]
    %include "unistd.inc"
    %include "sys/ipc.inc"
    %include "sys/stat.inc"
[list +]

    %define SHM_SIZE 1024  ; make it a 1K shared memory segment

section .bss
    key:        resq    1
    shmid:      resq    1
    data:       resq    1
    datalength: resq    1
    mode:       resq    1
    argc:       resq    1
    arg:        resq    1

section .rodata

    msgUsage:       db  "Usage: shmdemo [data_to_write]",10
    .length:        equ $-msgUsage
    errFtok:        db  "Cannot get IPC key from ftok function.",10
    .length:        equ $-errFtok
    errShmGet:      db  "Cannot create shared memory segment.",10
    .length:        equ $-errShmGet
    errShmAt:       db  "Cannot attach to shared memory segment.",10
    .length:        equ $-errShmAt
    errShmDet:      db  "Cannot detach from shared memory segment.",10
    .length:        equ $-errShmDet
    errRemShm:      db  "Shared Memory Segment cannot be removed.",10
    .length:        equ $-errRemShm
    msgShmRem:      db  "Shared Memory Segment is removed, check with ipcs -m",10
                    db  "eventually kill it with ipcrm -m [id]",10
    .length:        equ $-msgShmRem
    msgWriteTo:     db  "writing to segment : "
    .length:        equ $-msgWriteTo    
    msgReadFrom:    db  "segment contains   : "
    .length:        equ $-msgReadFrom
    eol:            db  10
    filename:       db  "shmdemo",0

section .data

    STAT        stat     

section .text
    global _start:
    pop     rax                                             ;get argc
    cmp     rax,2
    jg      error.usage                                     ;if to much arguments
	pop		rdi												;get programname
    ; if two or less arguments, save the argument count for later use
    mov     qword[argc], rax    

    ; make the key, we pop the pointer to this binary from stack to create a key
    mov		rdi,filename                                             ;path to file
    mov     rsi, 'R'                                        ;proj_id
    call    ftok                                            
    mov     qword[key], rax                                 ;save the retrieved key
    and     rax,rax
    js      error.ftok                                      ;on error just exit
    ; create and connect to the shared memory segment
    ;permissions are in octal notation
    syscall shmget, qword[key], SHM_SIZE, 0o644 | IPC_CREAT
    ; if no errors then the segment created, then we attach to it
    and     rax, rax
    jns     attach
    ; if an error occured, check if the segment already exist and attach
    syscall shmget, qword[key], SHM_SIZE
    ; if still an error then we cannot connect -> other problem
    and     rax, rax                                    ;there is an error
    js      error.shmget
    ; attach to the shared memory segment
    ; first save our id
    mov     qword[shmid], rax
    ; attach to the segment to get a pointer to it
    syscall shmat, qword[shmid], 0, 0
    and     rax, rax    
    js      error.shmat                                 ;there is an error
    ; save pointer to shared memory segment
    mov     qword[data], rax
    ; read or modify the segment, based on the commandline
    mov     rax, qword[argc]
    cmp     rax, 2
    jne     printSegment                                ;no second argument
    ; this code is additional to kill a shared memory segment
    pop     rdi
    mov     eax,dword[rdi]
    xor     eax,"kill"                                  ;in fact remove but kill is 4 bytes
    jz      killShm
    push    rdi
    ; if two argments where on the commandline and it isn't 'kill', the second is the data we want to write
    ; to our segment
    ; first write what we're gonna do
    syscall     write, stdout, msgWriteTo, msgWriteTo.length
    pop     rdi                                         ;get pointer to the data
    mov     qword[arg], rdi                             ;save for later use    
    call    dataLength                                  ;calculate length of data
    cmp     rax,1024                                    ;check length
    jle     storeLength                                 ;length is less then segment size
    mov     rax,1024                                    ;if more than 1024 bytes, just take first 1024 bytes
    mov     qword[datalength],rax                       ;store length of data
    ; write to stdout
    dec     rax                                         ;ignore trailing zero
    syscall write, stdout,qword[arg],qword[datalength]
    syscall write, stdout,eol,1
    ; copy data into shared memory segment
    mov     rsi,qword[arg]                              ;pointer to new data
    mov     rcx,qword[datalength]                       ;length of new data
    mov     rdi,qword[data]                             ;pointer to shared memory segment
    lodsb                                                   ;load byte from new data
    stosb                                                   ;and store in shared memory segment
    loopnz  repeat                                      ;and so on for entire new data size
    ; we print the content of the shared memory segment
    ; it's here we continue the program when only one argument is given
    ; print what we will show on screen
    syscall write, stdout, msgReadFrom, msgReadFrom.length
    ; check if rdi doesn't point to NULL, should not happen but in less controlled software
    ; it can happen.  Just to be sure...
    mov     rdi,qword[data]
    and     rdi, rdi
    jz      error.usage                                 ;if data doesn't points to a zero address
    call    dataLength                                  ;calculate length of it, that's why we keep
                                                            ;the trailing zero
    ; print the data in our shared memory segment on screen, remember rax has the length of it
    dec     rax                                         ;ignore trailing zero
    syscall write, stdout, qword[data], rax
    syscall write, stdout, eol,1                        ;and terminate the line

    ; detach from segment
    syscall shmdt,qword[data]
    js      error.detach
    syscall exit,0
    ; this is optional code to demonstrate how to remove a shared memory segment
    syscall shmctl,qword[shmid],IPC_RMID,0
    and     rax,rax
    js      error.remove
    syscall write,stdout,msgShmRem,msgShmRem.length
    syscall exit,0

; messages to print when something isn't what it's supposed to be    
    syscall write, stderr, errRemShm, errRemShm.length
    syscall exit,1
    syscall write, stderr, msgUsage, msgUsage.length
    syscall exit,1
    syscall write, stderr, errFtok, errFtok.length
    syscall exit,1     
    syscall write, stderr, errShmGet, errShmGet.length
    syscall exit,1
    syscall write, stderr, errShmAt, errShmAt.length
    syscall exit,1
    syscall write, stderr, errShmDet, errShmDet.length
    syscall exit,1

    ; in : rdi is pointer to zero terminated list of bytes
    push    rcx                         ;save rcx
    xor     rcx,rcx                     ;make the biggest number possible
    dec     rcx
    xor     al,al                       ;the zero we will scan for
    repne   scasb                       ;scan byte sequence until zero found
    not     rcx                         ;invert all and pass length
    mov     rax,rcx                     ;via rax to caller
    pop     rcx                         ;restore rcx
    ; key = ((st.st_ino & 0xffff) | ((st.st_dev & 0xff) << 16) | ((proj_id & 0xff) << 24));
    ; save the project id in r8 (will remain after syscalls)

    mov     r8,rsi
    ; open the file
    syscall open,rdi, O_RDONLY
    and     rax,rax
    js      .done
    syscall fstat,rax, stat
    and     rax,rax
    js      .done
    mov     rax,qword [stat.st_ino]     ;get the file size
    and     rax,0xFFFF
    mov     rbx,qword [stat.st_dev]
    and     rbx,0xFF
    shl     rbx,16
    or      rax,rbx
    and     r8,0xFF                                ; R8 = proj_id
    shl     r8,24
    or      rax,r8
    ; rax now contains a key which uniquely identifies the file.