file_copy.go 6.4 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176
  1. package file_copy
  2. import (
  3. "fmt"
  4. "github.com/aws/aws-sdk-go/service/s3"
  5. "github.com/ente-io/museum/ente"
  6. "github.com/ente-io/museum/pkg/controller"
  7. "github.com/ente-io/museum/pkg/repo"
  8. "github.com/ente-io/museum/pkg/utils/auth"
  9. "github.com/ente-io/museum/pkg/utils/s3config"
  10. enteTime "github.com/ente-io/museum/pkg/utils/time"
  11. "github.com/gin-contrib/requestid"
  12. "github.com/gin-gonic/gin"
  13. "github.com/sirupsen/logrus"
  14. "time"
  15. )
  16. type FileCopyController struct {
  17. S3Config *s3config.S3Config
  18. FileController *controller.FileController
  19. FileRepo *repo.FileRepository
  20. CollectionCtrl *controller.CollectionController
  21. ObjectRepo *repo.ObjectRepository
  22. }
  23. type copyS3ObjectReq struct {
  24. SourceS3Object ente.S3ObjectKey
  25. DestObjectKey string
  26. }
  27. type fileCopyInternal struct {
  28. SourceFile ente.File
  29. DestCollectionID int64
  30. // The FileKey is encrypted with the destination collection's key
  31. EncryptedFileKey string
  32. EncryptedFileKeyNonce string
  33. FileCopyReq *copyS3ObjectReq
  34. ThumbCopyReq *copyS3ObjectReq
  35. }
  36. func (fci fileCopyInternal) newFile(ownedID int64) ente.File {
  37. newFileAttributes := fci.SourceFile.File
  38. newFileAttributes.ObjectKey = fci.FileCopyReq.DestObjectKey
  39. newThumbAttributes := fci.SourceFile.Thumbnail
  40. newThumbAttributes.ObjectKey = fci.ThumbCopyReq.DestObjectKey
  41. return ente.File{
  42. OwnerID: ownedID,
  43. CollectionID: fci.DestCollectionID,
  44. EncryptedKey: fci.EncryptedFileKey,
  45. KeyDecryptionNonce: fci.EncryptedFileKeyNonce,
  46. File: newFileAttributes,
  47. Thumbnail: newThumbAttributes,
  48. Metadata: fci.SourceFile.Metadata,
  49. UpdationTime: enteTime.Microseconds(),
  50. IsDeleted: false,
  51. }
  52. }
  53. func (fc *FileCopyController) CopyFiles(c *gin.Context, req ente.CopyFileSyncRequest) (*ente.CopyResponse, error) {
  54. userID := auth.GetUserID(c.Request.Header)
  55. app := auth.GetApp(c)
  56. logger := logrus.WithFields(logrus.Fields{"req_id": requestid.Get(c), "user_id": userID})
  57. err := fc.CollectionCtrl.IsCopyAllowed(c, userID, req)
  58. if err != nil {
  59. return nil, err
  60. }
  61. fileIDs := req.FileIDs()
  62. fileToCollectionFileMap := make(map[int64]*ente.CollectionFileItem, len(req.CollectionFileItems))
  63. for _, collectionFileItem := range req.CollectionFileItems {
  64. fileToCollectionFileMap[collectionFileItem.ID] = &collectionFileItem
  65. }
  66. s3ObjectsToCopy, err := fc.ObjectRepo.GetObjectsForFileIDs(fileIDs)
  67. if err != nil {
  68. return nil, err
  69. }
  70. // note: this assumes that preview existingFilesToCopy for videos are not tracked inside the object_keys table
  71. if len(s3ObjectsToCopy) != 2*len(fileIDs) {
  72. return nil, ente.NewInternalError(fmt.Sprintf("expected %d objects, got %d", 2*len(fileIDs), len(s3ObjectsToCopy)))
  73. }
  74. // todo:(neeraj) if the total size is greater than 1GB, do an early check if the user can upload the existingFilesToCopy
  75. var totalSize int64
  76. for _, obj := range s3ObjectsToCopy {
  77. totalSize += obj.FileSize
  78. }
  79. logger.WithField("totalSize", totalSize).Info("total size of existingFilesToCopy to copy")
  80. // request the uploadUrls using existing method. This is to ensure that orphan objects are automatically cleaned up
  81. // todo:(neeraj) optimize this method by removing the need for getting a signed url for each object
  82. uploadUrls, err := fc.FileController.GetUploadURLs(c, userID, len(s3ObjectsToCopy), app)
  83. if err != nil {
  84. return nil, err
  85. }
  86. existingFilesToCopy, err := fc.FileRepo.GetFileAttributesForCopy(fileIDs)
  87. if err != nil {
  88. return nil, err
  89. }
  90. if len(existingFilesToCopy) != len(fileIDs) {
  91. return nil, ente.NewInternalError(fmt.Sprintf("expected %d existingFilesToCopy, got %d", len(fileIDs), len(existingFilesToCopy)))
  92. }
  93. fileOGS3Object := make(map[int64]*copyS3ObjectReq)
  94. fileThumbS3Object := make(map[int64]*copyS3ObjectReq)
  95. for i, s3Obj := range s3ObjectsToCopy {
  96. if s3Obj.Type == ente.FILE {
  97. fileOGS3Object[s3Obj.FileID] = &copyS3ObjectReq{
  98. SourceS3Object: s3Obj,
  99. DestObjectKey: uploadUrls[i].ObjectKey,
  100. }
  101. } else if s3Obj.Type == ente.THUMBNAIL {
  102. fileThumbS3Object[s3Obj.FileID] = &copyS3ObjectReq{
  103. SourceS3Object: s3Obj,
  104. DestObjectKey: uploadUrls[i].ObjectKey,
  105. }
  106. } else {
  107. return nil, ente.NewInternalError(fmt.Sprintf("unexpected object type %s", s3Obj.Type))
  108. }
  109. }
  110. fileCopyList := make([]fileCopyInternal, 0, len(existingFilesToCopy))
  111. for _, file := range existingFilesToCopy {
  112. fileCopy := fileCopyInternal{
  113. SourceFile: file,
  114. DestCollectionID: req.DstCollection,
  115. EncryptedFileKey: fileToCollectionFileMap[file.ID].EncryptedKey,
  116. EncryptedFileKeyNonce: fileToCollectionFileMap[file.ID].KeyDecryptionNonce,
  117. FileCopyReq: fileOGS3Object[file.ID],
  118. ThumbCopyReq: fileThumbS3Object[file.ID],
  119. }
  120. fileCopyList = append(fileCopyList, fileCopy)
  121. }
  122. oldToNewFileIDMap := make(map[int64]int64)
  123. for _, fileCopy := range fileCopyList {
  124. newFile, err := fc.createCopy(c, fileCopy, userID, app)
  125. if err != nil {
  126. return nil, err
  127. }
  128. oldToNewFileIDMap[fileCopy.SourceFile.ID] = newFile.ID
  129. }
  130. return &ente.CopyResponse{OldToNewFileIDMap: oldToNewFileIDMap}, nil
  131. }
  132. func (fc *FileCopyController) createCopy(c *gin.Context, fcInternal fileCopyInternal, userID int64, app ente.App) (*ente.File, error) {
  133. // using HotS3Client copy the File and Thumbnail
  134. s3Client := fc.S3Config.GetHotS3Client()
  135. hotBucket := fc.S3Config.GetHotBucket()
  136. err := copyS3Object(s3Client, hotBucket, fcInternal.FileCopyReq)
  137. if err != nil {
  138. return nil, err
  139. }
  140. err = copyS3Object(s3Client, hotBucket, fcInternal.ThumbCopyReq)
  141. if err != nil {
  142. return nil, err
  143. }
  144. file := fcInternal.newFile(userID)
  145. newFile, err := fc.FileController.Create(c, userID, file, "", app)
  146. if err != nil {
  147. return nil, err
  148. }
  149. return &newFile, nil
  150. }
  151. // Helper function for S3 object copying.
  152. func copyS3Object(s3Client *s3.S3, bucket *string, req *copyS3ObjectReq) error {
  153. copySource := fmt.Sprintf("%s/%s", *bucket, req.SourceS3Object.ObjectKey)
  154. copyInput := &s3.CopyObjectInput{
  155. Bucket: bucket,
  156. CopySource: &copySource,
  157. Key: &req.DestObjectKey,
  158. }
  159. start := time.Now()
  160. _, err := s3Client.CopyObject(copyInput)
  161. elapsed := time.Since(start)
  162. if err != nil {
  163. return fmt.Errorf("failed to copy (%s) from %s to %s: %w", req.SourceS3Object.Type, copySource, req.DestObjectKey, err)
  164. }
  165. logrus.WithField("duration", elapsed).WithField("size", req.SourceS3Object.FileSize).Infof("copied (%s) from %s to %s", req.SourceS3Object.Type, copySource, req.DestObjectKey)
  166. return nil
  167. }