mirror of
https://github.com/krau/SaveAny-Bot.git
synced 2026-07-12 16:02:05 +08:00
feat: refactor webdav client and implement custom upload stream handling
This commit is contained in:
1
go.mod
1
go.mod
@@ -12,7 +12,6 @@ require (
|
|||||||
github.com/rhysd/go-github-selfupdate v1.2.3
|
github.com/rhysd/go-github-selfupdate v1.2.3
|
||||||
github.com/spf13/cobra v1.8.1
|
github.com/spf13/cobra v1.8.1
|
||||||
github.com/spf13/viper v1.19.0
|
github.com/spf13/viper v1.19.0
|
||||||
github.com/studio-b12/gowebdav v0.10.0
|
|
||||||
golang.org/x/net v0.35.0
|
golang.org/x/net v0.35.0
|
||||||
golang.org/x/time v0.10.0
|
golang.org/x/time v0.10.0
|
||||||
)
|
)
|
||||||
|
|||||||
2
go.sum
2
go.sum
@@ -172,8 +172,6 @@ github.com/spf13/viper v1.19.0 h1:RWq5SEjt8o25SROyN3z2OrDB9l7RPd3lwTWU8EcEdcI=
|
|||||||
github.com/spf13/viper v1.19.0/go.mod h1:GQUN9bilAbhU/jgc1bKs99f/suXKeUMct8Adx5+Ntkg=
|
github.com/spf13/viper v1.19.0/go.mod h1:GQUN9bilAbhU/jgc1bKs99f/suXKeUMct8Adx5+Ntkg=
|
||||||
github.com/stretchr/testify v1.10.0 h1:Xv5erBjTwe/5IxqUQTdXv5kgmIvbHo3QQyRwhJsOfJA=
|
github.com/stretchr/testify v1.10.0 h1:Xv5erBjTwe/5IxqUQTdXv5kgmIvbHo3QQyRwhJsOfJA=
|
||||||
github.com/stretchr/testify v1.10.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY=
|
github.com/stretchr/testify v1.10.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY=
|
||||||
github.com/studio-b12/gowebdav v0.10.0 h1:Yewz8FFiadcGEu4hxS/AAJQlHelndqln1bns3hcJIYc=
|
|
||||||
github.com/studio-b12/gowebdav v0.10.0/go.mod h1:bHA7t77X/QFExdeAnDzK6vKM34kEZAcE1OX4MfiwjkE=
|
|
||||||
github.com/subosito/gotenv v1.6.0 h1:9NlTDc1FTs4qu0DDq7AEtTPNw6SVm7uBMsUCUjABIf8=
|
github.com/subosito/gotenv v1.6.0 h1:9NlTDc1FTs4qu0DDq7AEtTPNw6SVm7uBMsUCUjABIf8=
|
||||||
github.com/subosito/gotenv v1.6.0/go.mod h1:Dk4QP5c2W3ibzajGcXpNraDfq2IrhjMIvMSWPKKo0FU=
|
github.com/subosito/gotenv v1.6.0/go.mod h1:Dk4QP5c2W3ibzajGcXpNraDfq2IrhjMIvMSWPKKo0FU=
|
||||||
github.com/tcnksm/go-gitconfig v0.1.2 h1:iiDhRitByXAEyjgBqsKi9QU4o2TNtv9kPP3RgPgXBPw=
|
github.com/tcnksm/go-gitconfig v0.1.2 h1:iiDhRitByXAEyjgBqsKi9QU4o2TNtv9kPP3RgPgXBPw=
|
||||||
|
|||||||
70
storage/webdav/client.go
Normal file
70
storage/webdav/client.go
Normal file
@@ -0,0 +1,70 @@
|
|||||||
|
package webdav
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"fmt"
|
||||||
|
"io"
|
||||||
|
"net/http"
|
||||||
|
"strings"
|
||||||
|
)
|
||||||
|
|
||||||
|
type Client struct {
|
||||||
|
BaseURL string
|
||||||
|
Username string
|
||||||
|
Password string
|
||||||
|
httpClient *http.Client
|
||||||
|
}
|
||||||
|
|
||||||
|
func NewClient(baseURL, username, password string, httpClient *http.Client) *Client {
|
||||||
|
if !strings.HasSuffix(baseURL, "/") {
|
||||||
|
baseURL += "/"
|
||||||
|
}
|
||||||
|
if httpClient == nil {
|
||||||
|
httpClient = http.DefaultClient
|
||||||
|
}
|
||||||
|
return &Client{
|
||||||
|
BaseURL: baseURL,
|
||||||
|
Username: username,
|
||||||
|
Password: password,
|
||||||
|
httpClient: httpClient,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func (c *Client) doRequest(ctx context.Context, method, url string, body io.Reader) (*http.Response, error) {
|
||||||
|
req, err := http.NewRequestWithContext(ctx, method, url, body)
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
if c.Username != "" && c.Password != "" {
|
||||||
|
req.SetBasicAuth(c.Username, c.Password)
|
||||||
|
}
|
||||||
|
return c.httpClient.Do(req)
|
||||||
|
}
|
||||||
|
|
||||||
|
func (c *Client) MkDir(ctx context.Context, dirPath string) error {
|
||||||
|
url := c.BaseURL + dirPath
|
||||||
|
resp, err := c.doRequest(ctx, "MKCOL", url, nil)
|
||||||
|
if err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
defer resp.Body.Close()
|
||||||
|
|
||||||
|
if resp.StatusCode >= 200 && resp.StatusCode < 300 {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
return fmt.Errorf("MKCOL: %s", resp.Status)
|
||||||
|
}
|
||||||
|
|
||||||
|
func (c *Client) WriteFile(ctx context.Context, remotePath string, content io.Reader) error {
|
||||||
|
url := c.BaseURL + remotePath
|
||||||
|
resp, err := c.doRequest(ctx, "PUT", url, content)
|
||||||
|
if err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
defer resp.Body.Close()
|
||||||
|
|
||||||
|
if resp.StatusCode >= 200 && resp.StatusCode < 300 {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
return fmt.Errorf("PUT: %s", resp.Status)
|
||||||
|
}
|
||||||
@@ -1,50 +1,58 @@
|
|||||||
package webdav
|
package webdav
|
||||||
|
|
||||||
// TODO: gowebdav's WriteStream impl cause high memory usage, need to implement our own WriteStream
|
import (
|
||||||
// type WebdavWriter struct {
|
"context"
|
||||||
// pipeWriter *io.PipeWriter
|
"fmt"
|
||||||
// done chan error
|
"io"
|
||||||
// path string
|
"path"
|
||||||
// }
|
|
||||||
|
|
||||||
// func (w *WebdavWriter) Write(p []byte) (n int, err error) {
|
"github.com/krau/SaveAny-Bot/logger"
|
||||||
// return w.pipeWriter.Write(p)
|
)
|
||||||
// }
|
|
||||||
|
|
||||||
// func (w *WebdavWriter) Close() error {
|
type WebdavWriter struct {
|
||||||
// if err := w.pipeWriter.Close(); err != nil {
|
pipeWriter *io.PipeWriter
|
||||||
// return err
|
done chan error
|
||||||
// }
|
path string
|
||||||
// if err := <-w.done; err != nil {
|
}
|
||||||
// return fmt.Errorf("upload failed: %w", err)
|
|
||||||
// }
|
|
||||||
|
|
||||||
// return nil
|
func (w *WebdavWriter) Write(p []byte) (n int, err error) {
|
||||||
// }
|
return w.pipeWriter.Write(p)
|
||||||
|
}
|
||||||
|
|
||||||
// func (w *Webdav) NewUploadStream(ctx context.Context, storagePath string) (io.WriteCloser, error) {
|
func (w *WebdavWriter) Close() error {
|
||||||
// if err := w.client.MkdirAll(path.Dir(storagePath), os.ModePerm); err != nil {
|
if err := w.pipeWriter.Close(); err != nil {
|
||||||
// logger.L.Errorf("Failed to create directory %s: %v", path.Dir(storagePath), err)
|
return err
|
||||||
// return nil, ErrFailedToCreateDirectory
|
}
|
||||||
// }
|
if err := <-w.done; err != nil {
|
||||||
// pipeReader, pipeWriter := io.Pipe()
|
return fmt.Errorf("upload failed: %w", err)
|
||||||
// done := make(chan error, 1)
|
}
|
||||||
// go func() {
|
|
||||||
// defer func() {
|
|
||||||
// if err := recover(); err != nil {
|
|
||||||
// done <- fmt.Errorf("panic during upload: %v", err)
|
|
||||||
// }
|
|
||||||
// }()
|
|
||||||
|
|
||||||
// err := w.client.WriteStream(storagePath, pipeReader, os.ModePerm)
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
// pipeReader.Close()
|
func (w *Webdav) NewUploadStream(ctx context.Context, storagePath string) (io.WriteCloser, error) {
|
||||||
// done <- err
|
if err := w.client.MkDir(ctx, path.Dir(storagePath)); err != nil {
|
||||||
// }()
|
logger.L.Errorf("Failed to create directory %s: %v", path.Dir(storagePath), err)
|
||||||
|
return nil, ErrFailedToCreateDirectory
|
||||||
|
}
|
||||||
|
pipeReader, pipeWriter := io.Pipe()
|
||||||
|
done := make(chan error, 1)
|
||||||
|
go func() {
|
||||||
|
defer func() {
|
||||||
|
if err := recover(); err != nil {
|
||||||
|
done <- fmt.Errorf("panic during upload: %v", err)
|
||||||
|
}
|
||||||
|
}()
|
||||||
|
|
||||||
// return &WebdavWriter{
|
err := w.client.WriteFile(ctx, storagePath, pipeReader)
|
||||||
// pipeWriter: pipeWriter,
|
|
||||||
// done: done,
|
pipeReader.Close()
|
||||||
// path: storagePath,
|
done <- err
|
||||||
// }, nil
|
}()
|
||||||
// }
|
|
||||||
|
return &WebdavWriter{
|
||||||
|
pipeWriter: pipeWriter,
|
||||||
|
done: done,
|
||||||
|
path: storagePath,
|
||||||
|
}, nil
|
||||||
|
}
|
||||||
|
|||||||
@@ -3,6 +3,7 @@ package webdav
|
|||||||
import (
|
import (
|
||||||
"context"
|
"context"
|
||||||
"fmt"
|
"fmt"
|
||||||
|
"net/http"
|
||||||
"os"
|
"os"
|
||||||
"path"
|
"path"
|
||||||
"time"
|
"time"
|
||||||
@@ -10,12 +11,11 @@ import (
|
|||||||
"github.com/krau/SaveAny-Bot/config"
|
"github.com/krau/SaveAny-Bot/config"
|
||||||
"github.com/krau/SaveAny-Bot/logger"
|
"github.com/krau/SaveAny-Bot/logger"
|
||||||
"github.com/krau/SaveAny-Bot/types"
|
"github.com/krau/SaveAny-Bot/types"
|
||||||
"github.com/studio-b12/gowebdav"
|
|
||||||
)
|
)
|
||||||
|
|
||||||
type Webdav struct {
|
type Webdav struct {
|
||||||
config config.WebdavStorageConfig
|
config config.WebdavStorageConfig
|
||||||
client *gowebdav.Client
|
client *Client
|
||||||
}
|
}
|
||||||
|
|
||||||
func (w *Webdav) Init(cfg config.StorageConfig) error {
|
func (w *Webdav) Init(cfg config.StorageConfig) error {
|
||||||
@@ -27,12 +27,9 @@ func (w *Webdav) Init(cfg config.StorageConfig) error {
|
|||||||
return err
|
return err
|
||||||
}
|
}
|
||||||
w.config = *webdavConfig
|
w.config = *webdavConfig
|
||||||
client := gowebdav.NewClient(webdavConfig.URL, webdavConfig.Username, webdavConfig.Password)
|
w.client = NewClient(w.config.URL, w.config.Username, w.config.Password, &http.Client{
|
||||||
if err := client.Connect(); err != nil {
|
Timeout: time.Hour * 12,
|
||||||
return fmt.Errorf("failed to connect to webdav server: %w", err)
|
})
|
||||||
}
|
|
||||||
client.SetTimeout(12 * time.Hour)
|
|
||||||
w.client = client
|
|
||||||
return nil
|
return nil
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -46,7 +43,7 @@ func (w *Webdav) Name() string {
|
|||||||
|
|
||||||
func (w *Webdav) Save(ctx context.Context, filePath, storagePath string) error {
|
func (w *Webdav) Save(ctx context.Context, filePath, storagePath string) error {
|
||||||
logger.L.Infof("Saving file %s to %s", filePath, storagePath)
|
logger.L.Infof("Saving file %s to %s", filePath, storagePath)
|
||||||
if err := w.client.MkdirAll(path.Dir(storagePath), os.ModePerm); err != nil {
|
if err := w.client.MkDir(ctx, path.Dir(storagePath)); err != nil {
|
||||||
logger.L.Errorf("Failed to create directory %s: %v", path.Dir(storagePath), err)
|
logger.L.Errorf("Failed to create directory %s: %v", path.Dir(storagePath), err)
|
||||||
return ErrFailedToCreateDirectory
|
return ErrFailedToCreateDirectory
|
||||||
}
|
}
|
||||||
@@ -57,7 +54,7 @@ func (w *Webdav) Save(ctx context.Context, filePath, storagePath string) error {
|
|||||||
}
|
}
|
||||||
defer file.Close()
|
defer file.Close()
|
||||||
|
|
||||||
if err := w.client.WriteStream(storagePath, file, os.ModePerm); err != nil {
|
if err := w.client.WriteFile(ctx, storagePath, file); err != nil {
|
||||||
logger.L.Errorf("Failed to write file %s: %v", storagePath, err)
|
logger.L.Errorf("Failed to write file %s: %v", storagePath, err)
|
||||||
return ErrFailedToWriteFile
|
return ErrFailedToWriteFile
|
||||||
}
|
}
|
||||||
|
|||||||
Reference in New Issue
Block a user