diff --git a/core/src/main/java/google/registry/backup/BackupUtils.java b/core/src/main/java/google/registry/backup/BackupUtils.java index 6ccb882fa..8a036f319 100644 --- a/core/src/main/java/google/registry/backup/BackupUtils.java +++ b/core/src/main/java/google/registry/backup/BackupUtils.java @@ -56,12 +56,16 @@ public class BackupUtils { * *

The iterator reads from the stream on demand, and as such will fail if the stream is closed. */ - public static Iterator createDeserializingIterator(final InputStream input) { + public static Iterator createDeserializingIterator( + final InputStream input, boolean withAppIdOverride) { return new AbstractIterator() { @Override protected ImmutableObject computeNext() { EntityProto proto = new EntityProto(); if (proto.parseDelimitedFrom(input)) { // False means end of stream; other errors throw. + if (withAppIdOverride) { + proto = EntityImports.fixEntity(proto); + } return auditedOfy().load().fromEntity(EntityTranslator.createFromPb(proto)); } return endOfData(); @@ -70,6 +74,7 @@ public class BackupUtils { } public static ImmutableList deserializeEntities(byte[] bytes) { - return ImmutableList.copyOf(createDeserializingIterator(new ByteArrayInputStream(bytes))); + return ImmutableList.copyOf( + createDeserializingIterator(new ByteArrayInputStream(bytes), false)); } } diff --git a/core/src/main/java/google/registry/backup/CommitLogImports.java b/core/src/main/java/google/registry/backup/CommitLogImports.java index da7377f08..06429fc46 100644 --- a/core/src/main/java/google/registry/backup/CommitLogImports.java +++ b/core/src/main/java/google/registry/backup/CommitLogImports.java @@ -59,7 +59,7 @@ public final class CommitLogImports { InputStream inputStream) { try (AppEngineEnvironment appEngineEnvironment = new AppEngineEnvironment(); InputStream input = new BufferedInputStream(inputStream)) { - Iterator commitLogs = createDeserializingIterator(input); + Iterator commitLogs = createDeserializingIterator(input, false); checkState(commitLogs.hasNext()); checkState(commitLogs.next() instanceof CommitLogCheckpoint); diff --git a/core/src/main/java/google/registry/backup/EntityImports.java b/core/src/main/java/google/registry/backup/EntityImports.java new file mode 100644 index 000000000..1e87a3bef --- /dev/null +++ b/core/src/main/java/google/registry/backup/EntityImports.java @@ -0,0 +1,115 @@ +// Copyright 2021 The Nomulus Authors. All Rights Reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package google.registry.backup; + +import com.google.apphosting.api.ApiProxy; +import com.google.storage.onestore.v3.OnestoreEntity; +import com.google.storage.onestore.v3.OnestoreEntity.EntityProto; +import com.google.storage.onestore.v3.OnestoreEntity.Path; +import com.google.storage.onestore.v3.OnestoreEntity.Property.Meaning; +import com.google.storage.onestore.v3.OnestoreEntity.PropertyValue.ReferenceValue; +import java.nio.charset.StandardCharsets; +import java.util.Objects; + +/** Utilities for handling imported Datastore entities. */ +public class EntityImports { + + /** + * Transitively sets the {@code appId} of all keys in a foreign entity to that of the current + * system. + */ + public static EntityProto fixEntity(EntityProto entityProto) { + String currentAappId = ApiProxy.getCurrentEnvironment().getAppId(); + if (Objects.equals(currentAappId, entityProto.getKey().getApp())) { + return entityProto; + } + return fixEntity(entityProto, currentAappId); + } + + private static EntityProto fixEntity(EntityProto entityProto, String appId) { + if (entityProto.hasKey()) { + fixKey(entityProto, appId); + } + + for (OnestoreEntity.Property property : entityProto.mutablePropertys()) { + fixProperty(property, appId); + } + + for (OnestoreEntity.Property property : entityProto.mutableRawPropertys()) { + fixProperty(property, appId); + } + + // CommitLogMutation embeds an entity as bytes, which needs additional fixes. + if (isCommitLogMutation(entityProto)) { + fixMutationEntityProtoBytes(entityProto, appId); + } + return entityProto; + } + + private static boolean isCommitLogMutation(EntityProto entityProto) { + if (!entityProto.hasKey()) { + return false; + } + Path path = entityProto.getKey().getPath(); + if (path.elementSize() == 0) { + return false; + } + return Objects.equals( + path.getElement(path.elementSize() - 1).getType(StandardCharsets.UTF_8), + "CommitLogMutation"); + } + + private static void fixMutationEntityProtoBytes(EntityProto entityProto, String appId) { + for (OnestoreEntity.Property property : entityProto.mutableRawPropertys()) { + if (Objects.equals(property.getName(), "entityProtoBytes")) { + OnestoreEntity.PropertyValue value = property.getValue(); + EntityProto fixedProto = + fixEntity(bytesToEntityProto(value.getStringValueAsBytes()), appId); + value.setStringValueAsBytes(fixedProto.toByteArray()); + return; + } + } + } + + private static void fixKey(EntityProto entityProto, String appId) { + entityProto.getMutableKey().setApp(appId); + } + + private static void fixKey(ReferenceValue referenceValue, String appId) { + referenceValue.setApp(appId); + } + + private static void fixProperty(OnestoreEntity.Property property, String appId) { + OnestoreEntity.PropertyValue value = property.getMutableValue(); + if (value.hasReferenceValue()) { + fixKey(value.getMutableReferenceValue(), appId); + return; + } + if (property.getMeaningEnum().equals(Meaning.ENTITY_PROTO)) { + EntityProto embeddedProto = bytesToEntityProto(value.getStringValueAsBytes()); + fixEntity(embeddedProto, appId); + value.setStringValueAsBytes(embeddedProto.toByteArray()); + } + } + + private static EntityProto bytesToEntityProto(byte[] bytes) { + EntityProto entityProto = new EntityProto(); + boolean isParsed = entityProto.parseFrom(bytes); + if (!isParsed) { + throw new IllegalStateException("Failed to parse raw bytes as EntityProto."); + } + return entityProto; + } +} diff --git a/core/src/main/java/google/registry/backup/RestoreCommitLogsAction.java b/core/src/main/java/google/registry/backup/RestoreCommitLogsAction.java index 62b1c50d8..2f93667bd 100644 --- a/core/src/main/java/google/registry/backup/RestoreCommitLogsAction.java +++ b/core/src/main/java/google/registry/backup/RestoreCommitLogsAction.java @@ -104,7 +104,7 @@ public class RestoreCommitLogsAction implements Runnable { try (InputStream input = Channels.newInputStream( gcsService.openPrefetchingReadChannel(metadata.getFilename(), 0, BLOCK_SIZE))) { PeekingIterator commitLogs = - peekingIterator(createDeserializingIterator(input)); + peekingIterator(createDeserializingIterator(input, true)); lastCheckpoint = (CommitLogCheckpoint) commitLogs.next(); saveOfy(ImmutableList.of(lastCheckpoint)); // Save the checkpoint itself. while (commitLogs.hasNext()) { diff --git a/core/src/test/java/google/registry/backup/EntityImportsTest.java b/core/src/test/java/google/registry/backup/EntityImportsTest.java new file mode 100644 index 000000000..e5e105a3a --- /dev/null +++ b/core/src/test/java/google/registry/backup/EntityImportsTest.java @@ -0,0 +1,87 @@ +// Copyright 2021 The Nomulus Authors. All Rights Reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package google.registry.backup; + +import static com.google.common.truth.Truth.assertThat; +import static google.registry.model.ofy.ObjectifyService.auditedOfy; + +import com.google.appengine.api.datastore.DatastoreService; +import com.google.appengine.api.datastore.DatastoreServiceFactory; +import com.google.appengine.api.datastore.Entity; +import com.google.appengine.api.datastore.EntityTranslator; +import com.google.common.collect.ImmutableList; +import com.google.common.io.Resources; +import com.google.storage.onestore.v3.OnestoreEntity.EntityProto; +import google.registry.model.ofy.CommitLogCheckpoint; +import google.registry.testing.AppEngineExtension; +import google.registry.testing.DatastoreEntityExtension; +import java.io.InputStream; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Order; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.extension.RegisterExtension; + +public class EntityImportsTest { + + @RegisterExtension + @Order(value = 1) + final DatastoreEntityExtension datastoreEntityExtension = new DatastoreEntityExtension(); + + @RegisterExtension + final AppEngineExtension appEngine = + new AppEngineExtension.Builder().withDatastoreAndCloudSql().withoutCannedData().build(); + + private DatastoreService datastoreService; + + @BeforeEach + void beforeEach() { + datastoreService = DatastoreServiceFactory.getDatastoreService(); + } + + @Test + void importCommitLogs_keysFixed() throws Exception { + // Input resource is a standard commit log file whose entities has "AppId_1" as appId. The key + // fixes can be verified by checking that the appId of an imported entity's key has been updated + // to 'test' (which is set by AppEngineExtension) and/or that after persistence the imported + // entity can be loaded by Objectify. + try (InputStream commitLogInputStream = + Resources.getResource("google/registry/backup/commitlog.data").openStream()) { + ImmutableList entities = + loadEntityProtos(commitLogInputStream).stream() + .map(EntityImports::fixEntity) + .map(EntityTranslator::createFromPb) + .collect(ImmutableList.toImmutableList()); + // Verifies that the original appId has been overwritten. + assertThat(entities.get(0).getKey().getAppId()).isEqualTo("test"); + datastoreService.put(entities); + // Imported entity can be found by Ofy after appId conversion. + assertThat(auditedOfy().load().type(CommitLogCheckpoint.class).count()).isGreaterThan(0); + } + } + + private static ImmutableList loadEntityProtos(InputStream inputStream) { + ImmutableList.Builder protosBuilder = new ImmutableList.Builder<>(); + while (true) { + EntityProto proto = new EntityProto(); + boolean parsed = proto.parseDelimitedFrom(inputStream); + if (parsed && proto.isInitialized()) { + protosBuilder.add(proto); + } else { + break; + } + } + return protosBuilder.build(); + } +} diff --git a/core/src/test/java/google/registry/backup/RestoreCommitLogsActionTest.java b/core/src/test/java/google/registry/backup/RestoreCommitLogsActionTest.java index d935d7be4..b7001199d 100644 --- a/core/src/test/java/google/registry/backup/RestoreCommitLogsActionTest.java +++ b/core/src/test/java/google/registry/backup/RestoreCommitLogsActionTest.java @@ -34,9 +34,11 @@ import com.google.appengine.tools.cloudstorage.GcsService; import com.google.common.collect.ImmutableMap; import com.google.common.collect.ImmutableSet; import com.google.common.collect.Lists; +import com.google.common.io.Resources; import com.google.common.primitives.Longs; import com.googlecode.objectify.Key; import google.registry.model.ImmutableObject; +import google.registry.model.domain.DomainBase; import google.registry.model.ofy.CommitLogBucket; import google.registry.model.ofy.CommitLogCheckpoint; import google.registry.model.ofy.CommitLogCheckpointRoot; @@ -258,10 +260,40 @@ public class RestoreCommitLogsActionTest { assertInDatastore(CommitLogCheckpointRoot.create(now)); } + @Test + void testRestore_fromOtherProject() throws IOException { + // Input resource is a standard commit log file whose entities has "AppId_1" as appId. Among the + // entities are CommitLogMutations that have an embedded DomainBase and a ContactResource, both + // having "AppId_1" as appId. This test verifies that the embedded entities are properly + // imported, in particular, the domain's 'registrant' key can be used by Objectify to load the + // contact. + saveDiffFile( + gcsService, + Resources.toByteArray(Resources.getResource("google/registry/backup/commitlog.data")), + now); + action.run(); + auditedOfy().clearSessionCache(); + List domainBases = auditedOfy().load().type(DomainBase.class).list(); + assertThat(domainBases).hasSize(1); + DomainBase domainBase = domainBases.get(0); + // If the registrant is found, then the key instance in domainBase is fixed. + assertThat(auditedOfy().load().key(domainBase.getRegistrant().getOfyKey()).now()).isNotNull(); + } + static CommitLogCheckpoint createCheckpoint(DateTime now) { return CommitLogCheckpoint.create(now, toMap(getBucketIds(), x -> now)); } + static void saveDiffFile(GcsService gcsService, byte[] rawBytes, DateTime timestamp) + throws IOException { + gcsService.createOrReplace( + new GcsFilename(GCS_BUCKET, DIFF_FILE_PREFIX + timestamp), + new GcsFileOptions.Builder() + .addUserMetadata(LOWER_BOUND_CHECKPOINT, timestamp.minusMinutes(1).toString()) + .build(), + ByteBuffer.wrap(rawBytes)); + } + static Iterable saveDiffFile( GcsService gcsService, CommitLogCheckpoint checkpoint, ImmutableObject... entities) throws IOException { @@ -271,12 +303,7 @@ public class RestoreCommitLogsActionTest { for (ImmutableObject entity : allEntities) { serializeEntity(entity, output); } - gcsService.createOrReplace( - new GcsFilename(GCS_BUCKET, DIFF_FILE_PREFIX + now), - new GcsFileOptions.Builder() - .addUserMetadata(LOWER_BOUND_CHECKPOINT, now.minusMinutes(1).toString()) - .build(), - ByteBuffer.wrap(output.toByteArray())); + saveDiffFile(gcsService, output.toByteArray(), now); return allEntities; } diff --git a/core/src/test/resources/google/registry/backup/commitlog.data b/core/src/test/resources/google/registry/backup/commitlog.data new file mode 100644 index 000000000..391e8f5a3 Binary files /dev/null and b/core/src/test/resources/google/registry/backup/commitlog.data differ