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

store/src/java-test/com/zimbra/cs/store/MockStoreManager.java

@@ -27,6 +27,7 @@
 import java.util.HashMap;
 import java.util.Map;
 
+import com.zimbra.cs.volume.Volume;
 import org.apache.commons.io.output.ByteArrayOutputStream;
 
 import com.google.common.io.ByteStreams;
@@ -106,6 +107,11 @@ public StagedBlob stage(InputStream data, long actualSize, Mailbox mbox) throws
         return new MockStagedBlob(mbox, ByteStreams.toByteArray(data));
     }
 
+    @Override
+    public StagedBlob stage(InputStream data, long actualSize, Mailbox mbox, Volume volume) throws IOException {
+        return new MockStagedBlob(mbox, ByteStreams.toByteArray(data));
+    }
+
     @Override
     public StagedBlob stage(Blob blob, Mailbox mbox) {
         return new MockStagedBlob(mbox, ((MockBlob) blob).content);

store/src/java-test/com/zimbra/cs/store/StoreManagerAfterReset.java

@@ -18,6 +18,7 @@
 
 import com.zimbra.common.service.ServiceException;
 import com.zimbra.cs.mailbox.Mailbox;
+import com.zimbra.cs.volume.Volume;
 
 import java.io.IOException;
 import java.io.InputStream;
@@ -59,6 +60,11 @@ public StagedBlob stage(InputStream data, long actualSize, Mailbox mbox) throws
         return null;
     }
 
+    @Override
+    public StagedBlob stage(InputStream data, long actualSize, Mailbox mbox, Volume volume) throws IOException, ServiceException {
+        return null;
+    }
+
     @Override
     public StagedBlob stage(Blob blob, Mailbox mbox) throws IOException, ServiceException {
         return null;

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;
@@ -46,6 +47,7 @@ public abstract class StoreManager {
      */
     private static short currentVolumeId;
     private static Integer diskStreamingThreshold;
+    private static final int actualSizeforUnstagedBlob = -1;
 
 
     public static StoreManager getInstance () {
@@ -292,6 +294,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
@@ -303,7 +319,22 @@ public abstract StagedBlob stage(InputStream data, long actualSize, Mailbox mbox
      */
     public StagedBlob stage(InputStream data, Mailbox mbox)
     throws IOException, ServiceException {
-        return stage(data, -1, mbox);
+        return stage(data, actualSizeforUnstagedBlob, 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, actualSizeforUnstagedBlob, mbox, volume);
     }
 
     /**

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 {
+        throw ServiceException.UNSUPPORTED();
+    }
+
     /**
      * 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 (null != locator) {
+                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 (null != staged && null != staged.getLocator()) {
+                    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 (null != locator) {
+                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