Enable sending RPC marshalable objects *back* to their owner

src/StreamJsonRpc/IJsonRpcClientProxyInternal.cs

@@ -25,4 +25,9 @@ internal interface IJsonRpcClientProxyInternal : IJsonRpcClientProxy, IDisposabl
     /// Exceptions thrown from the handler propagate directly to the caller of the proxy, although the RPC request was already transmitted.
     /// </remarks>
     event EventHandler<string> CalledMethod;
+
+    /// <summary>
+    /// Gets the handle of the object that is being marshaled over RPC for which this is a proxy, if applicable.
+    /// </summary>
+    long? MarshaledObjectHandle { get; }
 }

src/StreamJsonRpc/JsonRpc.cs

@@ -725,9 +725,8 @@ public static T Attach<T>(Stream stream)
     public static T Attach<T>(Stream? sendingStream, Stream? receivingStream)
         where T : class
     {
-        TypeInfo proxyType = ProxyGeneration.Get(typeof(T).GetTypeInfo());
         var rpc = new JsonRpc(sendingStream, receivingStream);
-        T proxy = (T)Activator.CreateInstance(proxyType.AsType(), rpc, JsonRpcProxyOptions.Default, /* onDispose: */ null)!;
+        T proxy = rpc.CreateProxy<T>(null, JsonRpcProxyOptions.Default, null);
         rpc.StartListening();
         return proxy;
     }
@@ -762,9 +761,8 @@ public static T Attach<T>(IJsonRpcMessageHandler handler)
     public static T Attach<T>(IJsonRpcMessageHandler handler, JsonRpcProxyOptions? options)
         where T : class
     {
-        TypeInfo proxyType = ProxyGeneration.Get(typeof(T).GetTypeInfo());
         var rpc = new JsonRpc(handler);
-        T proxy = (T)Activator.CreateInstance(proxyType.AsType(), rpc, options ?? JsonRpcProxyOptions.Default, options?.OnDispose)!;
+        T proxy = rpc.CreateProxy<T>(null, options, null)!;
         rpc.StartListening();
         return proxy;
     }
@@ -789,9 +787,7 @@ public T Attach<T>()
     public T Attach<T>(JsonRpcProxyOptions? options)
         where T : class
     {
-        TypeInfo proxyType = ProxyGeneration.Get(typeof(T).GetTypeInfo());
-        T proxy = (T)Activator.CreateInstance(proxyType.AsType(), this, options ?? JsonRpcProxyOptions.Default, options?.OnDispose)!;
-        return proxy;
+        return this.CreateProxy<T>(null, options, null);
     }
 
     /// <summary>
@@ -810,9 +806,7 @@ public T Attach<T>(JsonRpcProxyOptions? options)
     public object Attach(Type interfaceType, JsonRpcProxyOptions? options)
     {
         Requires.NotNull(interfaceType, nameof(interfaceType));
-        TypeInfo proxyType = ProxyGeneration.Get(interfaceType.GetTypeInfo());
-        object proxy = Activator.CreateInstance(proxyType.AsType(), this, options ?? JsonRpcProxyOptions.Default, options?.OnDispose)!;
-        return proxy;
+        return this.CreateProxy(interfaceType.GetTypeInfo(), null, options ?? JsonRpcProxyOptions.Default, null)!;
     }
 
     /// <inheritdoc cref="AddLocalRpcTarget(object, JsonRpcTargetOptions?)"/>
@@ -1207,18 +1201,10 @@ internal static T MarshalLimitedArgument<T>(T marshaledObject)
         throw new NotImplementedException();
     }
 
