local-server/cmd/dkl-local-server/secret-store.go

package main

import (
	"encoding/json"
	"log"
	"net/http"
	"os"
	"path/filepath"
	"sort"
	"strings"
	"sync"

	restful "github.com/emicklei/go-restful"
	"m.cluseau.fr/go/httperr"

	"novit.tech/direktil/local-server/secretstore"
)

var secStore *secretstore.Store

func secStoreRoot() string            { return filepath.Join(*dataDir, "secrets") }
func secStorePath(name string) string { return filepath.Join(secStoreRoot(), name) }
func secKeysStorePath() string        { return secStorePath(".keys") }

func openSecretStore() {
	var err error

	keysPath := secKeysStorePath()

	if err := os.MkdirAll(filepath.Dir(filepath.Dir(keysPath)), 0755); err != nil {
		log.Fatal("failed to create dirs: ", err)
	}

	if err := os.MkdirAll(filepath.Dir(keysPath), 0700); err != nil {
		log.Fatal("failed to secret store dir: ", err)
	}

	secStore, err = secretstore.Open(keysPath)

	switch {
	case err == nil:
		wPublicState.Change(func(v *PublicState) {
			v.Store.New = false
			v.Store.Open = false
		})

	case os.IsNotExist(err):
		secStore = secretstore.New()
		wPublicState.Change(func(v *PublicState) {
			v.Store.New = true
			v.Store.Open = false
		})

	default:
		log.Fatal("failed to open keys store: ", err)
	}
}

var (
	unlockMutex = sync.Mutex{}

	ErrStoreAlreadyUnlocked = httperr.NewStd(1, http.StatusConflict, "store already unlocked")
	ErrInvalidPassphrase    = httperr.NewStd(2, http.StatusBadRequest, "invalid passphrase")
)

func unlockSecretStore(name string, passphrase []byte) (err httperr.Error) {
	unlockMutex.Lock()
	defer unlockMutex.Unlock()

	if secStore.Unlocked() {
		return ErrStoreAlreadyUnlocked
	}

	if secStore.IsNew() {
		err := secStore.Init(name, passphrase)
		if err != nil {
			return httperr.Internal(err)
		}

		err = secStore.SaveTo(secKeysStorePath())
		if err != nil {
			log.Print("secret store save error: ", err)
			secStore.Close()

			return httperr.Internal(err)
		}

	} else {
		if !secStore.Unlock([]byte(passphrase)) {
			return ErrInvalidPassphrase
		}
	}

	token := ""
	if err := readSecret("admin-token", &token); err != nil {
		if !os.IsNotExist(err) {
			log.Print("failed to read admin token: ", err)
			secStore.Close()

			return httperr.Internal(err)
		}

		token, err = newToken(32)
		if err != nil {
			secStore.Close()
			return httperr.Internal(err)
		}

		err = writeSecret("admin-token", token)
		if err != nil {
			log.Print("write error: ", err)
			secStore.Close()

			return httperr.Internal(err)
		}

		log.Print("wrote new admin token")
	}

	adminToken = token

	{
		token, err := newToken(16)
		if err != nil {
			secStore.Close()
			return httperr.Internal(err)
		}

		wState.Change(func(v *State) {
			v.Store.DownloadToken = token
		})
	}

	wPublicState.Change(func(v *PublicState) {
		v.Store.New = false
		v.Store.Open = true
	})

	go updateState()
	go migrateSecrets()

	return
}

func readSecret(name string, value any) (err error) {
	f, err := os.Open(secStorePath(name + ".data"))
	if err != nil {
		return
	}

	defer f.Close()

	in, err := secStore.NewReader(f)
	if err != nil {
		return
	}

	return json.NewDecoder(in).Decode(value)
}

func writeSecret(name string, value any) (err error) {
	path := secStorePath(name + ".data.new")

	if err = os.MkdirAll(filepath.Dir(path), 0700); err != nil {
		return
	}

	f, err := os.Create(path)
	if err != nil {
		return
	}

	err = func() (err error) {
		defer f.Close()

		out, err := secStore.NewWriter(f)
		if err != nil {
			return
		}

		return json.NewEncoder(out).Encode(value)
	}()

	if err != nil {
		return
	}

	err = os.Rename(f.Name(), secStorePath(name+".data"))
	if err != nil {
		return
	}

	go updateState()
	return
}

