mirror of
https://github.com/Dvorinka/Primora.git
synced 2026-06-03 20:13:01 +00:00
initiall commit
This commit is contained in:
Tomas Dvorak
@@ -0,0 +1,159 @@
|
||||
package storage
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/sha256"
|
||||
"encoding/hex"
|
||||
"fmt"
|
||||
"io"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
)
|
||||
|
||||
type PutResult struct {
|
||||
Path string
|
||||
SizeBytes int64
|
||||
SHA256Digest string
|
||||
}
|
||||
|
||||
type LocalStore struct {
|
||||
root string
|
||||
}
|
||||
|
||||
func NewLocalStore(root string) (*LocalStore, error) {
|
||||
if err := os.MkdirAll(root, 0o755); err != nil {
|
||||
return nil, fmt.Errorf("create storage root: %w", err)
|
||||
}
|
||||
return &LocalStore{root: root}, nil
|
||||
}
|
||||
|
||||
func (s *LocalStore) Put(ctx context.Context, bucketID, objectKey string, reader io.Reader) (PutResult, error) {
|
||||
cleanKey, err := sanitizeObjectKey(objectKey)
|
||||
if err != nil {
|
||||
return PutResult{}, err
|
||||
}
|
||||
dir := filepath.Join(s.root, bucketID)
|
||||
if err := os.MkdirAll(filepath.Dir(filepath.Join(dir, cleanKey)), 0o755); err != nil {
|
||||
return PutResult{}, fmt.Errorf("create object directory: %w", err)
|
||||
}
|
||||
path := filepath.Join(dir, cleanKey)
|
||||
tmpPath := path + ".tmp"
|
||||
file, err := os.Create(tmpPath)
|
||||
if err != nil {
|
||||
return PutResult{}, fmt.Errorf("create temp object: %w", err)
|
||||
}
|
||||
defer file.Close()
|
||||
|
||||
hasher := sha256.New()
|
||||
writer := io.MultiWriter(file, hasher)
|
||||
written, err := copyWithContext(ctx, writer, reader)
|
||||
if err != nil {
|
||||
_ = os.Remove(tmpPath)
|
||||
return PutResult{}, err
|
||||
}
|
||||
if err := file.Close(); err != nil {
|
||||
return PutResult{}, fmt.Errorf("close temp object: %w", err)
|
||||
}
|
||||
if err := os.Rename(tmpPath, path); err != nil {
|
||||
return PutResult{}, fmt.Errorf("rename object: %w", err)
|
||||
}
|
||||
return PutResult{
|
||||
Path: path,
|
||||
SizeBytes: written,
|
||||
SHA256Digest: hex.EncodeToString(hasher.Sum(nil)),
|
||||
}, nil
|
||||
}
|
||||
|
||||
func (s *LocalStore) Open(bucketID, objectKey string) (*os.File, string, error) {
|
||||
cleanKey, err := sanitizeObjectKey(objectKey)
|
||||
if err != nil {
|
||||
return nil, "", err
|
||||
}
|
||||
path := filepath.Join(s.root, bucketID, cleanKey)
|
||||
file, err := os.Open(path)
|
||||
if err != nil {
|
||||
return nil, "", fmt.Errorf("open object: %w", err)
|
||||
}
|
||||
return file, path, nil
|
||||
}
|
||||
|
||||
func (s *LocalStore) Delete(bucketID, objectKey string) error {
|
||||
cleanKey, err := sanitizeObjectKey(objectKey)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if err := os.Remove(filepath.Join(s.root, bucketID, cleanKey)); err != nil && !os.IsNotExist(err) {
|
||||
return fmt.Errorf("delete object: %w", err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *LocalStore) Move(bucketID, fromKey, toKey string) (string, error) {
|
||||
return s.MoveBetweenBuckets(bucketID, bucketID, fromKey, toKey)
|
||||
}
|
||||
|
||||
func (s *LocalStore) MoveBetweenBuckets(sourceBucketID, destinationBucketID, fromKey, toKey string) (string, error) {
|
||||
cleanFrom, err := sanitizeObjectKey(fromKey)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
cleanTo, err := sanitizeObjectKey(toKey)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
fromPath := filepath.Join(s.root, strings.TrimSpace(sourceBucketID), cleanFrom)
|
||||
toPath := filepath.Join(s.root, strings.TrimSpace(destinationBucketID), cleanTo)
|
||||
if err := os.MkdirAll(filepath.Dir(toPath), 0o755); err != nil {
|
||||
return "", fmt.Errorf("create destination directory: %w", err)
|
||||
}
|
||||
if err := os.Rename(fromPath, toPath); err != nil {
|
||||
return "", fmt.Errorf("move object: %w", err)
|
||||
}
|
||||
return toPath, nil
|
||||
}
|
||||
|
||||
func (s *LocalStore) DeleteBucket(bucketID string) error {
|
||||
bucketPath := filepath.Join(s.root, strings.TrimSpace(bucketID))
|
||||
if err := os.RemoveAll(bucketPath); err != nil {
|
||||
return fmt.Errorf("delete bucket directory: %w", err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func sanitizeObjectKey(key string) (string, error) {
|
||||
clean := filepath.Clean(strings.TrimSpace(key))
|
||||
if clean == "." || clean == "" || strings.HasPrefix(clean, "../") || strings.Contains(clean, "/../") || strings.HasPrefix(clean, "/") {
|
||||
return "", fmt.Errorf("invalid object key")
|
||||
}
|
||||
return clean, nil
|
||||
}
|
||||
|
||||
func copyWithContext(ctx context.Context, dst io.Writer, src io.Reader) (int64, error) {
|
||||
buffer := make([]byte, 32*1024)
|
||||
var written int64
|
||||
for {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return written, ctx.Err()
|
||||
default:
|
||||
}
|
||||
nr, er := src.Read(buffer)
|
||||
if nr > 0 {
|
||||
nw, ew := dst.Write(buffer[:nr])
|
||||
written += int64(nw)
|
||||
if ew != nil {
|
||||
return written, ew
|
||||
}
|
||||
if nr != nw {
|
||||
return written, io.ErrShortWrite
|
||||
}
|
||||
}
|
||||
if er != nil {
|
||||
if er == io.EOF {
|
||||
return written, nil
|
||||
}
|
||||
return written, er
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,181 @@
|
||||
package storage
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"crypto/sha256"
|
||||
"encoding/hex"
|
||||
"io"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestSanitizeObjectKey(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
valid, err := sanitizeObjectKey(" docs/readme.txt ")
|
||||
if err != nil {
|
||||
t.Fatalf("expected valid key, got error: %v", err)
|
||||
}
|
||||
if valid != "docs/readme.txt" {
|
||||
t.Fatalf("unexpected sanitized key: %s", valid)
|
||||
}
|
||||
|
||||
invalidKeys := []string{"", ".", "../secret", "/absolute/path", " /../../etc "}
|
||||
for _, key := range invalidKeys {
|
||||
if _, err := sanitizeObjectKey(key); err == nil {
|
||||
t.Fatalf("expected key %q to be invalid", key)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestLocalStorePutOpenDeleteRoundTrip(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
root := t.TempDir()
|
||||
store, err := NewLocalStore(root)
|
||||
if err != nil {
|
||||
t.Fatalf("new local store: %v", err)
|
||||
}
|
||||
|
||||
content := []byte("hello primora")
|
||||
put, err := store.Put(context.Background(), "bucket-a", "docs/hello.txt", bytes.NewReader(content))
|
||||
if err != nil {
|
||||
t.Fatalf("put object: %v", err)
|
||||
}
|
||||
if put.SizeBytes != int64(len(content)) {
|
||||
t.Fatalf("unexpected size: %d", put.SizeBytes)
|
||||
}
|
||||
expectedDigest := sha256.Sum256(content)
|
||||
if put.SHA256Digest != hex.EncodeToString(expectedDigest[:]) {
|
||||
t.Fatalf("unexpected digest: %s", put.SHA256Digest)
|
||||
}
|
||||
|
||||
file, path, err := store.Open("bucket-a", "docs/hello.txt")
|
||||
if err != nil {
|
||||
t.Fatalf("open object: %v", err)
|
||||
}
|
||||
defer file.Close()
|
||||
if filepath.Clean(path) != filepath.Clean(put.Path) {
|
||||
t.Fatalf("path mismatch: got %s, want %s", path, put.Path)
|
||||
}
|
||||
data, err := io.ReadAll(file)
|
||||
if err != nil {
|
||||
t.Fatalf("read object: %v", err)
|
||||
}
|
||||
if !bytes.Equal(data, content) {
|
||||
t.Fatalf("content mismatch: got %q, want %q", string(data), string(content))
|
||||
}
|
||||
|
||||
if err := store.Delete("bucket-a", "docs/hello.txt"); err != nil {
|
||||
t.Fatalf("delete object: %v", err)
|
||||
}
|
||||
if _, _, err := store.Open("bucket-a", "docs/hello.txt"); err == nil {
|
||||
t.Fatalf("expected open to fail after delete")
|
||||
}
|
||||
}
|
||||
|
||||
func TestLocalStoreDeleteBucketRemovesAllFiles(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
root := t.TempDir()
|
||||
store, err := NewLocalStore(root)
|
||||
if err != nil {
|
||||
t.Fatalf("new local store: %v", err)
|
||||
}
|
||||
|
||||
if _, err := store.Put(context.Background(), "bucket-z", "a.txt", bytes.NewReader([]byte("a"))); err != nil {
|
||||
t.Fatalf("put a.txt: %v", err)
|
||||
}
|
||||
if _, err := store.Put(context.Background(), "bucket-z", "nested/b.txt", bytes.NewReader([]byte("b"))); err != nil {
|
||||
t.Fatalf("put b.txt: %v", err)
|
||||
}
|
||||
|
||||
if err := store.DeleteBucket("bucket-z"); err != nil {
|
||||
t.Fatalf("delete bucket: %v", err)
|
||||
}
|
||||
|
||||
_, err = os.Stat(filepath.Join(root, "bucket-z"))
|
||||
if !os.IsNotExist(err) {
|
||||
t.Fatalf("expected bucket path to be removed, stat err: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestLocalStoreMoveObject(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
root := t.TempDir()
|
||||
store, err := NewLocalStore(root)
|
||||
if err != nil {
|
||||
t.Fatalf("new local store: %v", err)
|
||||
}
|
||||
|
||||
content := []byte("rename me")
|
||||
if _, err := store.Put(context.Background(), "bucket-r", "source/name.txt", bytes.NewReader(content)); err != nil {
|
||||
t.Fatalf("put source object: %v", err)
|
||||
}
|
||||
|
||||
newPath, err := store.Move("bucket-r", "source/name.txt", "dest/renamed.txt")
|
||||
if err != nil {
|
||||
t.Fatalf("move object: %v", err)
|
||||
}
|
||||
if filepath.Clean(newPath) != filepath.Clean(filepath.Join(root, "bucket-r", "dest/renamed.txt")) {
|
||||
t.Fatalf("unexpected moved path: %s", newPath)
|
||||
}
|
||||
|
||||
if _, _, err := store.Open("bucket-r", "source/name.txt"); err == nil {
|
||||
t.Fatalf("expected old key open to fail after move")
|
||||
}
|
||||
file, _, err := store.Open("bucket-r", "dest/renamed.txt")
|
||||
if err != nil {
|
||||
t.Fatalf("open moved object: %v", err)
|
||||
}
|
||||
defer file.Close()
|
||||
data, err := io.ReadAll(file)
|
||||
if err != nil {
|
||||
t.Fatalf("read moved object: %v", err)
|
||||
}
|
||||
if !bytes.Equal(data, content) {
|
||||
t.Fatalf("moved content mismatch: got %q, want %q", string(data), string(content))
|
||||
}
|
||||
}
|
||||
|
||||
func TestLocalStoreMoveObjectAcrossBuckets(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
root := t.TempDir()
|
||||
store, err := NewLocalStore(root)
|
||||
if err != nil {
|
||||
t.Fatalf("new local store: %v", err)
|
||||
}
|
||||
|
||||
content := []byte("cross bucket")
|
||||
if _, err := store.Put(context.Background(), "bucket-src", "folder/source.txt", bytes.NewReader(content)); err != nil {
|
||||
t.Fatalf("put source object: %v", err)
|
||||
}
|
||||
|
||||
newPath, err := store.MoveBetweenBuckets("bucket-src", "bucket-dst", "folder/source.txt", "archive/destination.txt")
|
||||
if err != nil {
|
||||
t.Fatalf("move across buckets: %v", err)
|
||||
}
|
||||
if filepath.Clean(newPath) != filepath.Clean(filepath.Join(root, "bucket-dst", "archive/destination.txt")) {
|
||||
t.Fatalf("unexpected moved path: %s", newPath)
|
||||
}
|
||||
|
||||
if _, _, err := store.Open("bucket-src", "folder/source.txt"); err == nil {
|
||||
t.Fatalf("expected source key open to fail after move")
|
||||
}
|
||||
file, _, err := store.Open("bucket-dst", "archive/destination.txt")
|
||||
if err != nil {
|
||||
t.Fatalf("open moved object: %v", err)
|
||||
}
|
||||
defer file.Close()
|
||||
data, err := io.ReadAll(file)
|
||||
if err != nil {
|
||||
t.Fatalf("read moved object: %v", err)
|
||||
}
|
||||
if !bytes.Equal(data, content) {
|
||||
t.Fatalf("moved content mismatch: got %q, want %q", string(data), string(content))
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user