126 lines
2.7 KiB
Go
126 lines
2.7 KiB
Go
package storage
|
|
|
|
import (
|
|
"fmt"
|
|
"io"
|
|
"mime/multipart"
|
|
"net/http"
|
|
"os"
|
|
"path/filepath"
|
|
"strings"
|
|
)
|
|
|
|
type Client struct {
|
|
rootDir string
|
|
maxObjectSize int64
|
|
}
|
|
|
|
func New(root string, maxObjectSize int64) (*Client, error) {
|
|
if maxObjectSize <= 0 {
|
|
return nil, fmt.Errorf("maxObjectSize must be greater than 0")
|
|
}
|
|
|
|
info, err := os.Stat(root)
|
|
if os.IsNotExist(err) {
|
|
return nil, fmt.Errorf("root directory does not exist: %s", root)
|
|
}
|
|
if err != nil {
|
|
return nil, fmt.Errorf("failed to stat root directory: %v", err)
|
|
}
|
|
if !info.IsDir() {
|
|
return nil, fmt.Errorf("root path is not a directory: %s", root)
|
|
}
|
|
|
|
testFile := filepath.Join(root, ".axiom-storage-check")
|
|
f, err := os.Create(testFile)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("root directory is not writable: %v", err)
|
|
}
|
|
f.Close()
|
|
os.Remove(testFile)
|
|
|
|
return &Client{
|
|
rootDir: root,
|
|
maxObjectSize: maxObjectSize,
|
|
}, nil
|
|
}
|
|
|
|
func (s *Client) getPath(bucket, key string) string {
|
|
return filepath.Join(s.rootDir, bucket, key)
|
|
}
|
|
|
|
func (s *Client) CreateBucket(bucket string) error {
|
|
path := filepath.Join(s.rootDir, bucket)
|
|
return os.MkdirAll(path, 0755)
|
|
}
|
|
|
|
func (s *Client) PutObject(bucket, key string, reader io.Reader) error {
|
|
fullPath := s.getPath(bucket, key)
|
|
|
|
if err := os.MkdirAll(filepath.Dir(fullPath), 0755); err != nil {
|
|
return err
|
|
}
|
|
|
|
f, err := os.Create(fullPath)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
defer f.Close()
|
|
|
|
limitedReader := io.LimitReader(reader, s.maxObjectSize+1)
|
|
written, err := io.Copy(f, limitedReader)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
if written > s.maxObjectSize {
|
|
os.Remove(fullPath)
|
|
return fmt.Errorf("object size exceeds maximum allowed: %d bytes", s.maxObjectSize)
|
|
}
|
|
|
|
return nil
|
|
}
|
|
|
|
func (s *Client) GetObject(bucket, key string) (*os.File, error) {
|
|
fullPath := s.getPath(bucket, key)
|
|
return os.Open(fullPath)
|
|
}
|
|
|
|
func (s *Client) DeleteObject(bucket, key string) error {
|
|
fullPath := s.getPath(bucket, key)
|
|
return os.Remove(fullPath)
|
|
}
|
|
|
|
func DetectFileType(file multipart.File) (string, []byte, error) {
|
|
const sniffLen = 512
|
|
buffer := make([]byte, sniffLen)
|
|
|
|
n, err := file.Read(buffer)
|
|
if err != nil && err != io.EOF {
|
|
return "", nil, err
|
|
}
|
|
|
|
if seeker, ok := file.(io.Seeker); ok {
|
|
_, _ = seeker.Seek(0, io.SeekStart)
|
|
} else {
|
|
return "", nil, fmt.Errorf("file is not seekable")
|
|
}
|
|
|
|
mimeType := http.DetectContentType(buffer[:n])
|
|
return mimeType, buffer[:n], nil
|
|
}
|
|
|
|
func VerifyFileType(file multipart.File, allowedMimeTypes []string) error {
|
|
mimeType, _, err := DetectFileType(file)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
for _, mt := range allowedMimeTypes {
|
|
if strings.EqualFold(mimeType, mt) {
|
|
return nil
|
|
}
|
|
}
|
|
|
|
return fmt.Errorf("invalid file type: %s", mimeType)
|
|
}
|