go-vise

commit c0b1f3302eeb366555bb62a32d097be4b6ce0081
parent 1d94353ca0898d7c993610a1556b43fea1bab5a3
Author: lash <dev@holbrook.no>
Date:   Fri, 30 Aug 2024 23:21:19 +0100

Implement prefix protection for static data in db

Diffstat:
M
db/db.go
 | 
12
++++++++++++
M
db/fs.go
 | 
4
++++
M
db/gdbm.go
 | 
3
+++
M
db/mem.go
 | 
4
++++
M
db/pg.go
 | 
4
++++
M
examples/db/main.go
 | 
3
++-
M
resource/db_test.go
 | 
11
+++++++++++

7 files changed, 40 insertions(+), 1 deletion(-)
diff --git a/db/db.go b/db/db.go
@@ -54,6 +54,7 @@ func ToDbKey(typ uint8, b []byte, l *lang.Language) []byte {
 type BaseDb struct {
 	pfx uint8
 	sid []byte
+	unsafe bool
 }
 
 // SetPrefix implements Db.
@@ -66,6 +67,17 @@ func(db *BaseDb) SetSession(sessionId string) {
 	db.sid = append([]byte(sessionId), 0x2E)
 }
 
+func(db *BaseDb) SetSafety(safe bool) {
+	db.unsafe = !safe
+}
+
+func(db *BaseDb) checkPut() bool {
+	if db.unsafe {
+		return true
+	}
+	return db.pfx > datatype_sessioned_threshold
+}
+
 // ToKey creates a DbKey within the current session context.
 func(db *BaseDb) ToKey(key []byte) ([]byte, error) {
 	var b []byte
diff --git a/db/fs.go b/db/fs.go
@@ -2,6 +2,7 @@ package db
 
 import (
 	"context"
+	"errors"
 	"io/ioutil"
 	"os"
 	"path"
@@ -46,6 +47,9 @@ func(fdb *FsDb) Get(ctx context.Context, key []byte) ([]byte, error) {
 
 // Put implements Db
 func(fdb *FsDb) Put(ctx context.Context, key []byte, val []byte) error {
+	if !fdb.checkPut() {
+		return errors.New("unsafe put and safety set")
+	}
 	fp, err := fdb.pathFor(key)
 	if err != nil {
 		return err
diff --git a/db/gdbm.go b/db/gdbm.go
@@ -37,6 +37,9 @@ func(gdb *GdbmDb) Connect(ctx context.Context, connStr string) error {
 
 // Put implements Db
 func(gdb *GdbmDb) Put(ctx context.Context, key []byte, val []byte) error {
+	if !gdb.checkPut() {
+		return errors.New("unsafe put and safety set")
+	}
 	k, err := gdb.ToKey(key)
 	if err != nil {
 		return err
diff --git a/db/mem.go b/db/mem.go
@@ -3,6 +3,7 @@ package db
 import (
 	"context"
 	"encoding/hex"
+	"errors"
 )
 
 // MemDb is a memory backend implementation of the Db interface.
@@ -50,6 +51,9 @@ func(mdb *MemDb) Get(ctx context.Context, key []byte) ([]byte, error) {
 
 // Put implements Db
 func(mdb *MemDb) Put(ctx context.Context, key []byte, val []byte) error {
+	if !mdb.checkPut() {
+		return errors.New("unsafe put and safety set")
+	}
 	k, err := mdb.toHexKey(key)
 	if err != nil {
 		return err
diff --git a/db/pg.go b/db/pg.go
@@ -2,6 +2,7 @@ package db
 
 import (
 	"context"
+	"errors"
 	"fmt"
 
 	"github.com/jackc/pgx/v5/pgxpool"
@@ -41,6 +42,9 @@ func(pdb *PgDb) Connect(ctx context.Context, connStr string) error {
 
 // Put implements Db.
 func(pdb *PgDb) Put(ctx context.Context, key []byte, val []byte) error {
+	if !pdb.checkPut() {
+		return errors.New("unsafe put and safety set")
+	}
 	k, err := pdb.ToKey(key)
 	if err != nil {
 		return err
diff --git a/examples/db/main.go b/examples/db/main.go
@@ -102,6 +102,7 @@ func main() {
 	store.Connect(ctx, dataDir)
 	store.SetSession("xyzzy")
 
+	store.SetSafety(false)
 	err := genCode(ctx, store)
 	if err != nil {
 		panic(err)
@@ -116,6 +117,7 @@ func main() {
 	if err != nil {
 		panic(err)
 	}
+	store.SetSafety(true)
 
 	tg, err := resource.NewDbFuncGetter(store, db.DATATYPE_TEMPLATE, db.DATATYPE_MENU, db.DATATYPE_BIN)
 	if err != nil {
@@ -125,7 +127,6 @@ func main() {
 	rs.WithTemplateGetter(tg.GetTemplate)
 	rs.WithMenuGetter(tg.GetMenu)
 	rs.WithCodeGetter(tg.GetCode)
-	rs.WithCodeGetter(tg.GetCode)
 	rs.AddLocalFunc("say", say)
 
 	ca := cache.NewCache()
diff --git a/resource/db_test.go b/resource/db_test.go
@@ -23,11 +23,18 @@ func TestDb(t *testing.T) {
 		t.Fatal("expected error")
 	}
 
+
 	store.SetPrefix(db.DATATYPE_TEMPLATE)
 	err = store.Put(ctx, []byte("foo"), []byte("bar"))
+	if err == nil {
+		t.Fatal("expected error")
+	}
+	store.SetSafety(false)
+	err = store.Put(ctx, []byte("foo"), []byte("bar"))
 	if err != nil {
 		t.Fatal(err)
 	}
+	store.SetSafety(true)
 	s, err = rs.GetTemplate(ctx, "foo")
 	if err != nil {
 		t.Fatal(err)
@@ -38,10 +45,12 @@ func TestDb(t *testing.T) {
 
 	// test support check
 	store.SetPrefix(db.DATATYPE_BIN)
+	store.SetSafety(false)
 	err = store.Put(ctx, []byte("xyzzy"), []byte("deadbeef"))
 	if err != nil {
 		t.Fatal(err)
 	}
+	store.SetSafety(true)
 
 	rs.WithCodeGetter(tg.GetCode)
 	b, err := rs.GetCode(ctx, "xyzzy")
@@ -69,10 +78,12 @@ func TestDb(t *testing.T) {
 		t.Fatal(err)
 	}
 	store.SetPrefix(db.DATATYPE_MENU)
+	store.SetSafety(false)
 	err = store.Put(ctx, []byte("inky"), []byte("pinky"))
 	if err != nil {
 		t.Fatal(err)
 	}
+	store.SetSafety(true)
 	rs.WithMenuGetter(tg.GetMenu)
 
 }