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 }