-    /// <summary>
-    /// Creates a JSON-RPC client proxy that implements a given set of interfaces.
-    /// </summary>
-    /// <param name="contractInterface">The interface that describes the functions available on the remote end.</param>
-    /// <param name="implementedOptionalInterfaces">Additional marshalable interfaces that the client proxy should implement.</param>
-    /// <param name="options">A set of customizations for how the client proxy is wired up. If <see langword="null" />, default options will be used.</param>
-    /// <returns>An instance of the generated proxy.</returns>
-    internal object Attach(Type contractInterface, (TypeInfo Type, int Code)[]? implementedOptionalInterfaces, JsonRpcProxyOptions? options)
+    /// <inheritdoc cref="CreateProxy(TypeInfo, ValueTuple{TypeInfo, int}[], JsonRpcProxyOptions?, long?)"/>
+    internal object Attach(Type contractInterface, (TypeInfo Type, int Code)[]? implementedOptionalInterfaces, JsonRpcProxyOptions? options, long? marshaledObjectHandle)
     {
-        TypeInfo proxyType = ProxyGeneration.Get(contractInterface.GetTypeInfo(), implementedOptionalInterfaces);
-        object proxy = Activator.CreateInstance(proxyType.AsType(), this, options ?? JsonRpcProxyOptions.Default, options?.OnDispose)!;
-        return proxy;
+        return this.CreateProxy(contractInterface.GetTypeInfo(), implementedOptionalInterfaces, options, marshaledObjectHandle);
     }
 
     /// <inheritdoc cref="RpcTargetInfo.AddLocalRpcMethod(MethodInfo, object?, JsonRpcMethodAttribute?, SynchronizationContext?)"/>
@@ -2694,6 +2680,31 @@ private void ThrowIfConfigurationLocked()
         }
     }
 
+    private T CreateProxy<T>((TypeInfo Type, int Code)[]? implementedOptionalInterfaces, JsonRpcProxyOptions? options, long? marshaledObjectHandle)
+        where T : class
+    {
+        return (T)this.CreateProxy(typeof(T).GetTypeInfo(), implementedOptionalInterfaces, options, marshaledObjectHandle);
+    }
+
+    /// <summary>
+    /// Creates a JSON-RPC client proxy that implements a given set of interfaces.
+    /// </summary>
+    /// <param name="contractInterface">The interface that describes the functions available on the remote end.</param>
+    /// <param name="implementedOptionalInterfaces">Additional marshalable interfaces that the client proxy should implement.</param>
+    /// <param name="options">A set of customizations for how the client proxy is wired up. If <see langword="null" />, default options will be used.</param>
+    /// <param name="marshaledObjectHandle">The handle to the remote object that is being marshaled via this proxy.</param>
+    /// <returns>An instance of the generated proxy.</returns>
+    private IJsonRpcClientProxyInternal CreateProxy(TypeInfo contractInterface, (TypeInfo Type, int Code)[]? implementedOptionalInterfaces, JsonRpcProxyOptions? options, long? marshaledObjectHandle)
+    {
+        TypeInfo proxyType = ProxyGeneration.Get(contractInterface, implementedOptionalInterfaces);
+        return (IJsonRpcClientProxyInternal)Activator.CreateInstance(
+            proxyType.AsType(),
+            this,
+            options ?? JsonRpcProxyOptions.Default,
+            marshaledObjectHandle,
+            options?.OnDispose)!;
+    }
+
     /// <summary>
     /// An object that correlates <see cref="JoinableTask"/> tokens within and between <see cref="JsonRpc"/> instances
     /// within a process that does <em>not</em> use <see cref="JoinableTaskFactory"/>,

src/StreamJsonRpc/ProxyGeneration.cs

@@ -108,6 +108,7 @@ internal static TypeInfo Get(TypeInfo contractInterface, (TypeInfo Type, int Cod
             FieldBuilder disposedField = proxyTypeBuilder.DefineField("disposed", typeof(bool), FieldAttributes.Private);
             FieldBuilder callingMethodField = proxyTypeBuilder.DefineField("callingMethod", typeof(EventHandler<string>), FieldAttributes.Private);
             FieldBuilder calledMethodField = proxyTypeBuilder.DefineField("calledMethod", typeof(EventHandler<string>), FieldAttributes.Private);
+            FieldBuilder marshaledObjectHandleField = proxyTypeBuilder.DefineField("marshaledObjectHandle", typeof(long?), fieldAttributes);
 
             VerifySupported(!FindAllOnThisAndOtherInterfaces(contractInterface, i => i.DeclaredProperties).Any(), Resources.UnsupportedPropertiesOnClientProxyInterface, contractInterface);
 
@@ -168,12 +169,12 @@ internal static TypeInfo Get(TypeInfo contractInterface, (TypeInfo Type, int Cod
                 }));
             }
 
