From 0a7700112f82e634a957685bee0cbaa3458f4945 Mon Sep 17 00:00:00 2001
From: Mel <einebeere@gmail.com>
Date: Sat, 28 May 2022 01:22:17 +0000
Subject: Harden VM Mem

---
 pkg/lang/vm/mem/cell.go   | 25 +++++++++++++++-
 pkg/lang/vm/mem/errors.go | 17 +++++++++++
 pkg/lang/vm/mem/mem.go    | 74 +++++++++++++++++++++++++++++++++--------------
 3 files changed, 94 insertions(+), 22 deletions(-)
 create mode 100644 pkg/lang/vm/mem/errors.go

(limited to 'pkg/lang/vm/mem')

diff --git a/pkg/lang/vm/mem/cell.go b/pkg/lang/vm/mem/cell.go
index 3f052d5..fcf9a50 100644
--- a/pkg/lang/vm/mem/cell.go
+++ b/pkg/lang/vm/mem/cell.go
@@ -13,8 +13,31 @@ const (
 	CellKindForbidden
 )
 
+func (c CellKind) String() string {
+	switch c {
+	case CellKindEmpty:
+		return "empty"
+	case CellKindString:
+		return "string"
+	case CellKindArray:
+		return "array"
+	case CellKindEnv:
+		return "env"
+	case CellKindOutlet:
+		return "outlet"
+	case CellKindForbidden:
+		return "forbidden"
+	default:
+		return "unknown"
+	}
+}
+
 type cell struct {
 	kind CellKind
 	refs int
-	data any
+	data CellData
+}
+
+type CellData interface {
+	DropCell(*Mem)
 }
diff --git a/pkg/lang/vm/mem/errors.go b/pkg/lang/vm/mem/errors.go
new file mode 100644
index 0000000..e7a3f8f
--- /dev/null
+++ b/pkg/lang/vm/mem/errors.go
@@ -0,0 +1,17 @@
+package mem
+
+import "errors"
+
+var (
+	ErrMemOverflow = errors.New("memory overflow, cannot allocate more than 10000 memory cells")
+
+	ErrFatalNonFreeCell = errors.New("non-free cell marked as free")
+)
+
+type ErrInvalidMemAccess struct {
+	Ptr Ptr
+}
+
+func (e ErrInvalidMemAccess) Error() string {
+	return "invalid memory access at " + e.Ptr.String()
+}
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
 }
-- 
cgit 1.4.1