Refactor the code

Un-expose the unnecessary methods
Update go mod
This commit is contained in:
Chun-Hung Tseng
2023-08-02 17:30:57 +02:00
parent 611ec433f8
commit 2fa36bc924
10 changed files with 895 additions and 865 deletions

View File

@@ -103,7 +103,7 @@ func RandomString(n int) string {
func createFolderExpectError(t *testing.T, ctx context.Context, protonDrive *ProtonDrive, parent, name string, expectedError error) {
parentLink := protonDrive.RootLink
if parent != "" {
targetFolderLink, err := protonDrive.SearchByNameRecursivelyFromRoot(ctx, parent, true, false)
targetFolderLink, err := protonDrive.searchByNameRecursivelyFromRoot(ctx, parent, true, false)
if err != nil {
t.Fatal(err)
}
@@ -129,7 +129,7 @@ func createFolder(t *testing.T, ctx context.Context, protonDrive *ProtonDrive, p
func uploadFileByReader(t *testing.T, ctx context.Context, protonDrive *ProtonDrive, parent, name string, in io.Reader, testParam int) {
parentLink := protonDrive.RootLink
if parent != "" {
targetFolderLink, err := protonDrive.SearchByNameRecursivelyFromRoot(ctx, parent, true, false)
targetFolderLink, err := protonDrive.searchByNameRecursivelyFromRoot(ctx, parent, true, false)
if err != nil {
t.Fatal(err)
}
@@ -151,7 +151,7 @@ func uploadFileByReader(t *testing.T, ctx context.Context, protonDrive *ProtonDr
func uploadFileByFilepathWithError(t *testing.T, ctx context.Context, protonDrive *ProtonDrive, parent, name string, filepath string, testParam int, expectedError error) {
parentLink := protonDrive.RootLink
if parent != "" {
targetFolderLink, err := protonDrive.SearchByNameRecursivelyFromRoot(ctx, parent, true, false)
targetFolderLink, err := protonDrive.searchByNameRecursivelyFromRoot(ctx, parent, true, false)
if err != nil {
t.Fatal(err)
}
@@ -194,7 +194,7 @@ func downloadFile(t *testing.T, ctx context.Context, protonDrive *ProtonDrive, p
func downloadFileWithOffset(t *testing.T, ctx context.Context, protonDrive *ProtonDrive, parent, name string, filepath string, data string, offset int64) {
parentLink := protonDrive.RootLink
if parent != "" {
targetFolderLink, err := protonDrive.SearchByNameRecursivelyFromRoot(ctx, parent, true, false)
targetFolderLink, err := protonDrive.searchByNameRecursivelyFromRoot(ctx, parent, true, false)
if err != nil {
t.Fatal(err)
}
@@ -208,7 +208,7 @@ func downloadFileWithOffset(t *testing.T, ctx context.Context, protonDrive *Prot
t.Fatalf("parentLink is not of folder type")
}
targetFileLink, err := protonDrive.SearchByNameRecursivelyFromRoot(ctx, name, false, false)
targetFileLink, err := protonDrive.searchByNameRecursivelyFromRoot(ctx, name, false, false)
if err != nil {
t.Fatal(err)
}
@@ -258,7 +258,7 @@ func downloadFileWithOffset(t *testing.T, ctx context.Context, protonDrive *Prot
}
func checkRevisions(protonDrive *ProtonDrive, ctx context.Context, t *testing.T, name string, totalRevisions, activeRevisions, draftRevisions, obseleteRevisions int) {
targetFileLink, err := protonDrive.SearchByNameRecursivelyFromRoot(ctx, name, false, true)
targetFileLink, err := protonDrive.searchByNameRecursivelyFromRoot(ctx, name, false, true)
if err != nil {
t.Fatal(err)
}
@@ -293,7 +293,7 @@ func checkRevisions(protonDrive *ProtonDrive, ctx context.Context, t *testing.T,
// During the integration test, the name much be unique since the link is returned by recursively search for the name from root
func deleteBySearchingFromRoot(t *testing.T, ctx context.Context, protonDrive *ProtonDrive, name string, isFolder bool, listAllActiveOrDraftFiles bool) {
targetLink, err := protonDrive.SearchByNameRecursivelyFromRoot(ctx, name, isFolder, listAllActiveOrDraftFiles)
targetLink, err := protonDrive.searchByNameRecursivelyFromRoot(ctx, name, isFolder, listAllActiveOrDraftFiles)
if err != nil {
t.Fatal(err)
}
@@ -317,7 +317,7 @@ func deleteBySearchingFromRoot(t *testing.T, ctx context.Context, protonDrive *P
func checkActiveFileListing(t *testing.T, ctx context.Context, protonDrive *ProtonDrive, expectedPaths []string) {
{
paths := make([]string, 0)
err := protonDrive.ListDirectoriesRecursively(ctx, protonDrive.MainShareKR, protonDrive.RootLink, false, -1, 0, true, "", &paths)
err := protonDrive.listDirectoriesRecursively(ctx, protonDrive.MainShareKR, protonDrive.RootLink, false, -1, 0, true, "", &paths)
if err != nil {
t.Fatal(err)
}
@@ -335,7 +335,7 @@ func checkActiveFileListing(t *testing.T, ctx context.Context, protonDrive *Prot
{
paths := make([]string, 0)
err := protonDrive.ListDirectoriesRecursively(ctx, protonDrive.MainShareKR, protonDrive.RootLink, false, -1, 0, false, "", &paths)
err := protonDrive.listDirectoriesRecursively(ctx, protonDrive.MainShareKR, protonDrive.RootLink, false, -1, 0, false, "", &paths)
if err != nil {
t.Fatal(err)
}
@@ -360,11 +360,11 @@ func checkActiveFileListing(t *testing.T, ctx context.Context, protonDrive *Prot
}
func moveFolder(t *testing.T, ctx context.Context, protonDrive *ProtonDrive, srcFolderName, dstParentFolderName string) {
targetSrcFolderLink, err := protonDrive.SearchByNameRecursivelyFromRoot(ctx, srcFolderName, true, false)
targetSrcFolderLink, err := protonDrive.searchByNameRecursivelyFromRoot(ctx, srcFolderName, true, false)
if err != nil {
t.Fatal(err)
}
targetDestFolderLink, err := protonDrive.SearchByNameRecursivelyFromRoot(ctx, dstParentFolderName, true, false)
targetDestFolderLink, err := protonDrive.searchByNameRecursivelyFromRoot(ctx, dstParentFolderName, true, false)
if err != nil {
t.Fatal(err)
}
@@ -379,11 +379,11 @@ func moveFolder(t *testing.T, ctx context.Context, protonDrive *ProtonDrive, src
}
func moveFile(t *testing.T, ctx context.Context, protonDrive *ProtonDrive, srcFileName, dstParentFolderName string) {
targetSrcFileLink, err := protonDrive.SearchByNameRecursivelyFromRoot(ctx, srcFileName, false, false)
targetSrcFileLink, err := protonDrive.searchByNameRecursivelyFromRoot(ctx, srcFileName, false, false)
if err != nil {
t.Fatal(err)
}
targetDestFolderLink, err := protonDrive.SearchByNameRecursivelyFromRoot(ctx, dstParentFolderName, true, false)
targetDestFolderLink, err := protonDrive.searchByNameRecursivelyFromRoot(ctx, dstParentFolderName, true, false)
if err != nil {
t.Fatal(err)
}

632
file.go
View File

@@ -1,21 +1,9 @@
package proton_api_bridge
import (
"bufio"
"bytes"
"context"
"crypto/sha1"
"crypto/sha256"
"encoding/base64"
"encoding/hex"
"io"
"log"
"mime"
"os"
"path/filepath"
"time"
"github.com/ProtonMail/gopenpgp/v2/crypto"
"github.com/henrybear327/go-proton-api"
"github.com/relvacode/iso8601"
)
@@ -27,86 +15,6 @@ type FileSystemAttrs struct {
Digests string // sha1 string
}
type FileDownloadReader struct {
protonDrive *ProtonDrive
ctx context.Context
data *bytes.Buffer
nodeKR *crypto.KeyRing
sessionKey *crypto.SessionKey
revision *proton.Revision
nextRevision int
isEOF bool
// TODO: integrity check if the entire file is read
}
func (r *FileDownloadReader) Read(p []byte) (int, error) {
if r.data.Len() == 0 {
// TODO: do we have memory sharing bug?
// to avoid sharing the underlying buffer array across re-population
r.data = bytes.NewBuffer(nil)
// we download and decrypt more content
err := r.populateBufferOnRead()
if err != nil {
return 0, err
}
if r.isEOF {
// if the file has been downloaded entirely, we return EOF
return 0, io.EOF
}
}
return r.data.Read(p)
}
func (r *FileDownloadReader) Close() error {
r.protonDrive = nil
return nil
}
func (reader *FileDownloadReader) populateBufferOnRead() error {
if len(reader.revision.Blocks) == 0 || len(reader.revision.Blocks) == reader.nextRevision {
reader.isEOF = true
return nil
}
offset := reader.nextRevision
for i := offset; i-offset < DOWNLOAD_BATCH_BLOCK_SIZE && i < len(reader.revision.Blocks); i++ {
// TODO: parallel download
blockReader, err := reader.protonDrive.c.GetBlock(reader.ctx, reader.revision.Blocks[i].BareURL, reader.revision.Blocks[i].Token)
if err != nil {
return err
}
defer blockReader.Close()
err = decryptBlockIntoBuffer(reader.sessionKey, reader.protonDrive.AddrKR, reader.nodeKR, reader.revision.Blocks[i].Hash, reader.revision.Blocks[i].EncSignature, reader.data, blockReader)
if err != nil {
return err
}
reader.nextRevision = i + 1
}
return nil
}
func (protonDrive *ProtonDrive) DownloadFileByID(ctx context.Context, linkID string, offset int64) (io.ReadCloser, int64, *FileSystemAttrs, error) {
/* It's like event system, we need to get the latest information before creating the move request! */
protonDrive.removeLinkIDFromCache(linkID, false)
link, err := protonDrive.getLink(ctx, linkID)
if err != nil {
return nil, 0, nil, err
}
return protonDrive.DownloadFile(ctx, link, offset)
}
func (protonDrive *ProtonDrive) GetRevisions(ctx context.Context, link *proton.Link, revisionType proton.RevisionState) ([]*proton.RevisionMetadata, error) {
revisions, err := protonDrive.c.ListRevisions(ctx, protonDrive.MainShare.ShareID, link.LinkID)
if err != nil {
@@ -226,543 +134,3 @@ func (protonDrive *ProtonDrive) GetActiveRevisionWithAttrs(ctx context.Context,
Digests: sha1Hash,
}, nil
}
func (protonDrive *ProtonDrive) DownloadFile(ctx context.Context, link *proton.Link, offset int64) (io.ReadCloser, int64, *FileSystemAttrs, error) {
if link.Type != proton.LinkTypeFile {
return nil, 0, nil, ErrLinkTypeMustToBeFileType
}
parentNodeKR, err := protonDrive.getLinkKRByID(ctx, link.ParentLinkID)
if err != nil {
return nil, 0, nil, err
}
nodeKR, err := link.GetKeyRing(parentNodeKR, protonDrive.AddrKR)
if err != nil {
return nil, 0, nil, err
}
sessionKey, err := link.GetSessionKey(nodeKR)
if err != nil {
return nil, 0, nil, err
}
revision, fileSystemAttrs, err := protonDrive.GetActiveRevisionWithAttrs(ctx, link)
if err != nil {
return nil, 0, nil, err
}
reader := &FileDownloadReader{
protonDrive: protonDrive,
ctx: ctx,
data: bytes.NewBuffer(nil),
nodeKR: nodeKR,
sessionKey: sessionKey,
revision: revision,
nextRevision: 0,
isEOF: false,
}
useFallbackDownload := false
if fileSystemAttrs != nil {
// based on offset, infer the nextRevision (0-based)
if fileSystemAttrs.BlockSizes == nil {
useFallbackDownload = true
} else {
// infer nextRevision
totalBytes := int64(0)
for i := 0; i < len(fileSystemAttrs.BlockSizes); i++ {
prevTotalBytes := totalBytes
totalBytes += fileSystemAttrs.BlockSizes[i]
if offset <= totalBytes {
offset = offset - prevTotalBytes
reader.nextRevision = i
break
}
}
// download will start from the specified block
n, err := io.CopyN(io.Discard, reader, offset)
if err != nil {
return nil, 0, nil, err
}
if int64(n) != offset {
return nil, 0, nil, ErrSeekOffsetAfterSkippingBlocks
}
}
}
if useFallbackDownload {
log.Println("Performing inefficient seek as metadata of encrypted file is missing")
n, err := io.CopyN(io.Discard, reader, offset)
if err != nil {
return nil, 0, nil, err
}
if int64(n) != offset {
return nil, 0, nil, ErrSeekOffsetAfterSkippingBlocks
}
}
return reader, link.Size, fileSystemAttrs, nil
}
func (protonDrive *ProtonDrive) UploadFileByReader(ctx context.Context, parentLinkID string, filename string, modTime time.Time, file io.Reader, testParam int) (string, *proton.RevisionXAttrCommon, error) {
parentLink, err := protonDrive.getLink(ctx, parentLinkID)
if err != nil {
return "", nil, err
}
return protonDrive.uploadFile(ctx, parentLink, filename, modTime, file, testParam)
}
func (protonDrive *ProtonDrive) UploadFileByPath(ctx context.Context, parentLink *proton.Link, filename string, filePath string, testParam int) (string, *proton.RevisionXAttrCommon, error) {
f, err := os.Open(filePath)
if err != nil {
return "", nil, err
}
defer f.Close()
info, err := os.Stat(filePath)
if err != nil {
return "", nil, err
}
in := bufio.NewReader(f)
return protonDrive.uploadFile(ctx, parentLink, filename, info.ModTime(), in, testParam)
}
func (protonDrive *ProtonDrive) handleRevisionConflict(ctx context.Context, link *proton.Link, createFileResp *proton.CreateFileRes) (string, bool, error) {
if link != nil {
linkID := link.LinkID
draftRevision, err := protonDrive.GetRevisions(ctx, link, proton.RevisionStateDraft)
if err != nil {
return "", false, err
}
// if we have a draft revision, depending on the user config, we can abort the upload or recreate a draft
// if we have no draft revision, then we can create a new draft revision directly (there is a restriction of 1 draft revision per file)
if len(draftRevision) > 0 {
// TODO: maintain clientUID to mark that this is our own draft (which can indicate failed upload attempt!)
if protonDrive.Config.ReplaceExistingDraft {
// Question: how do we observe for file upload cancellation -> clientUID?
// Random thoughts: if there are concurrent modification to the draft, the server should be able to catch this when commiting the revision
// since the manifestSignature (hash) will fail to match
// delete the draft revision (will fail if the file only have a draft but no active revisions)
if link.State == proton.LinkStateDraft {
// delete the link (skipping trash, otherwise it won't work)
err = protonDrive.c.DeleteChildren(ctx, protonDrive.MainShare.ShareID, link.ParentLinkID, linkID)
if err != nil {
return "", false, err
}
return "", true, nil
}
// delete the draft revision
err = protonDrive.c.DeleteRevision(ctx, protonDrive.MainShare.ShareID, linkID, draftRevision[0].ID)
if err != nil {
return "", false, err
}
} else {
// if there is a draft, based on the web behavior, it will ask if the user wants to replace the failed upload attempt
// current behavior, we report an error to not upload the file (conservative)
return "", false, ErrDraftExists
}
}
// create a new revision
newRevision, err := protonDrive.c.CreateRevision(ctx, protonDrive.MainShare.ShareID, linkID)
if err != nil {
return "", false, err
}
return newRevision.ID, false, nil
} else if createFileResp != nil {
return createFileResp.RevisionID, false, nil
} else {
// should not happen anymore, since the file search will include the draft now
return "", false, ErrInternalErrorOnFileUpload
}
}
func (protonDrive *ProtonDrive) createFileUploadDraft(ctx context.Context, parentLink *proton.Link, filename string, modTime time.Time, mimeType string) (string, string, *crypto.SessionKey, *crypto.KeyRing, error) {
parentNodeKR, err := protonDrive.getLinkKR(ctx, parentLink)
if err != nil {
return "", "", nil, nil, err
}
newNodeKey, newNodePassphraseEnc, newNodePassphraseSignature, err := generateNodeKeys(parentNodeKR, protonDrive.AddrKR)
if err != nil {
return "", "", nil, nil, err
}
createFileReq := proton.CreateFileReq{
ParentLinkID: parentLink.LinkID,
// Name string // Encrypted File Name
// Hash string // Encrypted File Name hash
MIMEType: mimeType, // MIME Type
// ContentKeyPacket string // The block's key packet, encrypted with the node key.
// ContentKeyPacketSignature string // Unencrypted signature of the content session key, signed with the NodeKey
NodeKey: newNodeKey, // The private NodeKey, used to decrypt any file/folder content.
NodePassphrase: newNodePassphraseEnc, // The passphrase used to unlock the NodeKey, encrypted by the owning Link/Share keyring.
NodePassphraseSignature: newNodePassphraseSignature, // The signature of the NodePassphrase
SignatureAddress: protonDrive.signatureAddress, // Signature email address used to sign passphrase and name
}
/* Name is encrypted using the parent's keyring, and signed with address key */
err = createFileReq.SetName(filename, protonDrive.AddrKR, parentNodeKR)
if err != nil {
return "", "", nil, nil, err
}
parentHashKey, err := parentLink.GetHashKey(parentNodeKR)
if err != nil {
return "", "", nil, nil, err
}
newNodeKR, err := getKeyRing(parentNodeKR, protonDrive.AddrKR, newNodeKey, newNodePassphraseEnc, newNodePassphraseSignature)
if err != nil {
return "", "", nil, nil, err
}
err = createFileReq.SetHash(filename, parentHashKey)
if err != nil {
return "", "", nil, nil, err
}
newSessionKey, err := createFileReq.SetContentKeyPacketAndSignature(newNodeKR)
if err != nil {
return "", "", nil, nil, err
}
createFileAction := func() (*proton.CreateFileRes, *proton.Link, error) {
createFileResp, err := protonDrive.c.CreateFile(ctx, protonDrive.MainShare.ShareID, createFileReq)
if err != nil {
// FIXME: check for duplicated filename by relying on checkAvailableHashes -> able to retrieve linkID too
// Also saving generating resources such as new nodeKR, etc.
if err != proton.ErrFileNameExist {
// other real error caught
return nil, nil, err
}
// search for the link within this folder which has an active/draft revision as we have a file creation conflict
link, err := protonDrive.SearchByNameInActiveFolder(ctx, parentLink, filename, true, false, proton.LinkStateActive)
if err != nil {
return nil, nil, err
}
if link == nil {
link, err = protonDrive.SearchByNameInActiveFolder(ctx, parentLink, filename, true, false, proton.LinkStateDraft)
if err != nil {
return nil, nil, err
}
if link == nil {
// we have a real problem here (unless the assumption is wrong)
// since we can't create a new file AND we can't locate a file with active/draft revision in it
return nil, nil, ErrCantLocateRevision
}
}
return nil, link, nil
}
return &createFileResp, nil, nil
}
createFileResp, link, err := createFileAction()
if err != nil {
return "", "", nil, nil, err
}
revisionID, shouldSubmitCreateFileRequestAgain, err := protonDrive.handleRevisionConflict(ctx, link, createFileResp)
if err != nil {
return "", "", nil, nil, err
}
if shouldSubmitCreateFileRequestAgain {
// the case where the link has only a draft but no active revision
// we need to delete the link and recreate one
createFileResp, link, err = createFileAction()
if err != nil {
return "", "", nil, nil, err
}
revisionID, _, err = protonDrive.handleRevisionConflict(ctx, link, createFileResp)
if err != nil {
return "", "", nil, nil, err
}
}
linkID := ""
if link != nil {
linkID = link.LinkID
// get original newSessionKey and newNodeKR
parentNodeKR, err = protonDrive.getLinkKRByID(ctx, link.ParentLinkID)
if err != nil {
return "", "", nil, nil, err
}
newNodeKR, err = link.GetKeyRing(parentNodeKR, protonDrive.AddrKR)
if err != nil {
return "", "", nil, nil, err
}
newSessionKey, err = link.GetSessionKey(newNodeKR)
if err != nil {
return "", "", nil, nil, err
}
} else {
linkID = createFileResp.ID
}
return linkID, revisionID, newSessionKey, newNodeKR, nil
}
func (protonDrive *ProtonDrive) uploadAndCollectBlockData(ctx context.Context, newSessionKey *crypto.SessionKey, newNodeKR *crypto.KeyRing, file io.Reader, linkID, revisionID string) ([]byte, int64, []int64, string, error) {
type PendingUploadBlocks struct {
blockUploadInfo proton.BlockUploadInfo
encData []byte
}
if newSessionKey == nil || newNodeKR == nil {
return nil, 0, nil, "", ErrMissingInputUploadAndCollectBlockData
}
totalFileSize := int64(0)
pendingUploadBlocks := make([]PendingUploadBlocks, 0)
manifestSignatureData := make([]byte, 0)
uploadPendingBlocks := func() error {
if len(pendingUploadBlocks) == 0 {
return nil
}
blockList := make([]proton.BlockUploadInfo, 0)
for i := range pendingUploadBlocks {
blockList = append(blockList, pendingUploadBlocks[i].blockUploadInfo)
}
blockUploadReq := proton.BlockUploadReq{
AddressID: protonDrive.MainShare.AddressID,
ShareID: protonDrive.MainShare.ShareID,
LinkID: linkID,
RevisionID: revisionID,
BlockList: blockList,
}
blockUploadResp, err := protonDrive.c.RequestBlockUpload(ctx, blockUploadReq)
if err != nil {
return err
}
errChan := make(chan error)
uploadBlockWrapper := func(ctx context.Context, errChan chan error, bareURL, token string, block io.Reader) {
// log.Println("Before semaphore")
if err := protonDrive.blockUploadSemaphore.Acquire(ctx, 1); err != nil {
errChan <- err
}
defer protonDrive.blockUploadSemaphore.Release(1)
// log.Println("After semaphore")
// defer log.Println("Release semaphore")
errChan <- protonDrive.c.UploadBlock(ctx, bareURL, token, block)
}
for i := range blockUploadResp {
go uploadBlockWrapper(ctx, errChan, blockUploadResp[i].BareURL, blockUploadResp[i].Token, bytes.NewReader(pendingUploadBlocks[i].encData))
}
for i := 0; i < len(blockUploadResp); i++ {
err := <-errChan
if err != nil {
return err
}
}
pendingUploadBlocks = pendingUploadBlocks[:0]
return nil
}
shouldContinue := true
sha1Digests := sha1.New()
blockSizes := make([]int64, 0)
for i := 1; shouldContinue; i++ {
if (i-1) > 0 && (i-1)%UPLOAD_BATCH_BLOCK_SIZE == 0 {
err := uploadPendingBlocks()
if err != nil {
return nil, 0, nil, "", err
}
}
// read at most data of size UPLOAD_BLOCK_SIZE
// for some reason, .Read might not actually read up to buffer size -> use io.ReadFull
data := make([]byte, UPLOAD_BLOCK_SIZE) // FIXME: get block size from the server config instead of hardcoding it
readBytes, err := io.ReadFull(file, data)
if err != nil {
if err == io.EOF || err == io.ErrUnexpectedEOF {
// might still have data to read!
if readBytes == 0 {
break
}
shouldContinue = false
} else {
// all other errors
return nil, 0, nil, "", err
}
}
data = data[:readBytes]
totalFileSize += int64(readBytes)
sha1Digests.Write(data)
blockSizes = append(blockSizes, int64(readBytes))
// encrypt data
dataPlainMessage := crypto.NewPlainMessage(data)
encData, err := newSessionKey.Encrypt(dataPlainMessage)
if err != nil {
return nil, 0, nil, "", err
}
encSignature, err := protonDrive.AddrKR.SignDetachedEncrypted(dataPlainMessage, newNodeKR)
if err != nil {
return nil, 0, nil, "", err
}
encSignatureStr, err := encSignature.GetArmored()
if err != nil {
return nil, 0, nil, "", err
}
h := sha256.New()
h.Write(encData)
hash := h.Sum(nil)
base64Hash := base64.StdEncoding.EncodeToString(hash)
if err != nil {
return nil, 0, nil, "", err
}
manifestSignatureData = append(manifestSignatureData, hash...)
pendingUploadBlocks = append(pendingUploadBlocks, PendingUploadBlocks{
blockUploadInfo: proton.BlockUploadInfo{
Index: i, // iOS drive: BE starts with 1
Size: int64(len(encData)),
EncSignature: encSignatureStr,
Hash: base64Hash,
},
encData: encData,
})
}
err := uploadPendingBlocks()
if err != nil {
return nil, 0, nil, "", err
}
sha1Hash := sha1Digests.Sum(nil)
sha1String := hex.EncodeToString(sha1Hash)
return manifestSignatureData, totalFileSize, blockSizes, sha1String, nil
}
func (protonDrive *ProtonDrive) commitNewRevision(ctx context.Context, nodeKR *crypto.KeyRing, xAttrCommon *proton.RevisionXAttrCommon, manifestSignatureData []byte, linkID, revisionID string) error {
manifestSignature, err := protonDrive.AddrKR.SignDetached(crypto.NewPlainMessage(manifestSignatureData))
if err != nil {
return err
}
manifestSignatureString, err := manifestSignature.GetArmored()
if err != nil {
return err
}
commitRevisionReq := proton.CommitRevisionReq{
ManifestSignature: manifestSignatureString,
SignatureAddress: protonDrive.signatureAddress,
}
err = commitRevisionReq.SetEncXAttrString(protonDrive.AddrKR, nodeKR, xAttrCommon)
if err != nil {
return err
}
err = protonDrive.c.CommitRevision(ctx, protonDrive.MainShare.ShareID, linkID, revisionID, commitRevisionReq)
if err != nil {
return err
}
return nil
}
// testParam is for integration test only
// 0 = normal mode
// 1 = up to create revision
// 2 = up to block upload
func (protonDrive *ProtonDrive) uploadFile(ctx context.Context, parentLink *proton.Link, filename string, modTime time.Time, file io.Reader, testParam int) (string, *proton.RevisionXAttrCommon, error) {
// TODO: if we should use github.com/gabriel-vasile/mimetype to detect the MIME type from the file content itself
// Note: this approach might cause the upload progress to display the "fake" progress, since we read in all the content all-at-once
// mimetype.SetLimit(0)
// mType := mimetype.Detect(fileContent)
// mimeType := mType.String()
// detect MIME type by looking at the filename only
mimeType := mime.TypeByExtension(filepath.Ext(filename))
if mimeType == "" {
// api requires a mime type passed in
mimeType = "text/plain"
}
/* step 1: create a draft */
linkID, revisionID, newSessionKey, newNodeKR, err := protonDrive.createFileUploadDraft(ctx, parentLink, filename, modTime, mimeType)
if err != nil {
return "", nil, err
}
if testParam == 1 {
return "", nil, nil
}
/* step 2: upload blocks and collect block data */
manifestSignature, fileSize, blockSizes, digests, err := protonDrive.uploadAndCollectBlockData(ctx, newSessionKey, newNodeKR, file, linkID, revisionID)
if err != nil {
return "", nil, err
}
if testParam == 2 {
// for integration tests
// we try to simulate blocks uploaded but not yet commited
return "", nil, nil
}
/* step 3: mark the file as active by commiting the revision */
xAttrCommon := &proton.RevisionXAttrCommon{
ModificationTime: modTime.Format("2006-01-02T15:04:05-0700"), /* ISO8601 */
Size: fileSize,
BlockSizes: blockSizes,
Digests: map[string]string{
"SHA1": digests,
},
}
err = protonDrive.commitNewRevision(ctx, newNodeKR, xAttrCommon, manifestSignature, linkID, revisionID)
if err != nil {
return "", nil, err
}
return linkID, xAttrCommon, nil
}
/*
There is a route that proton-go-api doesn't have - checkAvailableHashes.
This is used to quickly find the next available filename when the originally supplied filename is taken in the current folder.
Based on the code below, which is taken from the Proton iOS Drive app, we can infer that:
- when a file is to be uploaded && there is filename conflict after the first upload:
- on web, user will be prompted with a) overwrite b) keep both by appending filename with iteration number c) do nothing
- on the iOS client logic, we can see that when the filename conflict happens (after the upload attampt failed)
- the filename will be hashed by using filename + iteration
- 10 iterations will be done per batch, each iteration's hash will be sent to the server
- the server will return available hashes, and the client will take the lowest iteration as the filename to be used
- will be used to search for the next available filename (using hashes avoids the filename being known to the server)
*/

171
file_download.go Normal file
View File

@@ -0,0 +1,171 @@
package proton_api_bridge
import (
"bytes"
"context"
"io"
"log"
"github.com/ProtonMail/gopenpgp/v2/crypto"
"github.com/henrybear327/go-proton-api"
)
type FileDownloadReader struct {
protonDrive *ProtonDrive
ctx context.Context
data *bytes.Buffer
nodeKR *crypto.KeyRing
sessionKey *crypto.SessionKey
revision *proton.Revision
nextRevision int
isEOF bool
// TODO: integrity check if the entire file is read
}
func (r *FileDownloadReader) Read(p []byte) (int, error) {
if r.data.Len() == 0 {
// TODO: do we have memory sharing bug?
// to avoid sharing the underlying buffer array across re-population
r.data = bytes.NewBuffer(nil)
// we download and decrypt more content
err := r.populateBufferOnRead()
if err != nil {
return 0, err
}
if r.isEOF {
// if the file has been downloaded entirely, we return EOF
return 0, io.EOF
}
}
return r.data.Read(p)
}
func (r *FileDownloadReader) Close() error {
r.protonDrive = nil
return nil
}
func (reader *FileDownloadReader) populateBufferOnRead() error {
if len(reader.revision.Blocks) == 0 || len(reader.revision.Blocks) == reader.nextRevision {
reader.isEOF = true
return nil
}
offset := reader.nextRevision
for i := offset; i-offset < DOWNLOAD_BATCH_BLOCK_SIZE && i < len(reader.revision.Blocks); i++ {
// TODO: parallel download
blockReader, err := reader.protonDrive.c.GetBlock(reader.ctx, reader.revision.Blocks[i].BareURL, reader.revision.Blocks[i].Token)
if err != nil {
return err
}
defer blockReader.Close()
err = decryptBlockIntoBuffer(reader.sessionKey, reader.protonDrive.AddrKR, reader.nodeKR, reader.revision.Blocks[i].Hash, reader.revision.Blocks[i].EncSignature, reader.data, blockReader)
if err != nil {
return err
}
reader.nextRevision = i + 1
}
return nil
}
func (protonDrive *ProtonDrive) DownloadFileByID(ctx context.Context, linkID string, offset int64) (io.ReadCloser, int64, *FileSystemAttrs, error) {
/* It's like event system, we need to get the latest information before creating the move request! */
protonDrive.removeLinkIDFromCache(linkID, false)
link, err := protonDrive.getLink(ctx, linkID)
if err != nil {
return nil, 0, nil, err
}
return protonDrive.DownloadFile(ctx, link, offset)
}
func (protonDrive *ProtonDrive) DownloadFile(ctx context.Context, link *proton.Link, offset int64) (io.ReadCloser, int64, *FileSystemAttrs, error) {
if link.Type != proton.LinkTypeFile {
return nil, 0, nil, ErrLinkTypeMustToBeFileType
}
parentNodeKR, err := protonDrive.getLinkKRByID(ctx, link.ParentLinkID)
if err != nil {
return nil, 0, nil, err
}
nodeKR, err := link.GetKeyRing(parentNodeKR, protonDrive.AddrKR)
if err != nil {
return nil, 0, nil, err
}
sessionKey, err := link.GetSessionKey(nodeKR)
if err != nil {
return nil, 0, nil, err
}
revision, fileSystemAttrs, err := protonDrive.GetActiveRevisionWithAttrs(ctx, link)
if err != nil {
return nil, 0, nil, err
}
reader := &FileDownloadReader{
protonDrive: protonDrive,
ctx: ctx,
data: bytes.NewBuffer(nil),
nodeKR: nodeKR,
sessionKey: sessionKey,
revision: revision,
nextRevision: 0,
isEOF: false,
}
useFallbackDownload := false
if fileSystemAttrs != nil {
// based on offset, infer the nextRevision (0-based)
if fileSystemAttrs.BlockSizes == nil {
useFallbackDownload = true
} else {
// infer nextRevision
totalBytes := int64(0)
for i := 0; i < len(fileSystemAttrs.BlockSizes); i++ {
prevTotalBytes := totalBytes
totalBytes += fileSystemAttrs.BlockSizes[i]
if offset <= totalBytes {
offset = offset - prevTotalBytes
reader.nextRevision = i
break
}
}
// download will start from the specified block
n, err := io.CopyN(io.Discard, reader, offset)
if err != nil {
return nil, 0, nil, err
}
if int64(n) != offset {
return nil, 0, nil, ErrSeekOffsetAfterSkippingBlocks
}
}
}
if useFallbackDownload {
log.Println("Performing inefficient seek as metadata of encrypted file is missing")
n, err := io.CopyN(io.Discard, reader, offset)
if err != nil {
return nil, 0, nil, err
}
if int64(n) != offset {
return nil, 0, nil, ErrSeekOffsetAfterSkippingBlocks
}
}
return reader, link.Size, fileSystemAttrs, nil
}

479
file_upload.go Normal file
View File

@@ -0,0 +1,479 @@
package proton_api_bridge
import (
"bufio"
"bytes"
"context"
"crypto/sha1"
"crypto/sha256"
"encoding/base64"
"encoding/hex"
"io"
"mime"
"os"
"path/filepath"
"time"
"github.com/ProtonMail/gopenpgp/v2/crypto"
"github.com/henrybear327/go-proton-api"
)
func (protonDrive *ProtonDrive) handleRevisionConflict(ctx context.Context, link *proton.Link, createFileResp *proton.CreateFileRes) (string, bool, error) {
if link != nil {
linkID := link.LinkID
draftRevision, err := protonDrive.GetRevisions(ctx, link, proton.RevisionStateDraft)
if err != nil {
return "", false, err
}
// if we have a draft revision, depending on the user config, we can abort the upload or recreate a draft
// if we have no draft revision, then we can create a new draft revision directly (there is a restriction of 1 draft revision per file)
if len(draftRevision) > 0 {
// TODO: maintain clientUID to mark that this is our own draft (which can indicate failed upload attempt!)
if protonDrive.Config.ReplaceExistingDraft {
// Question: how do we observe for file upload cancellation -> clientUID?
// Random thoughts: if there are concurrent modification to the draft, the server should be able to catch this when commiting the revision
// since the manifestSignature (hash) will fail to match
// delete the draft revision (will fail if the file only have a draft but no active revisions)
if link.State == proton.LinkStateDraft {
// delete the link (skipping trash, otherwise it won't work)
err = protonDrive.c.DeleteChildren(ctx, protonDrive.MainShare.ShareID, link.ParentLinkID, linkID)
if err != nil {
return "", false, err
}
return "", true, nil
}
// delete the draft revision
err = protonDrive.c.DeleteRevision(ctx, protonDrive.MainShare.ShareID, linkID, draftRevision[0].ID)
if err != nil {
return "", false, err
}
} else {
// if there is a draft, based on the web behavior, it will ask if the user wants to replace the failed upload attempt
// current behavior, we report an error to not upload the file (conservative)
return "", false, ErrDraftExists
}
}
// create a new revision
newRevision, err := protonDrive.c.CreateRevision(ctx, protonDrive.MainShare.ShareID, linkID)
if err != nil {
return "", false, err
}
return newRevision.ID, false, nil
} else if createFileResp != nil {
return createFileResp.RevisionID, false, nil
} else {
// should not happen anymore, since the file search will include the draft now
return "", false, ErrInternalErrorOnFileUpload
}
}
func (protonDrive *ProtonDrive) createFileUploadDraft(ctx context.Context, parentLink *proton.Link, filename string, modTime time.Time, mimeType string) (string, string, *crypto.SessionKey, *crypto.KeyRing, error) {
parentNodeKR, err := protonDrive.getLinkKR(ctx, parentLink)
if err != nil {
return "", "", nil, nil, err
}
newNodeKey, newNodePassphraseEnc, newNodePassphraseSignature, err := generateNodeKeys(parentNodeKR, protonDrive.AddrKR)
if err != nil {
return "", "", nil, nil, err
}
createFileReq := proton.CreateFileReq{
ParentLinkID: parentLink.LinkID,
// Name string // Encrypted File Name
// Hash string // Encrypted File Name hash
MIMEType: mimeType, // MIME Type
// ContentKeyPacket string // The block's key packet, encrypted with the node key.
// ContentKeyPacketSignature string // Unencrypted signature of the content session key, signed with the NodeKey
NodeKey: newNodeKey, // The private NodeKey, used to decrypt any file/folder content.
NodePassphrase: newNodePassphraseEnc, // The passphrase used to unlock the NodeKey, encrypted by the owning Link/Share keyring.
NodePassphraseSignature: newNodePassphraseSignature, // The signature of the NodePassphrase
SignatureAddress: protonDrive.signatureAddress, // Signature email address used to sign passphrase and name
}
/* Name is encrypted using the parent's keyring, and signed with address key */
err = createFileReq.SetName(filename, protonDrive.AddrKR, parentNodeKR)
if err != nil {
return "", "", nil, nil, err
}
parentHashKey, err := parentLink.GetHashKey(parentNodeKR)
if err != nil {
return "", "", nil, nil, err
}
newNodeKR, err := getKeyRing(parentNodeKR, protonDrive.AddrKR, newNodeKey, newNodePassphraseEnc, newNodePassphraseSignature)
if err != nil {
return "", "", nil, nil, err
}
err = createFileReq.SetHash(filename, parentHashKey)
if err != nil {
return "", "", nil, nil, err
}
newSessionKey, err := createFileReq.SetContentKeyPacketAndSignature(newNodeKR)
if err != nil {
return "", "", nil, nil, err
}
createFileAction := func() (*proton.CreateFileRes, *proton.Link, error) {
createFileResp, err := protonDrive.c.CreateFile(ctx, protonDrive.MainShare.ShareID, createFileReq)
if err != nil {
// FIXME: check for duplicated filename by relying on checkAvailableHashes -> able to retrieve linkID too
// Also saving generating resources such as new nodeKR, etc.
if err != proton.ErrFileNameExist {
// other real error caught
return nil, nil, err
}
// search for the link within this folder which has an active/draft revision as we have a file creation conflict
link, err := protonDrive.SearchByNameInActiveFolder(ctx, parentLink, filename, true, false, proton.LinkStateActive)
if err != nil {
return nil, nil, err
}
if link == nil {
link, err = protonDrive.SearchByNameInActiveFolder(ctx, parentLink, filename, true, false, proton.LinkStateDraft)
if err != nil {
return nil, nil, err
}
if link == nil {
// we have a real problem here (unless the assumption is wrong)
// since we can't create a new file AND we can't locate a file with active/draft revision in it
return nil, nil, ErrCantLocateRevision
}
}
return nil, link, nil
}
return &createFileResp, nil, nil
}
createFileResp, link, err := createFileAction()
if err != nil {
return "", "", nil, nil, err
}
revisionID, shouldSubmitCreateFileRequestAgain, err := protonDrive.handleRevisionConflict(ctx, link, createFileResp)
if err != nil {
return "", "", nil, nil, err
}
if shouldSubmitCreateFileRequestAgain {
// the case where the link has only a draft but no active revision
// we need to delete the link and recreate one
createFileResp, link, err = createFileAction()
if err != nil {
return "", "", nil, nil, err
}
revisionID, _, err = protonDrive.handleRevisionConflict(ctx, link, createFileResp)
if err != nil {
return "", "", nil, nil, err
}
}
linkID := ""
if link != nil {
linkID = link.LinkID
// get original newSessionKey and newNodeKR
parentNodeKR, err = protonDrive.getLinkKRByID(ctx, link.ParentLinkID)
if err != nil {
return "", "", nil, nil, err
}
newNodeKR, err = link.GetKeyRing(parentNodeKR, protonDrive.AddrKR)
if err != nil {
return "", "", nil, nil, err
}
newSessionKey, err = link.GetSessionKey(newNodeKR)
if err != nil {
return "", "", nil, nil, err
}
} else {
linkID = createFileResp.ID
}
return linkID, revisionID, newSessionKey, newNodeKR, nil
}
func (protonDrive *ProtonDrive) uploadAndCollectBlockData(ctx context.Context, newSessionKey *crypto.SessionKey, newNodeKR *crypto.KeyRing, file io.Reader, linkID, revisionID string) ([]byte, int64, []int64, string, error) {
type PendingUploadBlocks struct {
blockUploadInfo proton.BlockUploadInfo
encData []byte
}
if newSessionKey == nil || newNodeKR == nil {
return nil, 0, nil, "", ErrMissingInputUploadAndCollectBlockData
}
totalFileSize := int64(0)
pendingUploadBlocks := make([]PendingUploadBlocks, 0)
manifestSignatureData := make([]byte, 0)
uploadPendingBlocks := func() error {
if len(pendingUploadBlocks) == 0 {
return nil
}
blockList := make([]proton.BlockUploadInfo, 0)
for i := range pendingUploadBlocks {
blockList = append(blockList, pendingUploadBlocks[i].blockUploadInfo)
}
blockUploadReq := proton.BlockUploadReq{
AddressID: protonDrive.MainShare.AddressID,
ShareID: protonDrive.MainShare.ShareID,
LinkID: linkID,
RevisionID: revisionID,
BlockList: blockList,
}
blockUploadResp, err := protonDrive.c.RequestBlockUpload(ctx, blockUploadReq)
if err != nil {
return err
}
errChan := make(chan error)
uploadBlockWrapper := func(ctx context.Context, errChan chan error, bareURL, token string, block io.Reader) {
// log.Println("Before semaphore")
if err := protonDrive.blockUploadSemaphore.Acquire(ctx, 1); err != nil {
errChan <- err
}
defer protonDrive.blockUploadSemaphore.Release(1)
// log.Println("After semaphore")
// defer log.Println("Release semaphore")
errChan <- protonDrive.c.UploadBlock(ctx, bareURL, token, block)
}
for i := range blockUploadResp {
go uploadBlockWrapper(ctx, errChan, blockUploadResp[i].BareURL, blockUploadResp[i].Token, bytes.NewReader(pendingUploadBlocks[i].encData))
}
for i := 0; i < len(blockUploadResp); i++ {
err := <-errChan
if err != nil {
return err
}
}
pendingUploadBlocks = pendingUploadBlocks[:0]
return nil
}
shouldContinue := true
sha1Digests := sha1.New()
blockSizes := make([]int64, 0)
for i := 1; shouldContinue; i++ {
if (i-1) > 0 && (i-1)%UPLOAD_BATCH_BLOCK_SIZE == 0 {
err := uploadPendingBlocks()
if err != nil {
return nil, 0, nil, "", err
}
}
// read at most data of size UPLOAD_BLOCK_SIZE
// for some reason, .Read might not actually read up to buffer size -> use io.ReadFull
data := make([]byte, UPLOAD_BLOCK_SIZE) // FIXME: get block size from the server config instead of hardcoding it
readBytes, err := io.ReadFull(file, data)
if err != nil {
if err == io.EOF || err == io.ErrUnexpectedEOF {
// might still have data to read!
if readBytes == 0 {
break
}
shouldContinue = false
} else {
// all other errors
return nil, 0, nil, "", err
}
}
data = data[:readBytes]
totalFileSize += int64(readBytes)
sha1Digests.Write(data)
blockSizes = append(blockSizes, int64(readBytes))
// encrypt data
dataPlainMessage := crypto.NewPlainMessage(data)
encData, err := newSessionKey.Encrypt(dataPlainMessage)
if err != nil {
return nil, 0, nil, "", err
}
encSignature, err := protonDrive.AddrKR.SignDetachedEncrypted(dataPlainMessage, newNodeKR)
if err != nil {
return nil, 0, nil, "", err
}
encSignatureStr, err := encSignature.GetArmored()
if err != nil {
return nil, 0, nil, "", err
}
h := sha256.New()
h.Write(encData)
hash := h.Sum(nil)
base64Hash := base64.StdEncoding.EncodeToString(hash)
if err != nil {
return nil, 0, nil, "", err
}
manifestSignatureData = append(manifestSignatureData, hash...)
pendingUploadBlocks = append(pendingUploadBlocks, PendingUploadBlocks{
blockUploadInfo: proton.BlockUploadInfo{
Index: i, // iOS drive: BE starts with 1
Size: int64(len(encData)),
EncSignature: encSignatureStr,
Hash: base64Hash,
},
encData: encData,
})
}
err := uploadPendingBlocks()
if err != nil {
return nil, 0, nil, "", err
}
sha1Hash := sha1Digests.Sum(nil)
sha1String := hex.EncodeToString(sha1Hash)
return manifestSignatureData, totalFileSize, blockSizes, sha1String, nil
}
func (protonDrive *ProtonDrive) commitNewRevision(ctx context.Context, nodeKR *crypto.KeyRing, xAttrCommon *proton.RevisionXAttrCommon, manifestSignatureData []byte, linkID, revisionID string) error {
manifestSignature, err := protonDrive.AddrKR.SignDetached(crypto.NewPlainMessage(manifestSignatureData))
if err != nil {
return err
}
manifestSignatureString, err := manifestSignature.GetArmored()
if err != nil {
return err
}
commitRevisionReq := proton.CommitRevisionReq{
ManifestSignature: manifestSignatureString,
SignatureAddress: protonDrive.signatureAddress,
}
err = commitRevisionReq.SetEncXAttrString(protonDrive.AddrKR, nodeKR, xAttrCommon)
if err != nil {
return err
}
err = protonDrive.c.CommitRevision(ctx, protonDrive.MainShare.ShareID, linkID, revisionID, commitRevisionReq)
if err != nil {
return err
}
return nil
}
// testParam is for integration test only
// 0 = normal mode
// 1 = up to create revision
// 2 = up to block upload
func (protonDrive *ProtonDrive) uploadFile(ctx context.Context, parentLink *proton.Link, filename string, modTime time.Time, file io.Reader, testParam int) (string, *proton.RevisionXAttrCommon, error) {
// TODO: if we should use github.com/gabriel-vasile/mimetype to detect the MIME type from the file content itself
// Note: this approach might cause the upload progress to display the "fake" progress, since we read in all the content all-at-once
// mimetype.SetLimit(0)
// mType := mimetype.Detect(fileContent)
// mimeType := mType.String()
// detect MIME type by looking at the filename only
mimeType := mime.TypeByExtension(filepath.Ext(filename))
if mimeType == "" {
// api requires a mime type passed in
mimeType = "text/plain"
}
/* step 1: create a draft */
linkID, revisionID, newSessionKey, newNodeKR, err := protonDrive.createFileUploadDraft(ctx, parentLink, filename, modTime, mimeType)
if err != nil {
return "", nil, err
}
if testParam == 1 {
return "", nil, nil
}
/* step 2: upload blocks and collect block data */
manifestSignature, fileSize, blockSizes, digests, err := protonDrive.uploadAndCollectBlockData(ctx, newSessionKey, newNodeKR, file, linkID, revisionID)
if err != nil {
return "", nil, err
}
if testParam == 2 {
// for integration tests
// we try to simulate blocks uploaded but not yet commited
return "", nil, nil
}
/* step 3: mark the file as active by commiting the revision */
xAttrCommon := &proton.RevisionXAttrCommon{
ModificationTime: modTime.Format("2006-01-02T15:04:05-0700"), /* ISO8601 */
Size: fileSize,
BlockSizes: blockSizes,
Digests: map[string]string{
"SHA1": digests,
},
}
err = protonDrive.commitNewRevision(ctx, newNodeKR, xAttrCommon, manifestSignature, linkID, revisionID)
if err != nil {
return "", nil, err
}
return linkID, xAttrCommon, nil
}
func (protonDrive *ProtonDrive) UploadFileByReader(ctx context.Context, parentLinkID string, filename string, modTime time.Time, file io.Reader, testParam int) (string, *proton.RevisionXAttrCommon, error) {
parentLink, err := protonDrive.getLink(ctx, parentLinkID)
if err != nil {
return "", nil, err
}
return protonDrive.uploadFile(ctx, parentLink, filename, modTime, file, testParam)
}
func (protonDrive *ProtonDrive) UploadFileByPath(ctx context.Context, parentLink *proton.Link, filename string, filePath string, testParam int) (string, *proton.RevisionXAttrCommon, error) {
f, err := os.Open(filePath)
if err != nil {
return "", nil, err
}
defer f.Close()
info, err := os.Stat(filePath)
if err != nil {
return "", nil, err
}
in := bufio.NewReader(f)
return protonDrive.uploadFile(ctx, parentLink, filename, info.ModTime(), in, testParam)
}
/*
There is a route that proton-go-api doesn't have - checkAvailableHashes.
This is used to quickly find the next available filename when the originally supplied filename is taken in the current folder.
Based on the code below, which is taken from the Proton iOS Drive app, we can infer that:
- when a file is to be uploaded && there is filename conflict after the first upload:
- on web, user will be prompted with a) overwrite b) keep both by appending filename with iteration number c) do nothing
- on the iOS client logic, we can see that when the filename conflict happens (after the upload attampt failed)
- the filename will be hashed by using filename + iteration
- 10 iterations will be done per batch, each iteration's hash will be sent to the server
- the server will return available hashes, and the client will take the lowest iteration as the filename to be used
- will be used to search for the next available filename (using hashes avoids the filename being known to the server)
*/

102
folder.go
View File

@@ -2,12 +2,8 @@ package proton_api_bridge
import (
"context"
"io"
"log"
"os"
"time"
"github.com/ProtonMail/gopenpgp/v2/crypto"
"github.com/henrybear327/go-proton-api"
)
@@ -64,102 +60,6 @@ func (protonDrive *ProtonDrive) ListDirectory(
return ret, nil
}
func (protonDrive *ProtonDrive) ListDirectoriesRecursively(
ctx context.Context,
parentNodeKR *crypto.KeyRing,
link *proton.Link,
download bool,
maxDepth, curDepth /* 0-based */ int,
excludeRoot bool,
pathSoFar string,
paths *[]string) error {
/*
Assumptions:
- we only care about the active ones
*/
if link.State != proton.LinkStateActive {
return nil
}
// log.Println("curDepth", curDepth, "pathSoFar", pathSoFar)
var currentPath = ""
if !(excludeRoot && curDepth == 0) {
name, err := link.GetName(parentNodeKR, protonDrive.AddrKR)
if err != nil {
return err
}
currentPath = pathSoFar + "/" + name
// log.Println("currentPath", currentPath)
if paths != nil {
*paths = append(*paths, currentPath)
}
}
if download {
if protonDrive.Config.DataFolderName == "" {
return ErrDataFolderNameIsEmpty
}
if link.Type == proton.LinkTypeFile {
log.Println("Downloading", currentPath)
defer log.Println("Completes downloading", currentPath)
reader, _, _, err := protonDrive.DownloadFile(ctx, link, 0)
if err != nil {
return err
}
byteArray, err := io.ReadAll(reader)
if err != nil {
return err
}
err = os.WriteFile("./"+protonDrive.Config.DataFolderName+"/"+currentPath, byteArray, 0777)
if err != nil {
return err
}
} else /* folder */ {
if !(excludeRoot && curDepth == 0) {
// log.Println("Creating folder", currentPath)
// defer log.Println("Completes creating folder", currentPath)
err := os.Mkdir("./"+protonDrive.Config.DataFolderName+"/"+currentPath, 0777)
if err != nil {
return err
}
}
}
}
if maxDepth == -1 || curDepth < maxDepth {
if link.Type == proton.LinkTypeFolder {
childrenLinks, err := protonDrive.c.ListChildren(ctx, protonDrive.MainShare.ShareID, link.LinkID, true)
if err != nil {
return err
}
// log.Printf("childrenLinks len = %v, %#v", len(childrenLinks), childrenLinks)
if childrenLinks != nil {
// get current node's keyring
linkKR, err := link.GetKeyRing(parentNodeKR, protonDrive.AddrKR)
if err != nil {
return err
}
for _, childLink := range childrenLinks {
err = protonDrive.ListDirectoriesRecursively(ctx, linkKR, &childLink, download, maxDepth, curDepth+1, excludeRoot, currentPath, paths)
if err != nil {
return err
}
}
}
}
}
return nil
}
func (protonDrive *ProtonDrive) CreateNewFolderByID(ctx context.Context, parentLinkID string, folderName string) (string, error) {
/* It's like event system, we need to get the latest information before creating the move request! */
protonDrive.removeLinkIDFromCache(parentLinkID, false)
@@ -257,7 +157,7 @@ func (protonDrive *ProtonDrive) MoveFileByID(ctx context.Context, srcLinkID, dst
}
func (protonDrive *ProtonDrive) MoveFile(ctx context.Context, srcLink *proton.Link, dstParentLink *proton.Link, dstName string) error {
return protonDrive.MoveFolder(ctx, srcLink, dstParentLink, dstName)
return protonDrive.moveLink(ctx, srcLink, dstParentLink, dstName)
}
func (protonDrive *ProtonDrive) MoveFolderByID(ctx context.Context, srcLinkID, dstParentLinkID, dstName string) error {

107
folder_recursive.go Normal file
View File

@@ -0,0 +1,107 @@
package proton_api_bridge
import (
"context"
"io"
"log"
"os"
"github.com/ProtonMail/gopenpgp/v2/crypto"
"github.com/henrybear327/go-proton-api"
)
func (protonDrive *ProtonDrive) listDirectoriesRecursively(
ctx context.Context,
parentNodeKR *crypto.KeyRing,
link *proton.Link,
download bool,
maxDepth, curDepth /* 0-based */ int,
excludeRoot bool,
pathSoFar string,
paths *[]string) error {
/*
Assumptions:
- we only care about the active ones
*/
if link.State != proton.LinkStateActive {
return nil
}
// log.Println("curDepth", curDepth, "pathSoFar", pathSoFar)
var currentPath = ""
if !(excludeRoot && curDepth == 0) {
name, err := link.GetName(parentNodeKR, protonDrive.AddrKR)
if err != nil {
return err
}
currentPath = pathSoFar + "/" + name
// log.Println("currentPath", currentPath)
if paths != nil {
*paths = append(*paths, currentPath)
}
}
if download {
if protonDrive.Config.DataFolderName == "" {
return ErrDataFolderNameIsEmpty
}
if link.Type == proton.LinkTypeFile {
log.Println("Downloading", currentPath)
defer log.Println("Completes downloading", currentPath)
reader, _, _, err := protonDrive.DownloadFile(ctx, link, 0)
if err != nil {
return err
}
byteArray, err := io.ReadAll(reader)
if err != nil {
return err
}
err = os.WriteFile("./"+protonDrive.Config.DataFolderName+"/"+currentPath, byteArray, 0777)
if err != nil {
return err
}
} else /* folder */ {
if !(excludeRoot && curDepth == 0) {
// log.Println("Creating folder", currentPath)
// defer log.Println("Completes creating folder", currentPath)
err := os.Mkdir("./"+protonDrive.Config.DataFolderName+"/"+currentPath, 0777)
if err != nil {
return err
}
}
}
}
if maxDepth == -1 || curDepth < maxDepth {
if link.Type == proton.LinkTypeFolder {
childrenLinks, err := protonDrive.c.ListChildren(ctx, protonDrive.MainShare.ShareID, link.LinkID, true)
if err != nil {
return err
}
// log.Printf("childrenLinks len = %v, %#v", len(childrenLinks), childrenLinks)
if childrenLinks != nil {
// get current node's keyring
linkKR, err := link.GetKeyRing(parentNodeKR, protonDrive.AddrKR)
if err != nil {
return err
}
for _, childLink := range childrenLinks {
err = protonDrive.listDirectoriesRecursively(ctx, linkKR, &childLink, download, maxDepth, curDepth+1, excludeRoot, currentPath, paths)
if err != nil {
return err
}
}
}
}
}
return nil
}

8
go.mod
View File

@@ -3,9 +3,9 @@ module github.com/henrybear327/Proton-API-Bridge
go 1.18
require (
github.com/ProtonMail/gluon v0.17.0
github.com/ProtonMail/gluon v0.17.1-0.20230724134000-308be39be96e
github.com/ProtonMail/gopenpgp/v2 v2.7.2
github.com/henrybear327/go-proton-api v0.0.0-20230725203741-316f1b3227a9
github.com/henrybear327/go-proton-api v0.0.0-20230802152927-59db7bb18c64
github.com/relvacode/iso8601 v1.3.0
golang.org/x/sync v0.3.0
)
@@ -28,8 +28,8 @@ require (
github.com/pkg/errors v0.9.1 // indirect
github.com/sirupsen/logrus v1.9.3 // indirect
golang.org/x/crypto v0.11.0 // indirect
golang.org/x/exp v0.0.0-20230725093048-515e97ebf090 // indirect
golang.org/x/net v0.12.0 // indirect
golang.org/x/exp v0.0.0-20230801115018-d63ba01acd4b // indirect
golang.org/x/net v0.13.0 // indirect
golang.org/x/sys v0.10.0 // indirect
golang.org/x/text v0.11.0 // indirect
)

18
go.sum
View File

@@ -2,8 +2,8 @@ github.com/Masterminds/semver/v3 v3.2.0 h1:3MEsd0SM6jqZojhjLWWeBY+Kcjy9i6MQAeY7Y
github.com/ProtonMail/bcrypt v0.0.0-20210511135022-227b4adcab57/go.mod h1:HecWFHognK8GfRDGnFQbW/LiV7A3MX3gZVs45vk5h8I=
github.com/ProtonMail/bcrypt v0.0.0-20211005172633-e235017c1baf h1:yc9daCCYUefEs69zUkSzubzjBbL+cmOXgnmt9Fyd9ug=
github.com/ProtonMail/bcrypt v0.0.0-20211005172633-e235017c1baf/go.mod h1:o0ESU9p83twszAU8LBeJKFAAMX14tISa0yk4Oo5TOqo=
github.com/ProtonMail/gluon v0.17.0 h1:QfMRUcXd47MANHmoerj1ZHXsNzfW9gjsLmF+7Dim5ZU=
github.com/ProtonMail/gluon v0.17.0/go.mod h1:Og5/Dz1MiGpCJn51XujZwxiLG7WzvvjE5PRpZBQmAHo=
github.com/ProtonMail/gluon v0.17.1-0.20230724134000-308be39be96e h1:lCsqUUACrcMC83lg5rTo9Y0PnPItE61JSfvMyIcANwk=
github.com/ProtonMail/gluon v0.17.1-0.20230724134000-308be39be96e/go.mod h1:Og5/Dz1MiGpCJn51XujZwxiLG7WzvvjE5PRpZBQmAHo=
github.com/ProtonMail/go-crypto v0.0.0-20230321155629-9a39f2531310/go.mod h1:8TI4H3IbrackdNgv+92dI+rhpCaLqM0IfpgCgenFvRE=
github.com/ProtonMail/go-crypto v0.0.0-20230717121422-5aa5874ade95 h1:KLq8BE0KwCL+mmXnjLWEAOYO+2l2AE4YMmqG1ZpZHBs=
github.com/ProtonMail/go-crypto v0.0.0-20230717121422-5aa5874ade95/go.mod h1:EjAoLdwvbIOoOQr3ihjnSoLZRtE8azugULFRteWMNc0=
@@ -49,10 +49,8 @@ github.com/go-resty/resty/v2 v2.7.0/go.mod h1:9PWDzw47qPphMRFfhsyk0NnSgvluHcljSM
github.com/goccy/go-json v0.10.2 h1:CrxCmQqYDkv1z7lO7Wbh2HN93uovUHgrECaO5ZrCXAU=
github.com/google/uuid v1.3.0 h1:t6JiXgmwXMjEs8VusXIJk2BXHsn+wx8BZdTaoZ5fu7I=
github.com/google/uuid v1.3.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
github.com/henrybear327/go-proton-api v0.0.0-20230725202642-a15fc6929b34 h1:uAwdnEOpYnAUJ71CQuPTS8dhOqywvaOub8tSHPqmUIU=
github.com/henrybear327/go-proton-api v0.0.0-20230725202642-a15fc6929b34/go.mod h1:l42xBSOrCmkAxzWUHcoUsG/cP8m1hMhV72GoChOX3bg=
github.com/henrybear327/go-proton-api v0.0.0-20230725203741-316f1b3227a9 h1:ITuA1x2yBCE4Wgae+ywZmezVuKOKhJ4nuDtColKE32c=
github.com/henrybear327/go-proton-api v0.0.0-20230725203741-316f1b3227a9/go.mod h1:eHoT/G3lmdatxoY04LSQthG/YnPmu4aebNZe8h+yvBc=
github.com/henrybear327/go-proton-api v0.0.0-20230802152927-59db7bb18c64 h1:s+tcvtvssdVK09u1fSBDk0g6F6fzPz+qDmPg+5kcU3c=
github.com/henrybear327/go-proton-api v0.0.0-20230802152927-59db7bb18c64/go.mod h1:w63MZuzufKcIZ93pwRgiOtxMXYafI8H74D77AxytOBc=
github.com/json-iterator/go v1.1.12 h1:PV8peI4a0ysnczrg+LtxykD8LfKY9ML6u2jnxaEnrnM=
github.com/klauspost/cpuid/v2 v2.2.4 h1:acbojRNwl3o09bUq+yDCtZFc1aiwaAAxtcn8YkZXnvk=
github.com/leodido/go-urn v1.2.4 h1:XlAE/cm/ms7TE/VMVoduSpNBoyc2dOxHs5MZSwAN63Q=
@@ -82,8 +80,8 @@ golang.org/x/crypto v0.3.1-0.20221117191849-2c476679df9a/go.mod h1:hebNnKkNXi2Uz
golang.org/x/crypto v0.7.0/go.mod h1:pYwdfH91IfpZVANVyUOhSIPZaFoJGxTFbZhFTx+dXZU=
golang.org/x/crypto v0.11.0 h1:6Ewdq3tDic1mg5xRO4milcWCfMVQhI4NkqWWvqejpuA=
golang.org/x/crypto v0.11.0/go.mod h1:xgJhtzW8F9jGdVFWZESrid1U1bjeNy4zgy5cRr/CIio=
golang.org/x/exp v0.0.0-20230725093048-515e97ebf090 h1:Di6/M8l0O2lCLc6VVRWhgCiApHV8MnQurBnFSHsQtNY=
golang.org/x/exp v0.0.0-20230725093048-515e97ebf090/go.mod h1:FXUEEKJgO7OQYeo8N01OfiKP8RXMtf6e8aTskBGqWdc=
golang.org/x/exp v0.0.0-20230801115018-d63ba01acd4b h1:r+vk0EmXNmekl0S0BascoeeoHk/L7wmaW2QF90K+kYI=
golang.org/x/exp v0.0.0-20230801115018-d63ba01acd4b/go.mod h1:FXUEEKJgO7OQYeo8N01OfiKP8RXMtf6e8aTskBGqWdc=
golang.org/x/mod v0.6.0-dev.0.20220419223038-86c51ed26bb4/go.mod h1:jJ57K6gSWd91VN4djpZkiMVwK6gcyfeH4XE8wZrZaV4=
golang.org/x/mod v0.8.0/go.mod h1:iBbtSCu2XBx23ZKBPSOrRkjjQPZFPuis4dIYUhu/chs=
golang.org/x/net v0.0.0-20190620200207-3b0461eec859/go.mod h1:z5CRVTTTmAJ677TzLLGU+0bjPO0LkuOLi4/5GtJWs/s=
@@ -96,8 +94,8 @@ golang.org/x/net v0.6.0/go.mod h1:2Tu9+aMcznHK/AK1HMvgo6xiTLG5rD5rZLDS+rp2Bjs=
golang.org/x/net v0.7.0/go.mod h1:2Tu9+aMcznHK/AK1HMvgo6xiTLG5rD5rZLDS+rp2Bjs=
golang.org/x/net v0.8.0/go.mod h1:QVkue5JL9kW//ek3r6jTKnTFis1tRmNAW2P1shuFdJc=
golang.org/x/net v0.9.0/go.mod h1:d48xBJpPfHeWQsugry2m+kC02ZBRGRgulfHnEXEuWns=
golang.org/x/net v0.12.0 h1:cfawfvKITfUsFCeJIHJrbSxpeu/E81khclypR0GVT50=
golang.org/x/net v0.12.0/go.mod h1:zEVYFnQC7m/vmpQFELhcD1EWkZlX69l4oqgmer6hfKA=
golang.org/x/net v0.13.0 h1:Nvo8UFsZ8X3BhAC9699Z1j7XQ3rsZnUUm7jfBEk1ueY=
golang.org/x/net v0.13.0/go.mod h1:zEVYFnQC7m/vmpQFELhcD1EWkZlX69l4oqgmer6hfKA=
golang.org/x/sync v0.0.0-20190423024810-112230192c58/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
golang.org/x/sync v0.0.0-20220722155255-886fb9371eb4/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
golang.org/x/sync v0.1.0/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=

106
search.go
View File

@@ -3,117 +3,13 @@ package proton_api_bridge
import (
"context"
"github.com/ProtonMail/gopenpgp/v2/crypto"
"github.com/henrybear327/go-proton-api"
)
/*
Observation: file name is unique, since it's checked (by hash) on the server
The filename is unique in a given folder, since it's checked (by using hash) on the server
*/
func (protonDrive *ProtonDrive) SearchByNameRecursivelyFromRoot(ctx context.Context, targetName string, isFolder bool, listAllActiveOrDraftFiles bool) (*proton.Link, error) {
var linkType proton.LinkType
if isFolder {
linkType = proton.LinkTypeFolder
} else {
linkType = proton.LinkTypeFile
}
return protonDrive.searchByNameRecursively(ctx, protonDrive.MainShareKR, protonDrive.RootLink, targetName, linkType, listAllActiveOrDraftFiles)
}
func (protonDrive *ProtonDrive) SearchByNameRecursivelyByID(ctx context.Context, folderLinkID string, targetName string, isFolder bool, listAllActiveOrDraftFiles bool) (*proton.Link, error) {
folderLink, err := protonDrive.getLink(ctx, folderLinkID)
if err != nil {
return nil, err
}
var linkType proton.LinkType
if isFolder {
linkType = proton.LinkTypeFolder
} else {
linkType = proton.LinkTypeFile
}
if folderLink.Type != proton.LinkTypeFolder {
return nil, ErrLinkTypeMustToBeFolderType
}
folderKeyRing, err := protonDrive.getLinkKRByID(ctx, folderLink.ParentLinkID)
if err != nil {
return nil, err
}
return protonDrive.searchByNameRecursively(ctx, folderKeyRing, folderLink, targetName, linkType, listAllActiveOrDraftFiles)
}
func (protonDrive *ProtonDrive) SearchByNameRecursively(ctx context.Context, folderLink *proton.Link, targetName string, isFolder bool, listAllActiveOrDraftFiles bool) (*proton.Link, error) {
var linkType proton.LinkType
if isFolder {
linkType = proton.LinkTypeFolder
} else {
linkType = proton.LinkTypeFile
}
if folderLink.Type != proton.LinkTypeFolder {
return nil, ErrLinkTypeMustToBeFolderType
}
folderKeyRing, err := protonDrive.getLinkKRByID(ctx, folderLink.ParentLinkID)
if err != nil {
return nil, err
}
return protonDrive.searchByNameRecursively(ctx, folderKeyRing, folderLink, targetName, linkType, listAllActiveOrDraftFiles)
}
func (protonDrive *ProtonDrive) searchByNameRecursively(
ctx context.Context,
parentNodeKR *crypto.KeyRing,
link *proton.Link,
targetName string,
linkType proton.LinkType,
listAllActiveOrDraftFiles bool) (*proton.Link, error) {
if listAllActiveOrDraftFiles {
if link.State != proton.LinkStateActive && link.State != proton.LinkStateDraft {
return nil, nil
}
} else if link.State != proton.LinkStateActive {
return nil, nil
}
name, err := link.GetName(parentNodeKR, protonDrive.AddrKR)
if err != nil {
return nil, err
}
if link.Type == linkType && name == targetName {
return link, nil
}
if link.Type == proton.LinkTypeFolder {
childrenLinks, err := protonDrive.c.ListChildren(ctx, protonDrive.MainShare.ShareID, link.LinkID, true)
if err != nil {
return nil, err
}
// log.Printf("childrenLinks len = %v, %#v", len(childrenLinks), childrenLinks)
// get current node's keyring
linkKR, err := link.GetKeyRing(parentNodeKR, protonDrive.AddrKR)
if err != nil {
return nil, err
}
for _, childLink := range childrenLinks {
ret, err := protonDrive.searchByNameRecursively(ctx, linkKR, &childLink, targetName, linkType, listAllActiveOrDraftFiles)
if err != nil {
return nil, err
}
if ret != nil {
return ret, nil
}
}
}
return nil, nil
}
// if the target isn't found, nil will be returned for both return values
func (protonDrive *ProtonDrive) SearchByNameInActiveFolderByID(ctx context.Context,
folderLinkID string,

111
search_recursive.go Normal file
View File

@@ -0,0 +1,111 @@
package proton_api_bridge
import (
"context"
"github.com/ProtonMail/gopenpgp/v2/crypto"
"github.com/henrybear327/go-proton-api"
)
func (protonDrive *ProtonDrive) searchByNameRecursivelyFromRoot(ctx context.Context, targetName string, isFolder bool, listAllActiveOrDraftFiles bool) (*proton.Link, error) {
var linkType proton.LinkType
if isFolder {
linkType = proton.LinkTypeFolder
} else {
linkType = proton.LinkTypeFile
}
return protonDrive.performSearchByNameRecursively(ctx, protonDrive.MainShareKR, protonDrive.RootLink, targetName, linkType, listAllActiveOrDraftFiles)
}
// func (protonDrive *ProtonDrive) searchByNameRecursivelyByID(ctx context.Context, folderLinkID string, targetName string, isFolder bool, listAllActiveOrDraftFiles bool) (*proton.Link, error) {
// folderLink, err := protonDrive.getLink(ctx, folderLinkID)
// if err != nil {
// return nil, err
// }
// var linkType proton.LinkType
// if isFolder {
// linkType = proton.LinkTypeFolder
// } else {
// linkType = proton.LinkTypeFile
// }
// if folderLink.Type != proton.LinkTypeFolder {
// return nil, ErrLinkTypeMustToBeFolderType
// }
// folderKeyRing, err := protonDrive.getLinkKRByID(ctx, folderLink.ParentLinkID)
// if err != nil {
// return nil, err
// }
// return protonDrive.performSearchByNameRecursively(ctx, folderKeyRing, folderLink, targetName, linkType, listAllActiveOrDraftFiles)
// }
func (protonDrive *ProtonDrive) SearchByNameRecursively(ctx context.Context, folderLink *proton.Link, targetName string, isFolder bool, listAllActiveOrDraftFiles bool) (*proton.Link, error) {
var linkType proton.LinkType
if isFolder {
linkType = proton.LinkTypeFolder
} else {
linkType = proton.LinkTypeFile
}
if folderLink.Type != proton.LinkTypeFolder {
return nil, ErrLinkTypeMustToBeFolderType
}
folderKeyRing, err := protonDrive.getLinkKRByID(ctx, folderLink.ParentLinkID)
if err != nil {
return nil, err
}
return protonDrive.performSearchByNameRecursively(ctx, folderKeyRing, folderLink, targetName, linkType, listAllActiveOrDraftFiles)
}
func (protonDrive *ProtonDrive) performSearchByNameRecursively(
ctx context.Context,
parentNodeKR *crypto.KeyRing,
link *proton.Link,
targetName string,
linkType proton.LinkType,
listAllActiveOrDraftFiles bool) (*proton.Link, error) {
if listAllActiveOrDraftFiles {
if link.State != proton.LinkStateActive && link.State != proton.LinkStateDraft {
return nil, nil
}
} else if link.State != proton.LinkStateActive {
return nil, nil
}
name, err := link.GetName(parentNodeKR, protonDrive.AddrKR)
if err != nil {
return nil, err
}
if link.Type == linkType && name == targetName {
return link, nil
}
if link.Type == proton.LinkTypeFolder {
childrenLinks, err := protonDrive.c.ListChildren(ctx, protonDrive.MainShare.ShareID, link.LinkID, true)
if err != nil {
return nil, err
}
// log.Printf("childrenLinks len = %v, %#v", len(childrenLinks), childrenLinks)
// get current node's keyring
linkKR, err := link.GetKeyRing(parentNodeKR, protonDrive.AddrKR)
if err != nil {
return nil, err
}
for _, childLink := range childrenLinks {
ret, err := protonDrive.performSearchByNameRecursively(ctx, linkKR, &childLink, targetName, linkType, listAllActiveOrDraftFiles)
if err != nil {
return nil, err
}
if ret != nil {
return ret, nil
}
}
}
return nil, nil
}