warpbox-dev/backend/libs/services/storage_smb.go

package services

import (
	"context"
	"io"
	"net"
	"os"
	"path"
	"strconv"
	"strings"
	"time"

	"github.com/hirochachacha/go-smb2"
)

type smbStorageBackend struct {
	cfg StorageBackendConfig
}

func (b smbStorageBackend) ID() string   { return b.cfg.ID }
func (b smbStorageBackend) Type() string { return StorageBackendSMB }

func (b smbStorageBackend) Put(ctx context.Context, key string, body io.Reader, _ int64, _ string) error {
	share, closer, err := b.share()
	if err != nil {
		return err
	}
	defer closer()
	if err := ctx.Err(); err != nil {
		return err
	}
	remotePath := b.remotePath(key)
	if err := share.MkdirAll(path.Dir(remotePath), 0o755); err != nil {
		return err
	}
	target, err := share.OpenFile(remotePath, os.O_CREATE|os.O_TRUNC|os.O_WRONLY, 0o644)
	if err != nil {
		return err
	}
	defer target.Close()
	_, err = io.Copy(target, body)
	return err
}

func (b smbStorageBackend) Get(ctx context.Context, key string) (StorageObject, error) {
	share, closer, err := b.share()
	if err != nil {
		return StorageObject{}, err
	}
	if err := ctx.Err(); err != nil {
		closer()
		return StorageObject{}, err
	}
	source, err := share.Open(b.remotePath(key))
	if err != nil {
		closer()
		return StorageObject{}, err
	}
	stat, err := source.Stat()
	if err != nil {
		source.Close()
		closer()
		return StorageObject{}, err
	}
	return StorageObject{Key: key, Size: stat.Size(), ModTime: stat.ModTime(), Body: closeWith(source, closer)}, nil
}

func (b smbStorageBackend) Delete(ctx context.Context, key string) error {
	share, closer, err := b.share()
	if err != nil {
		return err
	}
	defer closer()
	if err := ctx.Err(); err != nil {
		return err
	}
	if err := share.Remove(b.remotePath(key)); err != nil && !os.IsNotExist(err) {
		return err
	}
	return nil
}

func (b smbStorageBackend) DeletePrefix(ctx context.Context, prefix string) error {
	share, closer, err := b.share()
	if err != nil {
		return err
	}
	defer closer()
	if err := ctx.Err(); err != nil {
		return err
	}
	err = share.RemoveAll(b.remotePath(prefix))
	if err != nil && !os.IsNotExist(err) {
		return err
	}
	return nil
}

func (b smbStorageBackend) Usage(ctx context.Context) (int64, error) {
	share, closer, err := b.share()
	if err != nil {
		return 0, err
	}
	defer closer()
	if err := ctx.Err(); err != nil {
		return 0, err
	}
	return smbUsage(share, cleanRemoteRoot(b.cfg.RemotePath))
}

func (b smbStorageBackend) Test(ctx context.Context) error {
	key := ".warpbox-storage-test-" + randomID(6)
	if err := b.Put(ctx, key, strings.NewReader("ok"), 2, "text/plain"); err != nil {
		return err
	}
	return b.Delete(ctx, key)
}

func (b smbStorageBackend) share() (*smb2.Share, func(), error) {
	conn, err := net.DialTimeout("tcp", b.cfg.Host+":"+strconv.Itoa(b.cfg.Port), 15*time.Second)
	if err != nil {
		return nil, nil, err
	}
	dialer := &smb2.Dialer{
		Initiator: &smb2.NTLMInitiator{
			User:     b.cfg.Username,
			Password: b.cfg.Password,
			Domain:   b.cfg.Domain,
		},
	}
	session, err := dialer.Dial(conn)
	if err != nil {
		conn.Close()
		return nil, nil, err
	}
	share, err := session.Mount(b.cfg.Share)
	if err != nil {
		session.Logoff()
		conn.Close()
		return nil, nil, err
	}
	return share, func() {
		share.Umount()
		session.Logoff()
		conn.Close()
	}, nil
}

func (b smbStorageBackend) remotePath(key string) string {
	return strings.TrimPrefix(path.Join(cleanRemoteRoot(b.cfg.RemotePath), cleanObjectKey(key)), "/")
}

func smbUsage(share *smb2.Share, root string) (int64, error) {
	root = strings.TrimPrefix(root, "/")
	entries, err := share.ReadDir(root)
	if err != nil {
		if os.IsNotExist(err) {
			return 0, nil
		}
		return 0, err
	}
	var total int64
	for _, entry := range entries {
		item := path.Join(root, entry.Name())
		if entry.IsDir() {
			size, err := smbUsage(share, item)
			if err != nil {
				return 0, err
			}
			total += size
			continue
		}
		total += entry.Size()
	}
	return total, nil
}