ZBUG-4409: Missing Blobs on a restored account from backup with S3 storage

store/src/java/com/zimbra/cs/store/StoreManager.java

@@ -29,6 +29,7 @@
 import com.zimbra.cs.store.file.FileBlobStore;
 import com.zimbra.cs.store.helper.ClassHelper;
 import com.zimbra.cs.util.Zimbra;
+import com.zimbra.cs.volume.Volume;
 import com.zimbra.cs.volume.VolumeManager;
 
 import java.io.IOException;
@@ -292,6 +293,20 @@ public abstract Blob storeIncoming(InputStream data, boolean storeAsIs)
     public abstract StagedBlob stage(InputStream data, long actualSize, Mailbox mbox)
     throws IOException, ServiceException;
 
+    /**
+     * Stage an incoming <code>InputStream</code> to an
+     * appropriate place for subsequent storage in a <code>Mailbox</code> via
+     * {@link #link(StagedBlob, Mailbox, int, int)} or {@link #renameTo}.
+     *
+     * @param data the data stream
+     * @param actualSize the content size, or {@code -1} if the content size is not available
+     * @param callback callback, or {@code null}
+     * @param mbox the mailbox
+     * @param volume the volume
+     */
+    public abstract StagedBlob stage(InputStream data, long actualSize, Mailbox mbox, Volume volume)
+            throws IOException, ServiceException;
+
     /**
      * Stage an incoming <code>InputStream</code> to an
      * appropriate place for subsequent storage in a <code>Mailbox</code> via
@@ -306,6 +321,21 @@ public StagedBlob stage(InputStream data, Mailbox mbox)
         return stage(data, -1, mbox);
     }
 
+    /**
+     * Stage an incoming <code>InputStream</code> to an
+     * appropriate place for subsequent storage in a <code>Mailbox</code> via
+     * {@link #link(StagedBlob, Mailbox, int, int)} or {@link #renameTo}.
+     *
+     * @param data the data stream
+     * @param callback callback, or {@code null}
+     * @param mbox the mailbox
+     * @param volume the volume
+     */
+    public StagedBlob stage(InputStream data, Mailbox mbox, Volume volume)
+            throws IOException, ServiceException {
+        return stage(data, -1, mbox, volume);
+    }
+
     /**
      * Stage an incoming <code>Blob</code> (see {@link #storeIncoming}) to an
      * appropriate place for subsequent storage in a <code>Mailbox</code> via

store/src/java/com/zimbra/cs/store/external/ExternalBlobIO.java

@@ -38,6 +38,21 @@ public interface ExternalBlobIO {
      */
     String writeStreamToStore(InputStream in, long actualSize, Mailbox mbox) throws IOException, ServiceException;
 
