go-omni/storage/local/local.go

// Package local stores objects atomically in a private, flat managed directory.
// This directory must not be writable by untrusted local users/processes.
package local

import (
	"context"
	"crypto/sha256"
	"encoding/binary"
	"encoding/hex"
	"encoding/json"
	"errors"
	"git.yuntuo.hk/k2safe/go-omni/storage"
	"io"
	"os"
	"path/filepath"
	"strings"
)

type Store struct {
	root     string
	maxBytes int64
}

func New(root string, maxBytes int64) (*Store, error) {
	if root == "" || maxBytes < 0 || maxBytes > 1<<40 {
		return nil, errors.New("storage/local: invalid configuration")
	}
	if maxBytes == 0 {
		maxBytes = 32 << 20
	}
	if err := os.MkdirAll(root, 0700); err != nil {
		return nil, err
	}
	root, err := filepath.EvalSymlinks(root)
	if err != nil {
		return nil, err
	}
	root, err = filepath.Abs(root)
	if err != nil {
		return nil, err
	}
	info, err := os.Stat(root)
	if err != nil {
		return nil, err
	}
	if !info.IsDir() || info.Mode().Perm()&0022 != 0 {
		return nil, errors.New("storage/local: root must be a private writable directory")
	}
	return &Store{root: root, maxBytes: maxBytes}, nil
}
func (s *Store) path(key string) string {
	sum := sha256.Sum256([]byte(key))
	return filepath.Join(s.root, hex.EncodeToString(sum[:])+".object")
}
func (s *Store) Put(ctx context.Context, key string, reader io.Reader, size int64, contentType string) (storage.Info, error) {
	if err := storage.Check(ctx, key); err != nil {
		return storage.Info{}, err
	}
	if reader == nil || size < 0 || len(contentType) > 256 || strings.ContainsAny(contentType, "\r\n") {
		return storage.Info{}, errors.New("storage/local: invalid upload")
	}
	if size > s.maxBytes {
		return storage.Info{}, storage.ErrTooLarge
	}
	info := storage.Info{Key: key, Size: size, ContentType: contentType}
	metadata, _ := json.Marshal(info)
	f, err := os.CreateTemp(s.root, ".upload-")
	if err != nil {
		return storage.Info{}, err
	}
	defer os.Remove(f.Name())
	defer f.Close()
	if err = binary.Write(f, binary.BigEndian, uint32(len(metadata))); err != nil {
		return storage.Info{}, err
	}
	if _, err = f.Write(metadata); err != nil {
		return storage.Info{}, err
	}
	n, err := io.Copy(f, io.LimitReader(storage.ContextReader{Context: ctx, Reader: reader}, size+1))
	if err != nil {
		return storage.Info{}, err
	}
	if n != size {
		return storage.Info{}, errors.New("storage/local: declared size mismatch")
	}
	if err = ctx.Err(); err != nil {
		return storage.Info{}, err
	}
	if err = f.Sync(); err != nil {
		return storage.Info{}, err
	}
	stat, err := f.Stat()
	if err != nil {
		return storage.Info{}, err
	}
	if err = f.Close(); err != nil {
		return storage.Info{}, err
	}
	if err = os.Rename(f.Name(), s.path(key)); err != nil {
		return storage.Info{}, err
	}
	info.ModifiedAt = stat.ModTime()
	return info, nil
}
func (s *Store) Get(ctx context.Context, key string) (storage.Object, error) {
	if err := storage.Check(ctx, key); err != nil {
		return storage.Object{}, err
	}
	path := s.path(key)
	stat, err := os.Lstat(path)
	if os.IsNotExist(err) {
		return storage.Object{}, storage.ErrNotFound
	}
	if err != nil {
		return storage.Object{}, err
	}
	if !stat.Mode().IsRegular() {
		return storage.Object{}, errors.New("storage/local: not a regular managed object")
	}
	f, err := os.Open(path)
	if err != nil {
		return storage.Object{}, err
	}
	// A concurrent Put can atomically replace the path between Lstat/Open.
	// Validate size against the descriptor whose bytes we actually read.
	stat, err = f.Stat()
	if err != nil {
		f.Close()
		return storage.Object{}, err
	}
	ok := false
	defer func() {
		if !ok {
			f.Close()
		}
	}()
	var length uint32
	if err = binary.Read(f, binary.BigEndian, &length); err != nil {
		return storage.Object{}, err
	}
	if length > 4096 {
		return storage.Object{}, errors.New("storage/local: invalid object header")
	}
	data := make([]byte, length)
	if _, err = io.ReadFull(f, data); err != nil {
		return storage.Object{}, err
	}
	var info storage.Info
	if json.Unmarshal(data, &info) != nil || info.Key != key || info.Size < 0 || info.Size > s.maxBytes || int64(length)+4+info.Size != stat.Size() {
		return storage.Object{}, errors.New("storage/local: corrupt object")
	}
	info.ModifiedAt = stat.ModTime()
	ok = true
	return storage.Object{Info: info, Body: &body{Reader: storage.ContextReader{Context: ctx, Reader: io.LimitReader(f, info.Size)}, Closer: f}}, nil
}

type body struct {
	io.Reader
	io.Closer
}

func (s *Store) Delete(ctx context.Context, key string) error {
	if err := storage.Check(ctx, key); err != nil {
		return err
	}
	err := os.Remove(s.path(key))
	if os.IsNotExist(err) {
		return nil
	}
	return err
}