From 440187c2a82dd2e335c172484bc22364ee755dd6 Mon Sep 17 00:00:00 2001
From: Nero <41307858+nero@users.noreply.github.com>
Date: Wed, 13 Jan 2021 18:32:17 +0000
Subject: [PATCH] Harden drive.asm against unexpected DS and ES values

---
 boot/kernel.asm  |   2 +
 kernel/drive.asm | 116 +++++++++++++++++++++++++++--------------------
 kernel/fcb.asm   |  38 ++++++++--------
 3 files changed, 87 insertions(+), 69 deletions(-)

diff --git a/boot/kernel.asm b/boot/kernel.asm
index 908c5db..ba05ad7 100644
--- a/boot/kernel.asm
+++ b/boot/kernel.asm
@@ -38,7 +38,9 @@ relinit:	; print banner
 		db "RDOS ",2,2, 0x0A, 0x0D, 0
 
 		push dx
+		push ds
 		call dinit
+		pop ds
 		pop dx
 
 		mov word [0x21*4], int21
diff --git a/kernel/drive.asm b/kernel/drive.asm
index 75944f6..0d9a355 100644
--- a/kernel/drive.asm
+++ b/kernel/drive.asm
@@ -34,63 +34,66 @@ section .text
 
 		; initial setup for disk i/o
 dinit:		; copy previously set DPT to our data area
-		lds si, [4*0x1E]
+		push cs
+		pop es
+		lds si, [es:4*0x1E]
 		mov di, dpt
 		mov cx, 11
 		rep movsb
-		; restore DS
-		xor ax, ax
-		mov ds, ax
 		; set vector
-		mov word [4*0x1E], dpt
-		mov word [4*0x1E+2], ds
+		mov word [cs:4*0x1E], dpt
+		mov word [cs:4*0x1E+2], ds
 		; set to invalid drive
-		mov byte [dsknum], 0xFF
+		mov byte [cs:dsknum], 0xFF
 		ret
 
 		; int 13 stub
 int13:		; debug output
-		;push bx
-		;push dx
-		;push cx
-		;push ax
-		;call printf
-		;db "int13 CALL AX=",2," CX=",2," DX=",2," BX=",2,0x0A,0x0D,0
+		push bx
+		push dx
+		push cx
+		push ax
+		call printf
+		db "int13 CALL AX=",2," CX=",2," DX=",2," BX=",2,0x0A,0x0D,0
 		; do the call
 		int 0x13
-		;jc .err
+		jc .err
+		ret
+.err:		push ax
+		call printf
+		db "DISK ERR AX=",2,0x0A,0x0D,0
+		stc
 		ret
-.err:		;push ax
-		;call printf
-		;db "DISK ERR AX=",2,0x0A,0x0D,0
-		;stc
-		;ret
 
 		; get drive parameters
 		; read BIOS int13h/AH=8 values
 		; DPT data gets copied over our dpt
 		; CHS data gets inserted into our bpb
-getprm:		test byte [dskflag], 4
-		jnz .nodpt
+getprm:		; skip if EBIOS extension enabled
+		test byte [cs:dskflag], 4
+		jnz .ret
 		; do the query
 		mov ah, 8
-		mov dl, [biosnum]
+		mov dl, [cs:biosnum]
 		call int13
 		; bail out if error
-		jc .nodpt
+		jc .ret
+		; ignore CHS values if odd
+		test cl, cl
+		jz .nochs
 		; get and store sector number
 		and cx, 0x3F
-		mov word [bpb+BPBSPT], cx
+		mov word [cs:bpb+BPBSPT], cx
 		; get and store number of heads
 		xchg dl, dh
 		and dx, 0xFF
 		inc dx
-		mov [bpb+BPBNOS], dx
+		mov [cs:bpb+BPBNOS], dx
 		; test if DPT ptr is non-zero
 .nochs:		mov ax, es
 		or ax, di
 		test ax, ax