+    /**
+     * Write data to blob store
+     *
+     * @param in:         InputStream containing data to be written
+     * @param actualSize: size of data in stream, or -1 if size is unknown. To be used by implementor for optimization where possible
+     * @param mbox:       Mailbox which contains the blob. Can optionally be used by store for partitioning
+     * @param destVolId:  Destination Volume id on which we need to store the blob
+     * @return locator string for the stored blob, unique identifier created by storage protocol
+     * @throws IOException
+     * @throws ServiceException
+     */
+    default String writeStreamToStore(InputStream in, long actualSize, Mailbox mbox, short destVolId) throws IOException, ServiceException {
+        return null;
+    }
+
     /**
      * Create an input stream for reading data from blob store
      * @param locator: identifier string for the blob as returned from write operation

store/src/java/com/zimbra/cs/store/external/ExternalStoreManager.java

@@ -44,6 +44,7 @@
 import com.zimbra.cs.store.StagedBlob;
 import com.zimbra.cs.store.StoreManager;
 import com.zimbra.cs.store.file.VolumeMailboxBlob;
+import com.zimbra.cs.volume.Volume;
 
 /**
  * Abstract base class for external store integration.
@@ -272,6 +273,31 @@ public StagedBlob stage(Blob blob, Mailbox mbox) throws IOException, ServiceExce
         }
     }
 
+    public StagedBlob stage(Blob blob, Mailbox mbox, Volume volume) throws IOException, ServiceException {
+        if (supports(StoreFeature.RESUMABLE_UPLOAD) && blob instanceof ExternalUploadedBlob) {
+            ZimbraLog.store.debug("blob already uploaded, just need to commit");
+            String locator = ((ExternalResumableUpload) this).finishUpload((ExternalUploadedBlob) blob);
+            if (locator != null) {
+                ZimbraLog.store.debug("wrote to locator %s",locator);
+                localCache.put(locator, getContent(blob));
+            } else {
+                ZimbraLog.store.warn("blob staging returned null locator");
+            }
+            return new ExternalStagedBlob(mbox, blob.getDigest(), blob.getRawSize(), locator);
+        } else {
+            InputStream is = getContent(blob);
+            try {
+                StagedBlob staged = stage(is, blob.getRawSize(), mbox, volume);
+                if (staged != null && staged.getLocator() != null) {
+                    localCache.put(staged.getLocator(), getContent(blob));
+                }
+                return staged;
+            } finally {
+                ByteUtil.closeStream(is);
+            }
+        }
+    }
+
     @Override
     public StagedBlob stage(InputStream in, long actualSize, Mailbox mbox) throws ServiceException, IOException {
         if (actualSize < 0) {
@@ -303,7 +329,35 @@ public StagedBlob stage(InputStream in, long actualSize, Mailbox mbox) throws Se
         }
     }
 
+    public StagedBlob stage(InputStream in, long actualSize, Mailbox mbox, Volume volume) throws ServiceException, IOException {
+        if (actualSize < 0) {
+            Blob blob = storeIncoming(in);
+            try {
+                return stage(blob, mbox, volume);
+            } finally {
+                quietDelete(blob);
+            }
+        }
+        MessageDigest digest;
+        try {
+            digest = MessageDigest.getInstance("SHA-256");
+        } catch (NoSuchAlgorithmException e) {
+            throw ServiceException.FAILURE("SHA-256 digest not found", e);
+        }
+        ByteUtil.PositionInputStream pin = new ByteUtil.PositionInputStream(new DigestInputStream(in, digest));
 
+        try {
+            String locator = writeStreamToStore(pin, actualSize, mbox, volume.getId());
+            if (locator != null) {
+                ZimbraLog.store.debug("wrote to locator %s",locator);
+            } else {
+                ZimbraLog.store.warn("blob staging returned null locator");
+            }
+            return new ExternalStagedBlob(mbox, ByteUtil.encodeFSSafeBase64(digest.digest()), pin.getPosition(), locator);
+        } catch (IOException e) {
+            throw ServiceException.FAILURE("unable to stage blob", e);
+        }
+    }
     @Override
     public Blob storeIncoming(InputStream data, boolean storeAsIs) throws IOException,
     ServiceException {

store/src/java/com/zimbra/cs/store/file/FileBlobStore.java

@@ -118,6 +118,12 @@ public VolumeStagedBlob stage(InputStream in, long actualSize, Mailbox mbox)
         return new VolumeStagedBlob(mbox, (VolumeBlob) blob).markStagedDirectly();
     }
 
+    @Override
+    public StagedBlob stage(InputStream data, long actualSize, Mailbox mbox, Volume volume) throws IOException, ServiceException {
+        // mailbox store is on the same volume as incoming directory, so no need to stage the blob
+        throw ServiceException.FAILURE("Operation can not be completed because ExternalStoreManager is not available", null);
+    }
+
     @Override
     public VolumeStagedBlob stage(Blob blob, Mailbox mbox) throws IOException {
         // mailbox store is on the same volume as incoming directory, so no need to stage the blob