store.go

package main

import (
	"bytes"
	"crypto/sha1"
	"encoding/hex"
	"fmt"
	"io"
	"log"
	"os"
	"strings"
	"sync"
)

type PathTransformation func(string, string) PathKey

func DefaultPathTransformation(root string, key string) PathKey {
	return PathKey{PathName: key, FileName: key, dirs: []string{root + "/" + key}}
}
func SHA1PathTransformation(root string, key string) PathKey {
	hash := sha1.Sum([]byte(key))
	hashstr := hex.EncodeToString(hash[:])

	blocksize := 5
	hashstrlen := len(hashstr) / blocksize
	dirs := make([]string, 0)
	dirs = append(dirs, root)

	for i := 0; i < hashstrlen; i++ {
		from, to := i*blocksize, (i*blocksize)+blocksize
		dirs = append(dirs, hashstr[from:to])
	}

	return PathKey{
		PathName: strings.Join(dirs, "/"),
		FileName: hashstr,
		dirs:     dirs,
	}
}

type PathKey struct {
	PathName string
	FileName string
	dirs     []string
}

func (p *PathKey) FullPath() string {
	return fmt.Sprintf("%s/%s", p.PathName, p.FileName)
}

func (p *PathKey) FirstMultiFilesDir() (string, error) {
	paths := p.dirs

	for {
		dirpath := strings.Join(paths, "/")

		files, err := os.ReadDir(dirpath)
		if err != nil {
			return "", err
		}

		if len(files) > 1 {
			return dirpath, nil
		}

		if len(paths) == 1 {
			break
		}
		paths = paths[:len(paths)-1]
	}

	return paths[0], nil
}

const defaultRootFolderName = "root"

const ChunkSize = 4065

type Chunk struct {
	Key           string
	FileSize      int64
	Offset        int64
	Size          int64
	Data          []byte
	IsLast        bool
	Checksum      uint32
	FinalChecksum uint32
}

type StoreConfig struct {
	Root               string
	PathTransformation PathTransformation
}

type Store struct {
	config StoreConfig
	mutex  sync.Mutex
}

func NewStore(config StoreConfig) *Store {
	if config.PathTransformation == nil {
		config.PathTransformation = DefaultPathTransformation
	}

	if len(config.Root) == 0 {
		config.Root = defaultRootFolderName
	}

	return &Store{
		config: config,
	}
}

func (s *Store) readFile(key string) (io.ReadCloser, error) {
	pathkey := s.config.PathTransformation(s.config.Root, key)
	return os.Open(pathkey.FullPath())
}

func (s *Store) ReadChunk(key string, offset int64, size int) ([]byte, error) {
	pathkey := s.config.PathTransformation(s.config.Root, key)
	fullpath := pathkey.FullPath()

	s.mutex.Lock()
	defer s.mutex.Unlock()

	f, err := os.Open(fullpath)
	if err != nil {
		return nil, err
	}
	defer f.Close()

	_, err = f.Seek(offset, io.SeekStart)
	if err != nil {
		return nil, err
	}

	buff := make([]byte, size)

	n, err := f.Read(buff)
	if err != nil && err != io.EOF {
		return nil, err
	}

	return buff[:n], nil
}

func (s *Store) Read(key string) (io.Reader, error) {
	f, err := s.readFile(key)
	if err != nil {
		return nil, err
	}
	defer f.Close()

	buff := new(bytes.Buffer)
	_, err = io.Copy(buff, f)

	return buff, err
}

func (s *Store) WriteChunk(chunk Chunk) error {
	pathkey := s.config.PathTransformation(s.config.Root, chunk.Key)

	if err := os.MkdirAll(pathkey.PathName, os.ModePerm); err != nil {
		return err
	}

	fullpath := pathkey.FullPath()

	s.mutex.Lock()
	defer s.mutex.Unlock()

	f, err := os.OpenFile(fullpath, os.O_CREATE|os.O_WRONLY, 0666)
	if err != nil {
		return err
	}
	defer f.Close()

	_, err = f.Seek(chunk.Offset, io.SeekStart)
	if err != nil {
		return err
	}

	n, err := f.Write(chunk.Data)
	if err != nil {
		return err
	}

	if int64(n) != chunk.Size {
		return fmt.Errorf("incomplete write, wrote %d bytes, expected %d", n, chunk.Size)
	}

	return nil
}

func (s *Store) Write(key string, r io.Reader) error {
	pathkey := s.config.PathTransformation(s.config.Root, key)

	if err := os.MkdirAll(pathkey.PathName, os.ModePerm); err != nil {
		return err
	}

	fullpath := pathkey.FullPath()

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

	n, err := io.Copy(f, r)
	if err != nil {
		return err
	}

	log.Printf("Written (%d) bytes to disk: %s", n, fullpath)
	return nil
}

func (s *Store) Delete(key string) error {
	pathkey := s.config.PathTransformation(s.config.Root, key)

	path, err := pathkey.FirstMultiFilesDir()

	if err != nil || path == pathkey.PathName {
		return os.RemoveAll(pathkey.FullPath())
	}

	pathdirs := strings.Split(path, "/")
	pathdirs = append(pathdirs, pathkey.dirs[len(pathdirs)])
	path = strings.Join(pathdirs, "/")

	return os.RemoveAll(path)
}

func (s *Store) Clear() error {
	return os.RemoveAll(s.config.Root)
}

func (s *Store) Has(key string) (os.FileInfo, bool) {
	pathkey := s.config.PathTransformation(s.config.Root, key)

	info, err := os.Stat(pathkey.FullPath())
	if err != nil {
		return nil, false
	}

	return info, true
}