-		jz .nodpt
+		jz .ret
 		; DS:SI = ES:DI
 		mov ax, es
 		mov ds, ax
@@ -102,20 +105,21 @@ getprm:		test byte [dskflag], 4
 		; do the copy
 		mov cx, 11
 		rep movsb
-.nodpt:		; restore segment registers
-		xor ax, ax
-		mov ds, ax
-		mov es, ax
-		ret
+.ret:		ret
 
 		; log in drive
 		; IN	dl	drive number
-logdrv:		; dont do anything if drive already selected
+logdrv:		; DS := CS
+		push cs
+		pop ds
+		; dont do anything if drive already selected
 		cmp dl, [dsknum]
 		je logfdd.ret
 		; clear out current contents
 		push dx
+		push ds
 		call flush
+		pop ds
 		mov ax, 0xFFFF
 		mov [dskseek], ax
 		mov [dskseek+2], ax
@@ -130,16 +134,22 @@ logdrv:		; dont do anything if drive already selected
 logfdd:		; save info for bios
 		mov [biosnum], dl
 		; reset dpt to defaults
+		push ds
 		call getprm
+		pop ds
 		; set default geometry (1.44 MB floppy)
 		mov word [bpb+BPBNOS], 2
 		mov word [bpb+BPBSPT], 18
 		; load boot sector
+		push ds
 		xor ax, ax
 		xor dx, dx
 		call mapabs
+		pop ds
 		; copy bios parameter block
 		; TODO: guess from first byte of FAT if BPB invalid
+		push ds
+		pop es
 		mov si, dskbuf+BPBOFF
 		mov di, bpb
 		mov cx, BPBSIZ4
@@ -153,6 +163,7 @@ logfdd:		; save info for bios
 		mov [bpb+BPBHS+2], ax
 .ret:		ret
 
+		; assumes that DS == CS
 loghdd:		sub dl, 2
 		cmp dl, 4
 		jnc logerr
@@ -187,17 +198,20 @@ loghdd:		sub dl, 2
 		mov cl, 4
 		sal bx, cl
 		; bail out if no partition
-		cmp byte [dskbuf+0x1be+bx+4], 0
+		cmp byte [cs:dskbuf+0x1be+bx+4], 0
 		je logerr
 		; load partition offset
-		mov ax, [dskbuf+0x1be+bx+8]
-		mov dx, [dskbuf+0x1be+bx+8+2]
+		mov ax, [cs:dskbuf+0x1be+bx+8]
+		mov dx, [cs:dskbuf+0x1be+bx+8+2]
 		; save to to stack
 		push dx
 		push ax
 		; load vbr
 		call mapabs
 		; copy bpb
+		mov ax, cs
+		mov ds, ax
+		mov es, ax
 		mov si, dskbuf+BPBOFF
 		mov di, bpb
 		mov cx, BPBSIZ4
@@ -207,15 +221,15 @@ loghdd:		sub dl, 2
 		; fix partition offset
 		pop ax
 		pop dx
-		mov [bpb+BPBHS], ax
-		mov [bpb+BPBHS+2], dx
+		mov [cs:bpb+BPBHS], ax
+		mov [cs:bpb+BPBHS+2], dx
 		ret
 
 logerr:		stc
 		ret
 
 mapclu:		; counting from begin of cluster area
-		mov bx, [bpb+BPBRDE]
+		mov bx, [cs:bpb+BPBRDE]
 		mov cl, 4
 		shr bx, cl ; 32 bytes per entry
 		add ax, bx
@@ -223,22 +237,22 @@ mapclu:		; counting from begin of cluster area
 maprd:		; counting from beginning of dir
 		; add fat table sizes
 		xor ch, ch
-		mov cl, byte [bpb+BPBFN]
-.loop:		add ax, [bpb+BPBFS]
+		mov cl, byte [cs:bpb+BPBFN]
+.loop:		add ax, [cs:bpb+BPBFS]
 		adc dx, 0
 		loop .loop
 mapfat:		; counting from beginning of FAT
 		; add reserved sector count
