From 802c9083848ccb3961b4ca92dbe51382bf0be0e7 Mon Sep 17 00:00:00 2001 From: krau <71133316+krau@users.noreply.github.com> Date: Sat, 1 Mar 2025 12:06:55 +0800 Subject: [PATCH] feat: refactor webdav client and implement custom upload stream handling --- go.mod | 1 - go.sum | 2 - storage/webdav/client.go | 70 +++++++++++++++++++++++++++++++ storage/webdav/stream.go | 90 ++++++++++++++++++++++------------------ storage/webdav/webdav.go | 17 ++++---- 5 files changed, 126 insertions(+), 54 deletions(-) create mode 100644 storage/webdav/client.go diff --git a/go.mod b/go.mod index 95ebd3e..2667698 100644 --- a/go.mod +++ b/go.mod @@ -12,7 +12,6 @@ require ( github.com/rhysd/go-github-selfupdate v1.2.3 github.com/spf13/cobra v1.8.1 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/time v0.10.0 ) diff --git a/go.sum b/go.sum index a017ae5..f9b8233 100644 --- a/go.sum +++ b/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/stretchr/testify v1.10.0 h1:Xv5erBjTwe/5IxqUQTdXv5kgmIvbHo3QQyRwhJsOfJA= 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/go.mod h1:Dk4QP5c2W3ibzajGcXpNraDfq2IrhjMIvMSWPKKo0FU= github.com/tcnksm/go-gitconfig v0.1.2 h1:iiDhRitByXAEyjgBqsKi9QU4o2TNtv9kPP3RgPgXBPw= diff --git a/storage/webdav/client.go b/storage/webdav/client.go new file mode 100644 index 0000000..8092603 --- /dev/null +++ b/storage/webdav/client.go @@ -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) +} diff --git a/storage/webdav/stream.go b/storage/webdav/stream.go index d490c3a..ac4604a 100644 --- a/storage/webdav/stream.go +++ b/storage/webdav/stream.go @@ -1,50 +1,58 @@ package webdav -// TODO: gowebdav's WriteStream impl cause high memory usage, need to implement our own WriteStream -// type WebdavWriter struct { -// pipeWriter *io.PipeWriter -// done chan error -// path string -// } +import ( + "context" + "fmt" + "io" + "path" -// func (w *WebdavWriter) Write(p []byte) (n int, err error) { -// return w.pipeWriter.Write(p) -// } + "github.com/krau/SaveAny-Bot/logger" +) -// func (w *WebdavWriter) Close() error { -// if err := w.pipeWriter.Close(); err != nil { -// return err -// } -// if err := <-w.done; err != nil { -// return fmt.Errorf("upload failed: %w", err) -// } +type WebdavWriter struct { + pipeWriter *io.PipeWriter + done chan error + path string +} -// 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) { -// if err := w.client.MkdirAll(path.Dir(storagePath), os.ModePerm); 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) -// } -// }() +func (w *WebdavWriter) Close() error { + if err := w.pipeWriter.Close(); err != nil { + return err + } + if err := <-w.done; err != nil { + return fmt.Errorf("upload failed: %w", err) + } -// err := w.client.WriteStream(storagePath, pipeReader, os.ModePerm) + return nil +} -// pipeReader.Close() -// done <- err -// }() +func (w *Webdav) NewUploadStream(ctx context.Context, storagePath string) (io.WriteCloser, error) { + 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{ -// pipeWriter: pipeWriter, -// done: done, -// path: storagePath, -// }, nil -// } + err := w.client.WriteFile(ctx, storagePath, pipeReader) + + pipeReader.Close() + done <- err + }() + + return &WebdavWriter{ + pipeWriter: pipeWriter, + done: done, + path: storagePath, + }, nil +} diff --git a/storage/webdav/webdav.go b/storage/webdav/webdav.go index 27f0ded..7b7a2a9 100644 --- a/storage/webdav/webdav.go +++ b/storage/webdav/webdav.go @@ -3,6 +3,7 @@ package webdav import ( "context" "fmt" + "net/http" "os" "path" "time" @@ -10,12 +11,11 @@ import ( "github.com/krau/SaveAny-Bot/config" "github.com/krau/SaveAny-Bot/logger" "github.com/krau/SaveAny-Bot/types" - "github.com/studio-b12/gowebdav" ) type Webdav struct { config config.WebdavStorageConfig - client *gowebdav.Client + client *Client } func (w *Webdav) Init(cfg config.StorageConfig) error { @@ -27,12 +27,9 @@ func (w *Webdav) Init(cfg config.StorageConfig) error { return err } w.config = *webdavConfig - client := gowebdav.NewClient(webdavConfig.URL, webdavConfig.Username, webdavConfig.Password) - if err := client.Connect(); err != nil { - return fmt.Errorf("failed to connect to webdav server: %w", err) - } - client.SetTimeout(12 * time.Hour) - w.client = client + w.client = NewClient(w.config.URL, w.config.Username, w.config.Password, &http.Client{ + Timeout: time.Hour * 12, + }) return nil } @@ -46,7 +43,7 @@ func (w *Webdav) Name() string { func (w *Webdav) Save(ctx context.Context, filePath, storagePath string) error { 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) return ErrFailedToCreateDirectory } @@ -57,7 +54,7 @@ func (w *Webdav) Save(ctx context.Context, filePath, storagePath string) error { } 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) return ErrFailedToWriteFile }