diff --git a/src/dorkbox/network/connection/KryoExtra.java b/src/dorkbox/network/connection/KryoExtra.java index aeb7dc3a..f274122b 100644 --- a/src/dorkbox/network/connection/KryoExtra.java +++ b/src/dorkbox/network/connection/KryoExtra.java @@ -22,10 +22,13 @@ import org.bouncycastle.crypto.modes.GCMBlockCipher; import org.bouncycastle.crypto.params.ParametersWithIV; import com.esotericsoftware.kryo.Kryo; +import com.esotericsoftware.kryo.util.DefaultStreamFactory; +import com.esotericsoftware.kryo.util.MapReferenceResolver; import dorkbox.network.pipeline.ByteBufInput; import dorkbox.network.pipeline.ByteBufOutput; import dorkbox.network.serialization.CryptoSerializationManager; +import dorkbox.network.serialization.EditableDefaultClassResolver; import dorkbox.util.bytes.BigEndian; import dorkbox.util.bytes.OptimizeUtilsByteArray; import dorkbox.util.bytes.OptimizeUtilsByteBuf; @@ -85,6 +88,8 @@ class KryoExtra extends Kryo { public KryoExtra(final CryptoSerializationManager serializationManager) { + super(new EditableDefaultClassResolver(), new MapReferenceResolver(), new DefaultStreamFactory()); + this.serializationManager = serializationManager; } diff --git a/src/dorkbox/network/serialization/EditableDefaultClassResolver.java b/src/dorkbox/network/serialization/EditableDefaultClassResolver.java new file mode 100644 index 00000000..5da48053 --- /dev/null +++ b/src/dorkbox/network/serialization/EditableDefaultClassResolver.java @@ -0,0 +1,36 @@ +/* + * Copyright 2018 dorkbox, llc + * + * 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 dorkbox.network.serialization; + +import com.esotericsoftware.kryo.util.DefaultClassResolver; +import com.esotericsoftware.kryo.util.Util; + +public +class EditableDefaultClassResolver extends DefaultClassResolver { + public + EditableDefaultClassResolver() { + } + + public + void deleteRegistrationType(Class type) { + if (type.isPrimitive()) { + classToRegistration.remove(Util.getWrapperClass(type)); + } + else { + classToRegistration.remove(type); + } + } +} diff --git a/src/dorkbox/network/serialization/Serialization.java b/src/dorkbox/network/serialization/Serialization.java index d5d2e084..5efc0839 100644 --- a/src/dorkbox/network/serialization/Serialization.java +++ b/src/dorkbox/network/serialization/Serialization.java @@ -31,6 +31,7 @@ import org.slf4j.Logger; import com.esotericsoftware.kryo.ClassResolver; import com.esotericsoftware.kryo.Kryo; import com.esotericsoftware.kryo.KryoException; +import com.esotericsoftware.kryo.Registration; import com.esotericsoftware.kryo.Serializer; import com.esotericsoftware.kryo.factories.ReflectionSerializerFactory; import com.esotericsoftware.kryo.factories.SerializerFactory; @@ -381,9 +382,25 @@ class Serialization implements CryptoSerializationMa // the server will ONLY send this object to the client, where on the client it becomes the proxy/interface. RemoteImplClass remoteImplClass = (RemoteImplClass) clazz; - // registers the implementation, so that when it is WRITTEN, it becomes a "magic" proxy object - int id = kryo.register(remoteImplClass.implClass, remoteObjectSerializer) - .getId(); + int id; + + // check to see if the interface is already registered. If so, we override it with the implementation class + EditableDefaultClassResolver classResolver = (EditableDefaultClassResolver) kryo.getClassResolver(); + Registration registration = classResolver.getRegistration(remoteImplClass.ifaceClass); + if (registration != null) { + // override it with the implementation + int oldId = registration.getId(); + id = kryo.register(remoteImplClass.implClass, remoteObjectSerializer, oldId).getId(); + + // delete the old one (since it shouldn't be stored anywhere) + classResolver.deleteRegistrationType(remoteImplClass.ifaceClass); + } + else { + // we have a new registration. + + // registers the implementation, so that when it is WRITTEN, it becomes a "magic" proxy object + id = kryo.register(remoteImplClass.implClass, remoteObjectSerializer).getId(); + } // sets up the RMI, so when we receive the iface class from the client, we know what impl to use // if this is over-written, we don't care.