Use .ContinueWith instead of finally for removal.

src/Autofac.Integration.WebApi.Owin/DependencyScopeHandler.cs

@@ -42,13 +42,19 @@ protected override Task<HttpResponseMessage> SendAsync(HttpRequestMessage reques
         var dependencyScope = new AutofacWebApiDependencyScope(lifetimeScope);
         request.Properties[HttpPropertyKeys.DependencyScope] = dependencyScope;
 
-        try
-        {
-            return base.SendAsync(request, cancellationToken);
-        }
-        finally
-        {
-            request.Properties.Remove(HttpPropertyKeys.DependencyScope);
-        }
+        // Using .ContinueWith instead of try/finally because you can't
+        // have async/await in [SecurityCritical] or [SecuritySafeCritical]
+        // marked code.
+        return base
+            .SendAsync(request, cancellationToken)
+            .ContinueWith(
+                task =>
+                {
+                    request.Properties.Remove(HttpPropertyKeys.DependencyScope);
+                    return task.Result;
+                },
+                CancellationToken.None,
+                TaskContinuationOptions.ExecuteSynchronously,
+                TaskScheduler.FromCurrentSynchronizationContext());
     }
 }

test/Autofac.Integration.WebApi.Owin.Test/DependencyScopeHandlerFixture.cs

@@ -47,7 +47,7 @@ public async void AddsAutofacDependencyScopeToHttpRequestMessage()
     }
 
     [Fact]
-    public async void RemoveAutofacDependencyScopeFromHttpRequestMessageAfterLeaving()
+    public void RemoveAutofacDependencyScopeAfterTaskExecutes_Task()
     {
         var request = new HttpRequestMessage();
         var context = new OwinContext();
@@ -56,22 +56,115 @@ public async void RemoveAutofacDependencyScopeFromHttpRequestMessageAfterLeaving
         var container = new ContainerBuilder().Build();
         context.Set(Constants.OwinLifetimeScopeKey, container);
 
-        var fakeHandler = new FakeInnerHandler(_ => new HttpResponseMessage(HttpStatusCode.OK));
+        var flag = new AutoResetEvent(false);
+        var fakeHandler = new FakeInnerHandler(_ =>
+        {
+            return Task.Factory.StartNew(() =>
+            {
+                flag.WaitOne();
+                Assert.Contains(HttpPropertyKeys.DependencyScope, request.Properties);
+                return new HttpResponseMessage(HttpStatusCode.OK);
+            });
+        });
         var handler = new DependencyScopeHandler { InnerHandler = fakeHandler };
         var invoker = new HttpMessageInvoker(handler);
-        await invoker.SendAsync(request, CancellationToken.None);
+        var task = invoker.SendAsync(request, CancellationToken.None);
+
+        Assert.Contains(HttpPropertyKeys.DependencyScope, request.Properties);
+        flag.Set();
+        task.GetAwaiter().GetResult();
+        Assert.DoesNotContain(HttpPropertyKeys.DependencyScope, request.Properties);
+    }
+
+    [Fact]
+    public async Task RemoveAutofacDependencyScopeAfterTaskExecutes_AsyncAwait()
+    {
+        var request = new HttpRequestMessage();
+        var context = new OwinContext();
+        request.Properties.Add("MS_OwinContext", context);
+
+        var container = new ContainerBuilder().Build();
+        context.Set(Constants.OwinLifetimeScopeKey, container);
+
+        var flag = new AutoResetEvent(false);
+        var fakeHandler = new FakeInnerHandler(async _ => await Task.Factory.StartNew(() =>
+        {
+            flag.WaitOne();
+            Assert.Contains(HttpPropertyKeys.DependencyScope, request.Properties);
+            return new HttpResponseMessage(HttpStatusCode.OK);
+        }));
+
+        var handler = new DependencyScopeHandler { InnerHandler = fakeHandler };
+        var invoker = new HttpMessageInvoker(handler);
+        var task = invoker.SendAsync(request, CancellationToken.None);
+
+        Assert.Contains(HttpPropertyKeys.DependencyScope, request.Properties);
+        flag.Set();
+        await task;
+        Assert.DoesNotContain(HttpPropertyKeys.DependencyScope, request.Properties);
+    }
+
+    [Fact]
+    public void RemoveAutofacDependencyScopeEvenIfTaskThrows_Task()
+    {
+        var request = new HttpRequestMessage();
+        var context = new OwinContext();
+        request.Properties.Add("MS_OwinContext", context);
+
+        var container = new ContainerBuilder().Build();
+        context.Set(Constants.OwinLifetimeScopeKey, container);
+
+        var fakeHandler = new FakeInnerHandler(_ =>
+        {
+            return Task.Factory.StartNew<HttpResponseMessage>(() =>
+            {
+                throw new DivideByZeroException();
+            });
+        });
+        var handler = new DependencyScopeHandler { InnerHandler = fakeHandler };
+        var invoker = new HttpMessageInvoker(handler);
+        var task = invoker.SendAsync(request, CancellationToken.None);
+
+        Assert.Contains(HttpPropertyKeys.DependencyScope, request.Properties);
+        Assert.Throws<AggregateException>(() => task.GetAwaiter().GetResult());
+        Assert.DoesNotContain(HttpPropertyKeys.DependencyScope, request.Properties);
+    }
+
+    [Fact]
+    public async Task RemoveAutofacDependencyScopeEvenIfTaskThrows_AsyncAwait()
+    {
+        var request = new HttpRequestMessage();
+        var context = new OwinContext();
+        request.Properties.Add("MS_OwinContext", context);
+
+        var container = new ContainerBuilder().Build();
+        context.Set(Constants.OwinLifetimeScopeKey, container);
 
+        var fakeHandler = new FakeInnerHandler(async _ => await Task.Factory.StartNew<HttpResponseMessage>(() =>
+        {
+            throw new DivideByZeroException();
+        }));
+        var handler = new DependencyScopeHandler { InnerHandler = fakeHandler };
+        var invoker = new HttpMessageInvoker(handler);
+        var task = invoker.SendAsync(request, CancellationToken.None);
+
+        Assert.Contains(HttpPropertyKeys.DependencyScope, request.Properties);
+        await Assert.ThrowsAsync<AggregateException>(() => task);
         Assert.DoesNotContain(HttpPropertyKeys.DependencyScope, request.Properties);
     }
 
     private class FakeInnerHandler : HttpMessageHandler
     {
-        private readonly Func<HttpRequestMessage, HttpResponseMessage> _delegate;
+        private readonly Func<HttpRequestMessage, Task<HttpResponseMessage>> _delegate;
 
         public FakeInnerHandler(Func<HttpRequestMessage, HttpResponseMessage> @delegate)
-            => _delegate = @delegate;
+            : this(request => Task.FromResult(@delegate(request)))
+        {
+        }
+
+        public FakeInnerHandler(Func<HttpRequestMessage, Task<HttpResponseMessage>> @delegate) => _delegate = @delegate;
 
         protected override Task<HttpResponseMessage> SendAsync(HttpRequestMessage request, CancellationToken cancellationToken)
-            => Task.FromResult(_delegate(request));
+            => _delegate(request);
     }
 }