-            // .ctor(JsonRpc, JsonRpcProxyOptions, Action onDispose)
+            // .ctor(JsonRpc, JsonRpcProxyOptions, long? marshaledObjectHandle, Action onDispose)
             {
                 ConstructorBuilder ctor = proxyTypeBuilder.DefineConstructor(
                     MethodAttributes.Public | MethodAttributes.HideBySig | MethodAttributes.SpecialName | MethodAttributes.RTSpecialName,
                     CallingConventions.Standard,
-                    new Type[] { typeof(JsonRpc), typeof(JsonRpcProxyOptions), typeof(Action) });
+                    [typeof(JsonRpc), typeof(JsonRpcProxyOptions), typeof(long?), typeof(Action)]);
                 ILGenerator il = ctor.GetILGenerator();
 
                 // : base()
@@ -190,9 +191,14 @@ internal static TypeInfo Get(TypeInfo contractInterface, (TypeInfo Type, int Cod
                 il.Emit(OpCodes.Ldarg_2);
                 il.Emit(OpCodes.Stfld, optionsField);
 
-                // this.onDispose = onDispose;
+                // this.marshaledObjectHandle = marshaledObjectHandle;
                 il.Emit(OpCodes.Ldarg_0);
                 il.Emit(OpCodes.Ldarg_3);
+                il.Emit(OpCodes.Stfld, marshaledObjectHandleField);
+
+                // this.onDispose = onDispose;
+                il.Emit(OpCodes.Ldarg_0);
+                il.Emit(OpCodes.Ldarg, 4);
                 il.Emit(OpCodes.Stfld, onDisposeField);
 
                 // Emit IL that supports events.
@@ -206,7 +212,7 @@ internal static TypeInfo Get(TypeInfo contractInterface, (TypeInfo Type, int Cod
 
             ImplementDisposeMethod(proxyTypeBuilder, jsonRpcField, onDisposeField, disposedField);
             ImplementIsDisposedProperty(proxyTypeBuilder, jsonRpcField, disposedField);
-            ImplementIJsonRpcClientProxyInternal(proxyTypeBuilder, callingMethodField, calledMethodField);
+            ImplementIJsonRpcClientProxyInternal(proxyTypeBuilder, callingMethodField, calledMethodField, marshaledObjectHandleField);
 
             // IJsonRpcClientProxy.JsonRpc property
             {
@@ -486,7 +492,7 @@ private static void EmitRaiseCallEvent(ILGenerator il, FieldBuilder eventHandler
         il.MarkLabel(endOfSubroutine);
     }
 
-    private static void ImplementIJsonRpcClientProxyInternal(TypeBuilder proxyTypeBuilder, FieldBuilder callingMethodField, FieldBuilder calledMethodField)
+    private static void ImplementIJsonRpcClientProxyInternal(TypeBuilder proxyTypeBuilder, FieldBuilder callingMethodField, FieldBuilder calledMethodField, FieldBuilder marshaledObjectHandleField)
     {
         void AddEvent(FieldBuilder evtField, string eventName)
         {
@@ -509,6 +515,25 @@ void AddEvent(FieldBuilder evtField, string eventName)
 
         AddEvent(callingMethodField, nameof(IJsonRpcClientProxyInternal.CallingMethod));
         AddEvent(calledMethodField, nameof(IJsonRpcClientProxyInternal.CalledMethod));
+
+        // Implement the IJsonRpcClientProxyInternal.MarshaledObjectHandle property.
+        PropertyBuilder marshaledObjectHandleProperty = proxyTypeBuilder.DefineProperty(
+            $"{nameof(IJsonRpcClientProxyInternal)}.{nameof(IJsonRpcClientProxyInternal.MarshaledObjectHandle)}",
+            PropertyAttributes.None,
+            typeof(long?),
+            Type.EmptyTypes);
+        MethodBuilder marshaledObjectHandleGetter = proxyTypeBuilder.DefineMethod(
+            $"get_{nameof(IJsonRpcClientProxyInternal.MarshaledObjectHandle)}",
+            MethodAttributes.Private | MethodAttributes.Final | MethodAttributes.HideBySig | MethodAttributes.SpecialName | MethodAttributes.NewSlot | MethodAttributes.Virtual,
+            typeof(long?),
+            Type.EmptyTypes);
+        ILGenerator il = marshaledObjectHandleGetter.GetILGenerator();
+        il.Emit(OpCodes.Ldarg_0);
+        il.Emit(OpCodes.Ldfld, marshaledObjectHandleField);
+        il.Emit(OpCodes.Ret);
+
+        proxyTypeBuilder.DefineMethodOverride(marshaledObjectHandleGetter, typeof(IJsonRpcClientProxyInternal).GetTypeInfo().GetDeclaredProperty(nameof(IJsonRpcClientProxyInternal.MarshaledObjectHandle))!.GetMethod!);
+        marshaledObjectHandleProperty.SetGetMethod(marshaledObjectHandleGetter);
     }
 
     private static void LoadParameterTypeArrayField(TypeBuilder proxyTypeBuilder, ParameterInfo[] parameterInfos, ILGenerator il)

src/StreamJsonRpc/Reflection/MessageFormatterRpcMarshaledContextTracker.cs

@@ -5,6 +5,7 @@
 using System.Collections.Immutable;
 using System.Diagnostics.CodeAnalysis;
 using System.Globalization;
+using System.Net;
 using System.Reflection;
 using System.Runtime.Serialization;
 using Microsoft.VisualStudio.Threading;
@@ -204,6 +205,23 @@ internal MarshalToken GetToken(object marshaledObject, JsonRpcTargetOptions opti
             throw new NotSupportedException(Resources.CallScopedMarshaledObjectInReturnValueNotAllowed);
         }
 
+        if (marshaledObject is IJsonRpcClientProxyInternal { MarshaledObjectHandle: not null } proxy)
+        {
+            // Supporting passing of a marshaled object over RPC requires that we:
+            // 1. Distinguish passing it back to its original owner vs. a 3rd party over an independent RPC connection.
+            // 2. If back to the original owner, we need to reuse the same handle and pass other data so the receiver recognizes this case.
+            if (proxy.JsonRpc != this.jsonRpc)
+            {
+                throw new NotSupportedException("Forwarding an RPC marshaled object to a 3rd party is not supported.");
+            }
+
+            return new MarshalToken
+            {
+                Handle = proxy.MarshaledObjectHandle.Value,
+                Marshaled = (int)MarshalMode.MarshallingProxyBackToOwner,
+            };
+        }
+
         long handle = this.nextUniqueHandle++;
 
         IRpcMarshaledContext<object> context = JsonRpc.MarshalWithControlledLifetime(declaredType, marshaledObject, options);
@@ -276,7 +294,15 @@ internal MarshalToken GetToken(object marshaledObject, JsonRpcTargetOptions opti
 
         if ((MarshalMode)token.Value.Marshaled == MarshalMode.MarshallingProxyBackToOwner)
         {
-            throw new NotSupportedException("Receiving marshaled objects back to the owner is not yet supported.");
+            lock (this.marshaledObjects)
+            {
+                if (this.marshaledObjects.TryGetValue(token.Value.Handle, out (IRpcMarshaledContext<object> Context, IDisposable Revert) marshaled))
+                {
+                    return marshaled.Context.Proxy;
+                }
+            }
+
+            throw new ProtocolViolationException("Marshaled object \"returned\" with an unrecognized handle.");
         }
 
         RpcMarshalableAttribute synthesizedAttribute = new()
@@ -318,7 +344,8 @@ internal MarshalToken GetToken(object marshaledObject, JsonRpcTargetOptions opti
 
                     this.jsonRpc.NotifyWithParameterObjectAsync("$/releaseMarshaledObject", new { handle = token.Value.Handle, ownedBySender = false }).Forget();
                 },
-            });
+            },
+            token.Value.Handle);
         if (options.OnProxyConstructed is object)
         {
             options.OnProxyConstructed((IJsonRpcClientProxyInternal)result);

src/StreamJsonRpc/RpcMarshaledContext.cs

@@ -14,13 +14,8 @@ internal class RpcMarshaledContext<T> : IRpcMarshaledContext<T>
     /// <param name="options">The <see cref="JsonRpcTargetOptions"/> to use when adding this object as an RPC target.</param>
     internal RpcMarshaledContext(T value, JsonRpcTargetOptions options)
     {
-        if (value is IJsonRpcClientProxy)
-        {
-            // Supporting passing of a marshaled object over RPC requires that we:
-            // 1. Distinguish passing it back to its original owner vs. a 3rd party over an independent RPC connection.
-            // 2. If back to the original owner, we need to reuse the same handle and pass other data so the receiver recognizes this case.
-            throw new NotSupportedException("Marshaling a proxy back to its owner ");
-        }
+        // We shouldn't reach this point with a proxy.
+        Requires.Argument(value is not IJsonRpcClientProxyInternal, nameof(value), "Cannot marshal a proxy.");
 
         this.Proxy = value;
         this.JsonRpcTargetOptions = options;

test/StreamJsonRpc.Tests/DisposableProxyTests.cs

@@ -141,8 +141,8 @@ public async Task IDisposable_MarshaledBackAndForth()
     {
         IDisposable? disposable = await this.client.GetDisposableAsync().WithCancellation(this.TimeoutToken);
         Assert.NotNull(disposable);
-        var ex = await Assert.ThrowsAnyAsync<Exception>(() => this.client.AcceptProxyAsync(disposable!)).WithCancellation(this.TimeoutToken);
-        Assert.True(IsExceptionOrInnerOfType<NotSupportedException>(ex));
+        await this.client.AcceptProxyAsync(disposable).WithCancellation(this.TimeoutToken);
+        Assert.Same(this.server.ReturnedDisposable, this.server.ReceivedProxy);
     }
 
     protected abstract IJsonRpcMessageFormatter CreateFormatter();

test/StreamJsonRpc.Tests/MarshalableProxyTests.cs

@@ -588,7 +588,31 @@ public async Task IMarshalable_MarshaledBackAndForth()
     {
         IMarshalable? proxyMarshalable = await this.client.GetMarshalableAsync().WithCancellation(this.TimeoutToken);
         Assert.NotNull(proxyMarshalable);
-        var ex = await Assert.ThrowsAnyAsync<Exception>(() => this.client.AcceptProxyAsync(proxyMarshalable!)).WithCancellation(this.TimeoutToken);
+        await this.client.AcceptProxyAsync(proxyMarshalable).WithCancellation(this.TimeoutToken);
+        Assert.Same(this.server.ReturnedMarshalable, this.server.ReceivedProxy);
+    }
+
+    [Fact]
+    public async Task IMarshalable_MarshaledAndForwarded()
+    {
+        IMarshalable? proxyMarshalable = await this.client.GetMarshalableAsync().WithCancellation(this.TimeoutToken);
+        Assert.NotNull(proxyMarshalable);
+
+        // Try to send the proxy to a *different* server. This should fail.
+        var pipes = FullDuplexStream.CreatePipePair();
+        IServer client2 = JsonRpc.Attach<IServer>(new LengthHeaderMessageHandler(pipes.Item1, this.CreateFormatter()));
+        JsonRpc clientRpc2 = ((IJsonRpcClientProxy)this.client).JsonRpc;
+        Server server2 = new();
+        JsonRpc serverRpc2 = new JsonRpc(new LengthHeaderMessageHandler(pipes.Item2, this.CreateFormatter()));
+        serverRpc2.AddLocalRpcTarget(server2);
+        serverRpc2.TraceSource = new TraceSource("Server2", SourceLevels.Verbose);
+        clientRpc2.TraceSource = new TraceSource("Client2", SourceLevels.Verbose);
+        serverRpc2.TraceSource.Listeners.Add(new XunitTraceListener(this.Logger));
+        clientRpc2.TraceSource.Listeners.Add(new XunitTraceListener(this.Logger));
+        serverRpc2.StartListening();
+
+        Exception ex = await Assert.ThrowsAnyAsync<Exception>(() => client2.AcceptProxyAsync(proxyMarshalable)).WithCancellation(this.TimeoutToken);
+        this.Logger.WriteLine("Received exception: {0}", ex);
         Assert.True(IsExceptionOrInnerOfType<NotSupportedException>(ex));
     }