Installing custom PIT handler (x86 assembler)

.model small

.386

.stack 100h

.data
old_pit_seg     dw      ?
old_pit_off     dw      ?

test_string     db      "TEST $"

counter         dw      0

.code
start:
    xor     ax, ax
    mov     ds, ax
    mov     bx, 20h ; ds:bx points to PIT interrupt vector

    ; read handler address
    cli
    mov     ax, [bx] ; offset to PIT handler
    mov     cx, [bx + 2] ; segment
    sti

    ; push PIT handler segment:offset onto the stack
    push    ax
    push    cx

    ; get address of custom PIT Handler at tick
    mov     ax, offset tick
    mov     cx, cs

    ; set new handler in the interrupt vector table at ds:bx
    cli
    mov     [bx], ax
    mov     [bx + 2], cx
    sti

    ; point ds to program data segment
    mov     ax, @data
    mov     ds, ax

    ; return segment:offset for original PIT interrupt handler and store
    pop     cx
    pop     ax
    mov     old_pit_off, ax
    mov     old_pit_seg, cx

quit:
    ; quit the program
    mov ax,4c00h
    int 21h

tick: ; the new custom PIT interrupt handler (does not handle BIOS time)
    pushf
    pusha

    ; restore program data segment and print "TEST " at the current cursor pos
    mov     ax, @data
    mov     ds, ax
    mov     ax, 0900h
    lea     dx, test_string
    int     21h

    ; increment the counter until 20, if 20 then restore original PIT handler else exit interrupt
    inc     counter
    cmp     counter, 20
    jne     tick_exit

    xor     ax, ax
    mov     es, ax
    mov     bx, 20h ; es:bx now points to PIT vector in interrupt table

    ; get the original segment:offset for the PIT interrupt handler
    mov     ax, old_pit_off
    mov     cx, old_pit_seg

    ; restore original PIT interrupt vector
    cli
    mov     es:[bx], ax
    mov     es:[bx + 2], cx
    sti

tick_exit:
    ; acknowledge tick
    mov     al, 20h
    out     20h, al

    ; restore registers and flags
    popa
    popf
    iret ; return from interrupt

end start