1 files changed, 177 insertions, 107 deletions

diff --git a/internal/state/multi.go b/internal/state/multi.go
index f014b3d1..16afc086 100644
--- a/internal/state/multi.go
+++ b/internal/state/multi.go
@@ -3,54 +3,135 @@ package state
 import (
 	"encoding/gob"
 	"errors"
+	"fmt"
 	"io/fs"
 	"os"
 	"path"
 	"strconv"
 	"sync"
 	"syscall"
+
+	"git.ophivana.moe/security/fortify/fst"
+	"git.ophivana.moe/security/fortify/internal/fmsg"
 )
 
 // fine-grained locking and access
 type multiStore struct {
-	path []string
+	base string
 
-	// created/opened by prepare
-	lockfile *os.File
-	// enforce prepare method
-	init sync.Once
-	// error returned by prepare
-	initErr error
+	// initialised backends
+	backends *sync.Map
 
-	lock sync.Mutex
+	lock sync.RWMutex
 }
 
-func (s *multiStore) Do(f func(b Backend)) (bool, error) {
-	s.init.Do(s.prepare)
-	if s.initErr != nil {
-		return false, s.initErr
-	}
+func (s *multiStore) Do(aid int, f func(c Cursor)) (bool, error) {
+	s.lock.RLock()
+	defer s.lock.RUnlock()
 
-	s.lock.Lock()
-	defer s.lock.Unlock()
+	// load or initialise new backend
+	b := new(multiBackend)
+	if v, ok := s.backends.LoadOrStore(aid, b); ok {
+		b = v.(*multiBackend)
+	} else {
+		b.lock.Lock()
+		b.path = path.Join(s.base, strconv.Itoa(aid))
+
+		// ensure directory
+		if err := os.MkdirAll(b.path, 0700); err != nil && !errors.Is(err, fs.ErrExist) {
+			s.backends.CompareAndDelete(aid, b)
+			return false, err
+		}
+
+		// open locker file
+		if l, err := os.OpenFile(b.path+".lock", os.O_RDWR|os.O_CREATE, 0600); err != nil {
+			s.backends.CompareAndDelete(aid, b)
+			return false, err
+		} else {
+			b.lockfile = l
+		}
+		b.lock.Unlock()
+	}
 
-	// lock store
-	if err := s.lockFile(); err != nil {
+	// lock backend
+	if err := b.lockFile(); err != nil {
 		return false, err
 	}
 
-	// initialise new backend for caller
-	b := new(multiBackend)
-	b.path = path.Join(s.path...)
+	// expose backend methods without exporting the pointer
+	c := new(struct{ *multiBackend })
+	c.multiBackend = b
 	f(b)
-	// disable backend
-	b.lock.Lock()
+	// disable access to the backend on a best-effort basis
+	c.multiBackend = nil
 
-	// unlock store
-	return true, s.unlockFile()
+	// unlock backend
+	return true, b.unlockFile()
 }
 
-func (s *multiStore) lockFileAct(lt int) (err error) {
+func (s *multiStore) List() ([]int, error) {
+	var entries []os.DirEntry
+
+	// read base directory to get all aids
+	if v, err := os.ReadDir(s.base); err != nil && !errors.Is(err, os.ErrNotExist) {
+		return nil, err
+	} else {
+		entries = v
+	}
+
+	aidsBuf := make([]int, 0, len(entries))
+	for _, e := range entries {
+		// skip non-directories
+		if !e.IsDir() {
+			fmsg.VPrintf("skipped non-directory entry %q", e.Name())
+			continue
+		}
+
+		// skip non-numerical names
+		if v, err := strconv.Atoi(e.Name()); err != nil {
+			fmsg.VPrintf("skipped non-aid entry %q", e.Name())
+			continue
+		} else {
+			if v < 0 || v > 9999 {
+				fmsg.VPrintf("skipped out of bounds entry %q", e.Name())
+				continue
+			}
+
+			aidsBuf = append(aidsBuf, v)
+		}
+	}
+
+	return append([]int(nil), aidsBuf...), nil
+}
+
+func (s *multiStore) Close() error {
+	s.lock.Lock()
+	defer s.lock.Unlock()
+
+	var errs []error
+	s.backends.Range(func(_, value any) bool {
+		b := value.(*multiBackend)
+		errs = append(errs, b.close())
+		return true
+	})
+
+	return errors.Join(errs...)
+}
+
+type multiBackend struct {
+	path string
+
+	// created/opened by prepare
+	lockfile *os.File
+
+	lock sync.RWMutex
+}
+
+func (b *multiBackend) filename(id *fst.ID) string {
+	return path.Join(b.path, id.String())
+}
+
+func (b *multiBackend) lockFileAct(lt int) (err error) {
 	op := "LockAct"
 	switch lt {
 	case syscall.LOCK_EX:
@@ -60,7 +141,7 @@ func (s *multiStore) lockFileAct(lt int) (err error) {
 	}
 
 	for {
-		err = syscall.Flock(int(s.lockfile.Fd()), lt)
+		err = syscall.Flock(int(b.lockfile.Fd()), lt)
 		if !errors.Is(err, syscall.EINTR) {
 			break
 		}
@@ -68,103 +149,80 @@ func (s *multiStore) lockFileAct(lt int) (err error) {
 	if err != nil {
 		return &fs.PathError{
 			Op:   op,
-			Path: s.lockfile.Name(),
+			Path: b.lockfile.Name(),
 			Err:  err,
 		}
 	}
 	return nil
 }
 
-func (s *multiStore) lockFile() error {
-	return s.lockFileAct(syscall.LOCK_EX)
-}
-
-func (s *multiStore) unlockFile() error {
-	return s.lockFileAct(syscall.LOCK_UN)
-}
-
-func (s *multiStore) prepare() {
-	s.initErr = func() error {
-		prefix := path.Join(s.path...)
-		// ensure directory
-		if err := os.MkdirAll(prefix, 0700); err != nil && !errors.Is(err, fs.ErrExist) {
-			return err
-		}
-
-		// open locker file
-		if f, err := os.OpenFile(prefix+".lock", os.O_RDWR|os.O_CREATE, 0600); err != nil {
-			return err
-		} else {
-			s.lockfile = f
-		}
-
-		return nil
-	}()
-}
-
-func (s *multiStore) Close() error {
-	s.lock.Lock()
-	defer s.lock.Unlock()
-
-	err := s.lockfile.Close()
-	if err == nil || errors.Is(err, os.ErrInvalid) || errors.Is(err, os.ErrClosed) {
-		return nil
-	}
-	return err
-}
-
-type multiBackend struct {
-	path string
-
-	lock sync.RWMutex
+func (b *multiBackend) lockFile() error {
+	return b.lockFileAct(syscall.LOCK_EX)
 }
 
-func (b *multiBackend) filename(pid int) string {
-	return path.Join(b.path, strconv.Itoa(pid))
+func (b *multiBackend) unlockFile() error {
+	return b.lockFileAct(syscall.LOCK_UN)
 }
 
 // reads all launchers in simpleBackend
 // file contents are ignored if decode is false
-func (b *multiBackend) load(decode bool) ([]*State, error) {
+func (b *multiBackend) load(decode bool) (Entries, error) {
 	b.lock.RLock()
 	defer b.lock.RUnlock()
 
-	var (
-		r []*State
-		f *os.File
-	)
-
-	// read directory contents, should only contain files named after PIDs
+	// read directory contents, should only contain files named after ids
+	var entries []os.DirEntry
 	if pl, err := os.ReadDir(b.path); err != nil {
 		return nil, err
 	} else {
-		for _, e := range pl {
-			// run in a function to better handle file closing
-			if err = func() error {
-				// open state file for reading
-				if f, err = os.Open(path.Join(b.path, e.Name())); err != nil {
-					return err
-				} else {
-					defer func() {
-						if f.Close() != nil {
-							// unreachable
-							panic("foreign state file closed prematurely")
-						}
-					}()
+		entries = pl
+	}
+
+	// allocate as if every entry is valid
+	// since that should be the case assuming no external interference happens
+	r := make(Entries, len(entries))
+
+	for _, e := range entries {
+		if e.IsDir() {
+			return nil, fmt.Errorf("unexpected directory %q in store", e.Name())
+		}
+
+		id := new(fst.ID)
+		if err := fst.ParseAppID(id, e.Name()); err != nil {
+			return nil, err
+		}
+
+		// run in a function to better handle file closing
+		if err := func() error {
+			// open state file for reading
+			if f, err := os.Open(path.Join(b.path, e.Name())); err != nil {
+				return err
+			} else {
+				defer func() {
+					if f.Close() != nil {
+						// unreachable
+						panic("foreign state file closed prematurely")
+					}
+				}()
 
-					var s State
-					r = append(r, &s)
+				s := new(State)
+				r[*id] = s
 
-					// append regardless, but only parse if required, used to implement Len
-					if decode {
-						return gob.NewDecoder(f).Decode(&s)
-					} else {
-						return nil
+				// append regardless, but only parse if required, used to implement Len
+				if decode {
+					if err = gob.NewDecoder(f).Decode(s); err != nil {
+						return err
+					}
+
+					if s.ID != *id {
+						return fmt.Errorf("state entry %s has unexpected id %s", id, &s.ID)
 					}
 				}
-			}(); err != nil {
-				return nil, err
+
+				return nil
 			}
+		}(); err != nil {
+			return nil, err
 		}
 	}
 
@@ -180,7 +238,7 @@ func (b *multiBackend) Save(state *State) error {
 		return errors.New("state does not contain config")
 	}
 
-	statePath := b.filename(state.PID)
+	statePath := b.filename(&state.ID)
 
 	// create and open state data file
 	if f, err := os.OpenFile(statePath, os.O_RDWR|os.O_CREATE|os.O_EXCL, 0600); err != nil {
@@ -197,14 +255,14 @@ func (b *multiBackend) Save(state *State) error {
 	}
 }
 
-func (b *multiBackend) Destroy(pid int) error {
+func (b *multiBackend) Destroy(id fst.ID) error {
 	b.lock.Lock()
 	defer b.lock.Unlock()
 
-	return os.Remove(b.filename(pid))
+	return os.Remove(b.filename(&id))
 }
 
-func (b *multiBackend) Load() ([]*State, error) {
+func (b *multiBackend) Load() (Entries, error) {
 	return b.load(true)
 }
 
@@ -214,9 +272,21 @@ func (b *multiBackend) Len() (int, error) {
 	return len(rn), err
 }
 
-// NewSimple returns an instance of a file-based store.
-func NewSimple(runDir string, prefix ...string) Store {
+func (b *multiBackend) close() error {
+	b.lock.Lock()
+	defer b.lock.Unlock()
+
+	err := b.lockfile.Close()
+	if err == nil || errors.Is(err, os.ErrInvalid) || errors.Is(err, os.ErrClosed) {
+		return nil
+	}
+	return err
+}
+
+// NewMulti returns an instance of the multi-file store.
+func NewMulti(runDir string) Store {
 	b := new(multiStore)
-	b.path = append([]string{runDir, "state"}, prefix...)
+	b.base = path.Join(runDir, "state")
+	b.backends = new(sync.Map)
 	return b
 }