Harden VM Mem

1 files changed, 53 insertions, 21 deletions

diff --git a/pkg/lang/vm/mem/mem.go b/pkg/lang/vm/mem/mem.go
index 4cb6fc4..bdcf01f 100644
--- a/pkg/lang/vm/mem/mem.go
+++ b/pkg/lang/vm/mem/mem.go
@@ -2,7 +2,7 @@ package mem
 
 type Mem struct {
 	cells []cell
-	free  []int
+	free  []Ptr
 }
 
 func New() Mem {
@@ -13,69 +13,101 @@ func New() Mem {
 
 	return Mem{
 		cells: cells,
-		free:  make([]int, 0),
+		free:  make([]Ptr, 0),
 	}
 }
 
-func (m *Mem) Allocate(kind CellKind) Ptr {
+func (m *Mem) Allocate(kind CellKind) (Ptr, error) {
 	if len(m.free) > 0 {
 		idx := m.free[len(m.free)-1]
 		m.free = m.free[:len(m.free)-1]
 
 		if m.cells[idx].kind != CellKindEmpty {
-			panic("invalid free cell")
+			return NullPtr, ErrFatalNonFreeCell
 		}
 
 		m.cells[idx].kind = kind
-		return Ptr(idx)
+		return Ptr(idx), nil
 	} else {
+		if len(m.cells) > 10000 {
+			return NullPtr, ErrMemOverflow
+		}
+
 		idx := len(m.cells)
 		m.cells = append(m.cells, cell{kind: kind, refs: 1})
-		return Ptr(idx)
+		return Ptr(idx), nil
 	}
 
 }
 
-func (m *Mem) Set(ptr Ptr, v any) {
-	if ptr >= Ptr(len(m.cells)) {
-		panic("out of bounds")
+func (m *Mem) Set(ptr Ptr, v CellData) error {
+	if err := m.validPtr(ptr); err != nil {
+		return err
 	}
 
 	m.cells[ptr].data = v
+	return nil
 }
 
-func (m *Mem) Get(ptr Ptr) any {
-	if ptr >= Ptr(len(m.cells)) {
-		panic("out of bounds")
+func (m *Mem) Get(ptr Ptr) (CellData, error) {
+	if err := m.validPtr(ptr); err != nil {
+		return nil, err
 	}
 
-	return m.cells[ptr].data
+	return m.cells[ptr].data, nil
 }
 
 func (m *Mem) Is(ptr Ptr, kind CellKind) bool {
 	if ptr >= Ptr(len(m.cells)) {
-		panic("out of bounds")
+		return false
 	}
 
 	return m.cells[ptr].kind == kind
 }
 
-func (m *Mem) Retain(ptr Ptr) {
+func (m *Mem) Kind(ptr Ptr) CellKind {
 	if ptr >= Ptr(len(m.cells)) {
-		panic("out of bounds")
+		return CellKindForbidden
+	}
+
+	return m.cells[ptr].kind
+}
+
+func (m *Mem) Retain(ptr Ptr) error {
+	if err := m.validPtr(ptr); err != nil {
+		return err
 	}
 
 	m.cells[ptr].refs++
+	return nil
 }
 
-func (m *Mem) Release(ptr Ptr) {
-	if ptr >= Ptr(len(m.cells)) {
-		panic("out of bounds")
+func (m *Mem) Release(ptr Ptr) error {
+	if err := m.validPtr(ptr); err != nil {
+		return err
 	}
 
 	m.cells[ptr].refs--
 	if m.cells[ptr].refs == 0 {
-		m.cells[ptr].kind = CellKindEmpty
-		m.free = append(m.free, int(ptr))
+		c := m.cells[ptr].data
+		c.DropCell(m)
+
+		m.cells[ptr] = cell{}
+		m.free = append(m.free, ptr)
 	}
+
+	return nil
+}
+
+func (m *Mem) validPtr(ptr Ptr) error {
+	if ptr >= Ptr(len(m.cells)) {
+		return ErrInvalidMemAccess{ptr}
+	}
+
+	kind := m.cells[ptr].kind
+	if kind == CellKindForbidden || kind == CellKindEmpty {
+		return ErrInvalidMemAccess{ptr}
+	}
+
+	return nil
 }