From e60a87781695f06cd98068c8335233f051af8b5b Mon Sep 17 00:00:00 2001 From: Roman Tomjak <6570684+romantomjak@users.noreply.github.com> Date: Sat, 31 Oct 2020 09:17:24 +0000 Subject: [PATCH] implement ghetto file upload for small files --- Makefile | 2 +- README.md | 7 +- b2/client.go | 16 ++- b2/file.go | 79 ++++++++++++++- command/bucket_create_test.go | 1 + command/commands.go | 5 + command/put.go | 183 ++++++++++++++++++++++++++++++++++ command/put_test.go | 108 ++++++++++++++++++++ 8 files changed, 393 insertions(+), 8 deletions(-) create mode 100644 command/put.go create mode 100644 command/put_test.go diff --git a/Makefile b/Makefile index 4c47904..dfbdead 100644 --- a/Makefile +++ b/Makefile @@ -1,6 +1,6 @@ SHELL = bash PROJECT_ROOT := $(patsubst %/,%,$(dir $(abspath $(lastword $(MAKEFILE_LIST))))) -VERSION := 0.4.0 +VERSION := 0.5.0 GIT_COMMIT := $(shell git rev-parse --short HEAD) GO_PKGS := $(shell go list ./...) diff --git a/README.md b/README.md index 968bda3..f0084b6 100644 --- a/README.md +++ b/README.md @@ -10,10 +10,14 @@ This project is in development phase. You can try it with latest release version ## Installation +Download and install using go get: + ```sh go get -u github.com/romantomjak/b2 ``` +or grab a binary from [releases](https://github.com/romantomjak/b2/releases/latest) section! + ## Usage ```sh @@ -33,6 +37,7 @@ Available commands are: create Create a new bucket get Download files list List files and buckets + put Upload files version Prints the client version ``` @@ -47,7 +52,7 @@ This is how far I've gotten: - [x] List all buckets - [ ] Update settings for a bucket - [x] List files in a bucket -- [ ] Upload small files +- [x] Upload small files (<100 MB) - [ ] Upload large files - [x] Download a file diff --git a/b2/client.go b/b2/client.go index d15e0b4..7a8fb40 100644 --- a/b2/client.go +++ b/b2/client.go @@ -197,15 +197,21 @@ func (c *Client) newRequest(ctx context.Context, method, path string, body inter u := c.baseURL.ResolveReference(rel) - buf := new(bytes.Buffer) + var b io.Reader if body != nil { - err := json.NewEncoder(buf).Encode(body) - if err != nil { - return nil, err + if r, ok := body.(io.Reader); ok { + b = r + } else { + buf := new(bytes.Buffer) + err := json.NewEncoder(buf).Encode(body) + if err != nil { + return nil, err + } + b = buf } } - req, err := http.NewRequestWithContext(ctx, method, u.String(), buf) + req, err := http.NewRequestWithContext(ctx, method, u.String(), b) if err != nil { return nil, err } diff --git a/b2/file.go b/b2/file.go index ae62182..dadaa1d 100644 --- a/b2/file.go +++ b/b2/file.go @@ -2,12 +2,17 @@ package b2 import ( "context" + "crypto/sha1" + "fmt" "io" "net/http" + "net/url" + "os" ) const ( - listFilesURL = "b2api/v2/b2_list_file_names" + listFilesURL = "b2api/v2/b2_list_file_names" + fileUploadURL = "b2api/v2/b2_get_upload_url" ) // File describes a File or a Folder in a Bucket @@ -38,6 +43,18 @@ type fileListRoot struct { NextFileName string `json:"nextFileName"` } +// UploadAuthorizationRequest represents a request to obtain a URL for uploading files +type UploadAuthorizationRequest struct { + BucketID string `json:"bucketId"` +} + +// UploadAuthorization contains the information for uploading a file +type UploadAuthorization struct { + BucketID string `json:"bucketId"` + UploadURL string `json:"uploadUrl"` + Token string `json:"authorizationToken"` +} + // FileService handles communication with the File related methods of the // B2 API type FileService struct { @@ -74,3 +91,63 @@ func (s *FileService) Download(ctx context.Context, url string, w io.Writer) (*h return resp, err } + +// UploadAuthorization returns the information for uploading a file +func (s *FileService) UploadAuthorization(ctx context.Context, uploadAuthorizationRequest *UploadAuthorizationRequest) (*UploadAuthorization, *http.Response, error) { + req, err := s.client.NewRequest(ctx, http.MethodPost, fileUploadURL, uploadAuthorizationRequest) + if err != nil { + return nil, nil, err + } + + auth := new(UploadAuthorization) + resp, err := s.client.Do(req, auth) + if err != nil { + return nil, resp, err + } + + return auth, resp, nil +} + +// Upload a file +func (s *FileService) Upload(ctx context.Context, uploadAuthorization *UploadAuthorization, src, dst string) (*File, *http.Response, error) { + f, err := os.Open(src) + if err != nil { + return nil, nil, err + } + defer f.Close() + + info, err := f.Stat() + if err != nil { + return nil, nil, err + } + + hash := sha1.New() + _, err = io.Copy(hash, f) + if err != nil { + return nil, nil, err + } + sha1 := fmt.Sprintf("%x", hash.Sum(nil)) + + f.Seek(0, 0) + + req, err := s.client.NewRequest(ctx, http.MethodPost, uploadAuthorization.UploadURL, f) + if err != nil { + return nil, nil, err + } + + req.ContentLength = info.Size() + + req.Header.Set("Authorization", uploadAuthorization.Token) + req.Header.Set("X-Bz-File-Name", url.QueryEscape(dst)) + req.Header.Set("Content-Type", "b2/x-auto") + req.Header.Set("X-Bz-Content-Sha1", sha1) + req.Header.Set("X-Bz-Info-src_last_modified_millis", fmt.Sprintf("%d", info.ModTime().Unix()*1000)) + + file := new(File) + resp, err := s.client.Do(req, file) + if err != nil { + return nil, resp, err + } + + return file, resp, nil +} diff --git a/command/bucket_create_test.go b/command/bucket_create_test.go index 4bb2f4c..e6502bd 100644 --- a/command/bucket_create_test.go +++ b/command/bucket_create_test.go @@ -77,6 +77,7 @@ func TestCreateBucketCommand_BucketCreateRequest(t *testing.T) { } _ = cmd.Run([]string{"my-bucket"}) + // TODO: write bucket response // return code is ignored on purpose here. // fake b2_create_bucket handler is not writing the response, so // the command will fail and return 1 diff --git a/command/commands.go b/command/commands.go index 5fe9b3e..1c003a7 100644 --- a/command/commands.go +++ b/command/commands.go @@ -27,6 +27,11 @@ func Commands(ui cli.Ui) map[string]cli.CommandFactory { baseCommand: baseCommand, }, nil }, + "put": func() (cli.Command, error) { + return &PutCommand{ + baseCommand: baseCommand, + }, nil + }, "version": func() (cli.Command, error) { return &VersionCommand{ baseCommand: baseCommand, diff --git a/command/put.go b/command/put.go new file mode 100644 index 0000000..d90c031 --- /dev/null +++ b/command/put.go @@ -0,0 +1,183 @@ +package command + +import ( + "context" + "errors" + "fmt" + "os" + "path" + "strings" + + "github.com/romantomjak/b2/b2" +) + +type PutCommand struct { + *baseCommand +} + +func (c *PutCommand) Help() string { + helpText := ` +Usage: b2 put + + Uploads the contents of source to destination. If destination + contains a trailing slash it is treated as a directory and + file is uploaded keeping the original filename. + +General Options: + + ` + c.generalOptions() + return strings.TrimSpace(helpText) +} + +func (c *PutCommand) Synopsis() string { + return "Upload files" +} + +func (c *PutCommand) Name() string { return "put" } + +func (c *PutCommand) Run(args []string) int { + flags := c.flagSet() + flags.Usage = func() { c.ui.Output(c.Help()) } + + if err := flags.Parse(args); err != nil { + return 1 + } + + // Check that we got both arguments + args = flags.Args() + numArgs := len(args) + if numArgs != 2 { + c.ui.Error("This command takes two arguments: and ") + return 1 + } + + // Check that source file exists + if !fileExists(args[0]) { + c.ui.Error(fmt.Sprintf("File does not exist: %s", args[0])) + return 1 + } + + // FIXME: remove when large file upload is implemented + err := checkMaxFileSize(args[0]) + if err != nil { + c.ui.Error("Large file upload is not yet implemented. Maximum file size is 100 MB") + return 1 + } + + bucketName, filePrefix := destinationBucketAndFilename(args[0], args[1]) + + // TODO: caching bucket name:id mappings could save this request + bucket, err := c.findBucketByName(bucketName) + if err != nil { + c.ui.Error(fmt.Sprintf("Error: %v", err)) + return 1 + } + + // Create a client + client, err := c.Client() + if err != nil { + c.ui.Error(fmt.Sprintf("Error: %v", err)) + return 1 + } + + // Request upload url + ctx := context.TODO() + + uploadAuthReq := &b2.UploadAuthorizationRequest{ + BucketID: bucket.ID, + } + uploadAuth, _, err := client.File.UploadAuthorization(ctx, uploadAuthReq) + if err != nil { + c.ui.Error(fmt.Sprintf("Error: %v", err)) + return 1 + } + + _, _, err = client.File.Upload(ctx, uploadAuth, args[0], filePrefix) + if err != nil { + c.ui.Error(err.Error()) + return 1 + } + + c.ui.Output(fmt.Sprintf("Uploaded %q to %q", args[0], path.Join(bucket.Name, filePrefix))) + + return 0 +} + +func (c *PutCommand) findBucketByName(name string) (*b2.Bucket, error) { + client, err := c.Client() + if err != nil { + c.ui.Error(fmt.Sprintf("Error: %v", err)) + return nil, err + } + + req := &b2.BucketListRequest{ + AccountID: client.AccountID, + Name: name, + } + + ctx := context.TODO() + + buckets, _, err := client.Bucket.List(ctx, req) + if err != nil { + return nil, err + } + + if len(buckets) == 0 { + return nil, fmt.Errorf("bucket with name %q was not found", name) + } + + return &buckets[0], nil +} + +// fileExists checks if a file exists and is not a directory +func fileExists(filename string) bool { + info, err := os.Stat(filename) + if os.IsNotExist(err) { + return false + } + return !info.IsDir() +} + +// checkMaxFileSize checks that file is a "small" file +func checkMaxFileSize(filename string) error { + var maxFileSize int64 = 100 << (10 * 2) // 100 mb + + info, err := os.Stat(filename) + if err != nil { + return err + } + + if info.Size() > maxFileSize { + return errors.New("file is too big") + } + + return nil +} + +// destinationBucketAndFilename returns upload bucket and filePrefix +// +// b2 does not have a concept of folders, so if destination contains +// a trailing slash it is treated as a directory and file is uploaded +// keeping the original filename. If destination is simply a bucket +// name, it is asumed the destination is "/" and filename is preserved +func destinationBucketAndFilename(source, destination string) (string, string) { + originalFilename := path.Base(source) + + destinationParts := strings.SplitN(destination, "/", 2) + bucketName := destinationParts[0] + filePrefix := "" + + if len(destinationParts) > 1 { + if strings.HasSuffix(destinationParts[1], "/") { + filePrefix = path.Join(destinationParts[1], originalFilename) + } else { + filePrefix = destinationParts[1] + } + } + + if filePrefix == "" { + filePrefix = originalFilename + } + + return bucketName, filePrefix +} diff --git a/command/put_test.go b/command/put_test.go new file mode 100644 index 0000000..ada18f4 --- /dev/null +++ b/command/put_test.go @@ -0,0 +1,108 @@ +package command + +import ( + "fmt" + "io/ioutil" + "net/http" + "os" + "path" + "testing" + + "github.com/mitchellh/cli" + "github.com/romantomjak/b2/b2" + "github.com/romantomjak/b2/testutil" + "github.com/stretchr/testify/assert" +) + +func TestPutCommand_FilePrefix(t *testing.T) { + filenameTests := []struct { + originalFilename string + destinationFilename string + filePrefix string + }{ + {"file1.txt", "bucket", "file1.txt"}, + {"file1.txt", "bucket/", "file1.txt"}, + {"file1.txt", "bucket/dir", "dir"}, + {"file1.txt", "bucket/dir/", "dir/file1.txt"}, + {"file1.txt", "bucket/dir/a", "dir/a"}, + {"file1.txt", "bucket/dir/a/", "dir/a/file1.txt"}, + } + for _, tt := range filenameTests { + t.Run(tt.destinationFilename, func(t *testing.T) { + _, filePrefix := destinationBucketAndFilename(tt.originalFilename, tt.destinationFilename) + assert.Equal(t, tt.filePrefix, filePrefix) + }) + } +} + +func TestPutCommand_CanUploadFile(t *testing.T) { + server, mux := testutil.NewServer() + defer server.Close() + + mux.HandleFunc("/b2api/v2/b2_list_buckets", func(w http.ResponseWriter, r *http.Request) { + fmt.Fprint(w, `{ + "buckets": [ + { + "accountId": "30f20426f0b1", + "bucketId": "87ba238875c6214145260818", + "bucketInfo": {}, + "bucketName": "Secret-Documents", + "bucketType": "allPrivate", + "lifecycleRules": [] + } ] + }`) + }) + + mux.HandleFunc("/b2api/v2/b2_get_upload_url", func(w http.ResponseWriter, r *http.Request) { + fmt.Fprintf(w, `{ + "bucketId": "87ba238875c6214145260818", + "uploadUrl": "%s", + "authorizationToken": "some-secret-token" + }`, server.URL) + }) + + mux.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) { + fileBytes, err := ioutil.ReadAll(r.Body) + defer r.Body.Close() + + assert.NoError(t, err) + assert.Equal(t, []byte("This file is not empty."), fileBytes) + + fmt.Fprint(w, `{ + "fileId": "4_h4a48fe8875c6214145260818_f000000000000472a_d20140104_m032022_c001_v0000123_t0104", + "fileName": "typing_test.txt", + "accountId": "d522aa47a10f", + "bucketId": "4a48fe8875c6214145260818", + "contentLength": 46, + "contentSha1": "bae5ed658ab3546aee12f23f36392f35dba1ebdd", + "contentType": "text/plain", + "fileInfo": { + "author": "unknown" + } + }`) + }) + + cache, _ := b2.NewInMemoryCache() + client, _ := b2.NewClient("key-id", "key-secret", b2.SetBaseURL(server.URL), b2.SetCache(cache)) + + ui := cli.NewMockUi() + cmd := &PutCommand{ + baseCommand: &baseCommand{ui: ui, client: client}, + } + + tmpFile, _ := ioutil.TempFile(os.TempDir(), "b2-cli-test-") + defer os.Remove(tmpFile.Name()) + + tmpFile.Write([]byte("This file is not empty.")) + tmpFile.Close() + + src := tmpFile.Name() + dst := "Secret-Documents" + + code := cmd.Run([]string{src, dst}) + assert.Equal(t, 0, code) + + out := ui.OutputWriter.String() + filename := fmt.Sprintf("%s/%s", dst, path.Base(tmpFile.Name())) + assert.Contains(t, out, fmt.Sprintf("Uploaded %q to %q", src, filename)) +}