diff --git a/file_upload.go b/file_upload.go index 14f7a58..de28bc0 100644 --- a/file_upload.go +++ b/file_upload.go @@ -8,6 +8,7 @@ import ( "crypto/sha256" "encoding/base64" "encoding/hex" + "fmt" "io" "mime" "os" @@ -74,10 +75,32 @@ func (protonDrive *ProtonDrive) handleRevisionConflict(ctx context.Context, link } } -func (protonDrive *ProtonDrive) createFileUploadDraft(ctx context.Context, parentLink *proton.Link, filename string, modTime time.Time, mimeType string) (string, string, *crypto.SessionKey, *crypto.KeyRing, error) { +type verificationData struct { + sessionKey *crypto.SessionKey + verificationCode []byte +} + +func (verifier *verificationData) verifyBlock(encData []byte) ([]byte, error) { + _, err := verifier.sessionKey.Decrypt(encData) + if err != nil { + return nil, fmt.Errorf("verification failed: %w", err) + } + + result := make([]byte, len(verifier.verificationCode)) + for i, verCodeByte := range verifier.verificationCode { + var encDataByte byte + if i < len(encData) { + encDataByte = encData[i] + } + result[i] = verCodeByte ^ encDataByte + } + return result, nil +} + +func (protonDrive *ProtonDrive) createFileUploadDraft(ctx context.Context, parentLink *proton.Link, filename string, modTime time.Time, mimeType string) (string, string, *crypto.SessionKey, *crypto.KeyRing, *verificationData, error) { parentNodeKR, err := protonDrive.getLinkKR(ctx, parentLink) if err != nil { - return "", "", nil, nil, err + return "", "", nil, nil, nil, err } /* @@ -86,7 +109,7 @@ func (protonDrive *ProtonDrive) createFileUploadDraft(ctx context.Context, paren */ newNodeKey, newNodePassphraseEnc, newNodePassphraseSignature, err := generateNodeKeys(parentNodeKR, protonDrive.DefaultAddrKR) if err != nil { - return "", "", nil, nil, err + return "", "", nil, nil, nil, err } createFileReq := proton.CreateFileReq{ @@ -112,7 +135,7 @@ func (protonDrive *ProtonDrive) createFileUploadDraft(ctx context.Context, paren */ err = createFileReq.SetName(filename, protonDrive.DefaultAddrKR, parentNodeKR) if err != nil { - return "", "", nil, nil, err + return "", "", nil, nil, nil, err } /* @@ -121,17 +144,17 @@ func (protonDrive *ProtonDrive) createFileUploadDraft(ctx context.Context, paren */ signatureVerificationKR, err := protonDrive.getSignatureVerificationKeyring([]string{parentLink.SignatureEmail}, parentNodeKR) if err != nil { - return "", "", nil, nil, err + return "", "", nil, nil, nil, err } parentHashKey, err := parentLink.GetHashKey(parentNodeKR, signatureVerificationKR) if err != nil { - return "", "", nil, nil, err + return "", "", nil, nil, nil, err } /* Use parent's hash key */ err = createFileReq.SetHash(filename, parentHashKey) if err != nil { - return "", "", nil, nil, err + return "", "", nil, nil, nil, err } /* @@ -140,7 +163,7 @@ func (protonDrive *ProtonDrive) createFileUploadDraft(ctx context.Context, paren */ newNodeKR, err := getKeyRing(parentNodeKR, protonDrive.DefaultAddrKR, newNodeKey, newNodePassphraseEnc, newNodePassphraseSignature) if err != nil { - return "", "", nil, nil, err + return "", "", nil, nil, nil, err } /* @@ -149,7 +172,7 @@ func (protonDrive *ProtonDrive) createFileUploadDraft(ctx context.Context, paren */ newSessionKey, err := createFileReq.SetContentKeyPacketAndSignature(newNodeKR) if err != nil { - return "", "", nil, nil, err + return "", "", nil, nil, nil, err } createFileAction := func() (*proton.CreateFileRes, *proton.Link, error) { @@ -190,12 +213,12 @@ func (protonDrive *ProtonDrive) createFileUploadDraft(ctx context.Context, paren createFileResp, link, err := createFileAction() if err != nil { - return "", "", nil, nil, err + return "", "", nil, nil, nil, err } revisionID, shouldSubmitCreateFileRequestAgain, err := protonDrive.handleRevisionConflict(ctx, link, createFileResp) if err != nil { - return "", "", nil, nil, err + return "", "", nil, nil, nil, err } if shouldSubmitCreateFileRequestAgain { @@ -203,12 +226,12 @@ func (protonDrive *ProtonDrive) createFileUploadDraft(ctx context.Context, paren // we need to delete the link and recreate one createFileResp, link, err = createFileAction() if err != nil { - return "", "", nil, nil, err + return "", "", nil, nil, nil, err } revisionID, _, err = protonDrive.handleRevisionConflict(ctx, link, createFileResp) if err != nil { - return "", "", nil, nil, err + return "", "", nil, nil, nil, err } } @@ -219,34 +242,54 @@ func (protonDrive *ProtonDrive) createFileUploadDraft(ctx context.Context, paren // get original sessionKey and nodeKR for the current link parentNodeKR, err = protonDrive.getLinkKRByID(ctx, link.ParentLinkID) if err != nil { - return "", "", nil, nil, err + return "", "", nil, nil, nil, err } signatureVerificationKR, err := protonDrive.getSignatureVerificationKeyring([]string{link.SignatureEmail}) if err != nil { - return "", "", nil, nil, err + return "", "", nil, nil, nil, err } newNodeKR, err = link.GetKeyRing(parentNodeKR, signatureVerificationKR) if err != nil { - return "", "", nil, nil, err + return "", "", nil, nil, nil, err } newSessionKey, err = link.GetSessionKey(newNodeKR) if err != nil { - return "", "", nil, nil, err + return "", "", nil, nil, nil, err } } else { linkID = createFileResp.ID } - return linkID, revisionID, newSessionKey, newNodeKR, nil + verification, err := protonDrive.c.VerifyRevision(ctx, protonDrive.MainShare.ShareID, linkID, revisionID) + if err != nil { + return "", "", nil, nil, nil, err + } + encSessionKey, err := base64.StdEncoding.DecodeString(verification.ContentKeyPacket) + if err != nil { + return "", "", nil, nil, nil, err + } + verifierSessionKey, err := newNodeKR.DecryptSessionKey(encSessionKey) + if err != nil { + return "", "", nil, nil, nil, err + } + verificationCode, err := base64.StdEncoding.DecodeString(verification.VerificationCode) + if err != nil { + return "", "", nil, nil, nil, err + } + + return linkID, revisionID, newSessionKey, newNodeKR, &verificationData{ + sessionKey: verifierSessionKey, + verificationCode: verificationCode, + }, 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) { +func (protonDrive *ProtonDrive) uploadAndCollectBlockData(ctx context.Context, newSessionKey *crypto.SessionKey, verifier *verificationData, 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 { + if newSessionKey == nil || newNodeKR == nil || verifier == nil { return nil, 0, nil, "", ErrMissingInputUploadAndCollectBlockData } @@ -360,10 +403,13 @@ func (protonDrive *ProtonDrive) uploadAndCollectBlockData(ctx context.Context, n h.Write(encData) hash := h.Sum(nil) base64Hash := base64.StdEncoding.EncodeToString(hash) + manifestSignatureData = append(manifestSignatureData, hash...) + + verificationToken, err := verifier.verifyBlock(encData) if err != nil { return nil, 0, nil, "", err } - manifestSignatureData = append(manifestSignatureData, hash...) + verificationTokenBase64 := base64.StdEncoding.EncodeToString(verificationToken) pendingUploadBlocks = append(pendingUploadBlocks, PendingUploadBlocks{ blockUploadInfo: proton.BlockUploadInfo{ @@ -371,6 +417,9 @@ func (protonDrive *ProtonDrive) uploadAndCollectBlockData(ctx context.Context, n Size: int64(len(encData)), EncSignature: encSignatureStr, Hash: base64Hash, + Verifier: proton.BlockUploadInfoVerifier{ + Token: verificationTokenBase64, + }, }, encData: encData, }) @@ -432,7 +481,7 @@ func (protonDrive *ProtonDrive) uploadFile(ctx context.Context, parentLink *prot } /* step 1: create a draft */ - linkID, revisionID, newSessionKey, newNodeKR, err := protonDrive.createFileUploadDraft(ctx, parentLink, filename, modTime, mimeType) + linkID, revisionID, newSessionKey, newNodeKR, verifier, err := protonDrive.createFileUploadDraft(ctx, parentLink, filename, modTime, mimeType) if err != nil { return "", nil, err } @@ -442,7 +491,7 @@ func (protonDrive *ProtonDrive) uploadFile(ctx context.Context, parentLink *prot } /* step 2: upload blocks and collect block data */ - manifestSignature, fileSize, blockSizes, digests, err := protonDrive.uploadAndCollectBlockData(ctx, newSessionKey, newNodeKR, file, linkID, revisionID) + manifestSignature, fileSize, blockSizes, digests, err := protonDrive.uploadAndCollectBlockData(ctx, newSessionKey, verifier, newNodeKR, file, linkID, revisionID) if err != nil { return "", nil, err }