var secL sync.Mutex

func updateSecret[T any](name string, update func(*T)) (err error) {
	secL.Lock()
	defer secL.Unlock()

	v := new(T)
	err = readSecret(name, v)
	if err != nil {
		if !os.IsNotExist(err) {
			return
		}
		err = nil
	}

	update(v)

	return writeSecret(name, *v)
}

func updateSecretWithKey[T any](name, key string, update func(v *T)) (err error) {
	secL.Lock()
	defer secL.Unlock()

	kvs := map[string]*T{}

	err = readSecret(name, &kvs)
	if err != nil {
		if !os.IsNotExist(err) {
			return
		}
		err = nil
	}

	update(kvs[key])

	return writeSecret(name, kvs)
}

type KVSecrets[T any] struct{ Name string }

func (s KVSecrets[T]) Data() (kvs map[string]T, err error) {
	kvs = make(map[string]T)
	err = readSecret(s.Name, &kvs)
	if err != nil {
		if !os.IsNotExist(err) {
			return
		}
		err = nil
	}
	return
}

type KV[T any] struct {
	K string
	V T
}

func (s KVSecrets[T]) List(prefix string) (list []KV[T], err error) {
	data, err := s.Data()
	if err != nil {
		return
	}

	list = make([]KV[T], 0, len(data))

	for k, v := range data {
		if !strings.HasPrefix(k, prefix) {
			continue
		}
		list = append(list, KV[T]{k, v})
	}

	sort.Slice(list, func(i, j int) bool { return list[i].K < list[j].K })

	return
}

func (s KVSecrets[T]) Keys(prefix string) (keys []string, err error) {
	kvs, err := s.Data()
	if err != nil {
		return
	}

	keys = make([]string, 0, len(kvs))

	for k := range kvs {
		if !strings.HasPrefix(k, prefix) {
			continue
		}
		keys = append(keys, k[len(prefix):])
	}

	sort.Strings(keys)

	return
}

func (s KVSecrets[T]) Get(key string) (v T, found bool, err error) {
	kvs, err := s.Data()
	if err != nil {
		return
	}

	v, found = kvs[key]
	return
}

func (s KVSecrets[T]) Put(key string, v T) (err error) {
	secL.Lock()
	defer secL.Unlock()

	kvs, err := s.Data()
	if err != nil {
		return
	}

	kvs[key] = v
	err = writeSecret(s.Name, kvs)
	return
}

func (s KVSecrets[T]) Del(key string) (err error) {
	secL.Lock()
	defer secL.Unlock()

	kvs, err := s.Data()
	if err != nil {
		return
	}

	delete(kvs, key)
	err = writeSecret(s.Name, kvs)
	return
}

func (s KVSecrets[T]) GetOrCreate(key string, create func() (T, error)) (v T, err error) {
	v, found, err := s.Get(key)
	if err != nil {
		return
	}

	if !found {
		v, err = create()
		if err != nil {
			return
		}

		err = s.Put(key, v)
	}

	return
}

func (s KVSecrets[T]) WsList(resp *restful.Response, prefix string) {
	keys, err := s.Keys(prefix)
	if err != nil {
		wsError(resp, err)
		return
	}

	resp.WriteEntity(keys)
}

func (s KVSecrets[T]) WsGet(resp *restful.Response, key string) {
	keys, found, err := s.Get(key)
	if err != nil {
		wsError(resp, err)
		return
	}

	if !found {
		wsNotFound(resp)
		return
	}

	resp.WriteEntity(keys)
}

func (s KVSecrets[T]) WsPut(req *restful.Request, resp *restful.Response, key string) {
	v := new(T)
	err := req.ReadEntity(v)
	if err != nil {
		wsBadRequest(resp, err.Error())
		return
	}

	err = s.Put(key, *v)
	if err != nil {
		wsError(resp, err)
		return
	}
}

func (s KVSecrets[T]) WsDel(req *restful.Request, resp *restful.Response, key string) {
	err := s.Del(key)
	if err != nil {
		wsError(resp, err)
		return
	}
}