package store
import (
"fmt"
"os"
"path/filepath"
"sync"
"time"
"gopkg.in/yaml.v3"
)
type persistedData struct {
Server ServerState `yaml:"server"`
Keys []AccessKey `yaml:"access_keys"`
}
type YAMLFileStore struct {
mu sync.RWMutex
filePath string
data persistedData
}
func NewYAMLFileStore(filePath string, serverID, serverName, hostname, defaultCipher string, defaultPort int) (*YAMLFileStore, error) {
s := &YAMLFileStore{filePath: filePath}
if err := os.MkdirAll(filepath.Dir(filePath), 0700); err != nil {
return nil, fmt.Errorf("creating state directory: %w", err)
}
data, err := os.ReadFile(filePath)
if err != nil {
if !os.IsNotExist(err) {
return nil, fmt.Errorf("reading state file: %w", err)
}
// Initialize with defaults.
s.data = persistedData{
Server: ServerState{
ID: serverID,
Name: serverName,
Hostname: hostname,
DefaultPort: defaultPort,
DefaultCipher: defaultCipher,
PortForNewAccessKeys: defaultPort,
CreatedTimestampMs: time.Now().UnixMilli(),
},
}
return s, s.persist()
}
if err := yaml.Unmarshal(data, &s.data); err != nil {
return nil, fmt.Errorf("parsing state file: %w", err)
}
return s, nil
}
func (s *YAMLFileStore) ListKeys() []AccessKey {
s.mu.RLock()
defer s.mu.RUnlock()
out := make([]AccessKey, len(s.data.Keys))
copy(out, s.data.Keys)
return out
}
func (s *YAMLFileStore) GetKey(id string) (AccessKey, bool) {
s.mu.RLock()
defer s.mu.RUnlock()
for _, k := range s.data.Keys {
if k.ID == id {
return k, true
}
}
return AccessKey{}, false
}
func (s *YAMLFileStore) CreateKey(ak AccessKey) error {
s.mu.Lock()
defer s.mu.Unlock()
for _, k := range s.data.Keys {
if k.ID == ak.ID {
return fmt.Errorf("key %s already exists", ak.ID)
}
}
s.data.Keys = append(s.data.Keys, ak)
return s.persist()
}
func (s *YAMLFileStore) UpdateKey(id string, fn func(*AccessKey)) error {
s.mu.Lock()
defer s.mu.Unlock()
for i := range s.data.Keys {
if s.data.Keys[i].ID == id {
fn(&s.data.Keys[i])
return s.persist()
}
}
return fmt.Errorf("key %s not found", id)
}
func (s *YAMLFileStore) DeleteKey(id string) error {
s.mu.Lock()
defer s.mu.Unlock()
for i, k := range s.data.Keys {
if k.ID == id {
s.data.Keys = append(s.data.Keys[:i], s.data.Keys[i+1:]...)
return s.persist()
}
}
return fmt.Errorf("key %s not found", id)
}
func (s *YAMLFileStore) GetServer() ServerState {
s.mu.RLock()
defer s.mu.RUnlock()
return s.data.Server
}
func (s *YAMLFileStore) UpdateServer(fn func(*ServerState)) error {
s.mu.Lock()
defer s.mu.Unlock()
fn(&s.data.Server)
return s.persist()
}
func (s *YAMLFileStore) persist() error {
data, err := yaml.Marshal(s.data)
if err != nil {
return fmt.Errorf("marshaling state: %w", err)
}
tmpPath := s.filePath + ".tmp"
if err := os.WriteFile(tmpPath, data, 0600); err != nil {
return fmt.Errorf("writing temp state file: %w", err)
}
if err := os.Rename(tmpPath, s.filePath); err != nil {
return fmt.Errorf("renaming state file: %w", err)
}
return nil
}