-		add ax, [bpb+BPBRSC]
+		add ax, [cs:bpb+BPBRSC]
 		adc dx, 0
 map:		; count from partition start
-		add ax, [bpb+BPBHS]
-		add dx, [bpb+BPBHS+2]
+		add ax, [cs:bpb+BPBHS]
+		add dx, [cs:bpb+BPBHS+2]
 mapabs:		; absolute sector count
 		; skip doing a read if sector number matches
-		cmp ax, [dskseek]
+		cmp ax, [cs:dskseek]
 		jne read
-		cmp dx, [dskseek+2]
+		cmp dx, [cs:dskseek+2]
 		jne read
 		ret
 
@@ -248,8 +262,8 @@ read:		push ax
 		pop dx
 		pop ax
 		; store the sector number
-		mov [dskseek], ax
-		mov [dskseek+2], dx
+		mov [cs:dskseek], ax
+		mov [cs:dskseek+2], dx
 		; do the actual read
 
 		; low level read and write
@@ -259,6 +273,10 @@ _read:		mov ch, 2
 		db 0x3D ; cmp ax, imm16: causes next instr to be skipped
 _write:		mov ch, 3
 		mov cl, 1 ; read len
+		; DS := ES := CS
+		mov ax, cs
+		mov ds, ax
+		mov es, ax
 		; check if ebios supported
 		test byte [dskflag], 4
 		jz .l00
@@ -322,11 +340,11 @@ _write:		mov ch, 3
 		ret
 
 		; mark dskbuf as containing unwritten changes
-dirty:		or byte [dskflag], 1 ; dirty
+dirty:		or byte [cs:dskflag], 1 ; dirty
 		ret
 
 		; flush buffer if dirty
-flush:		test byte [dskflag], 1
+flush:		test byte [cs:dskflag], 1
 		jz .ret
 		; TODO: error handling & retries
 		jmp _write
diff --git a/kernel/fcb.asm b/kernel/fcb.asm
index 87b8131..9ab93b6 100644
--- a/kernel/fcb.asm
+++ b/kernel/fcb.asm
@@ -3,7 +3,7 @@ section .text
 		; set default drive for current process
 		; IN	dl	drive number
 setdd:		push ds
-		mov ds, [ss:curpsp]
+		mov ds, [cs:curpsp]
 		mov [PSPDD], dl
 		pop ds
 		ret
@@ -11,11 +11,23 @@ setdd:		push ds
 		; get default drive for current process
 		; OUT	al	drive number
 getdd:		push ds
-		mov ds, [ss:curpsp]
+		mov ds, [cs:curpsp]
 		mov al, [PSPDD]
 		pop ds
 		ret
 
+		; auto-complete drive field in fcb
+		; IN	ds:dx	far ptr FCB
+		; OUT	si	copy of dx
+fixfcb:		mov si, dx
+		cmp byte [si], 0
+		jne .ret
+		call getdd
+		inc dl
+		mov byte [si], dl
+.ret:		mov dx, si
+		ret
+
 		; Load root directory entry
 		; IN	ax	number of directory entry
 lddir:		push ax
@@ -56,22 +68,8 @@ fnfile:		mov ax, [es:bx+FCBDEN]
 .err:		stc
 		ret
 
-		; initialize a FCB for directory scanning
-open:		mov word [es:bx+FCBDEN], 0
-.search:	call fnfile
-		mov al, 0xFF
-		jc .ret
-		push si
-		lea di, [bx+1]
-		mov cx, 11
-		rep cmpsb
-		pop si
-		jne .search
-		mov al, 0
-.ret:		ret
-
-putfn:		mov cx, 11
-.loop:		lodsb
-		call pputc
-		loop .loop
+		; Open file using FCB
+		; IN	DS:DX	far ptr to FCB
+open:		call fixfcb
+		int 3
 		ret