Wcf and CallContext

WcfContextTest.cs

using System;
using System.Runtime.Remoting.Messaging;
using System.Runtime.Serialization;
using System.ServiceModel;
using System.ServiceModel.Channels;
using System.ServiceModel.Description;
using System.ServiceModel.Dispatcher;

namespace WcfContextTest
{
    class Program
    {
        static void Main(string[] args)
        {
            const string address = "net.tcp://127.0.0.1:1234/";
            ServiceHost host = new ServiceHost(typeof(TestImpl));
            host.AddServiceEndpoint(typeof(ITest), new NetTcpBinding(), address);
            host.Open();

            var cf = new ChannelFactory<ITest>(new NetTcpBinding(), new EndpointAddress(address));
            foreach (var op in cf.Endpoint.Contract.Operations)
                op.OperationBehaviors.Add(new ContextInjectBehavior());
            var c = cf.CreateChannel();

            Console.WriteLine("Without context:");
            Console.WriteLine("Response is: {0}", c.Test(new TestReq { Id = 1 }).Name);
            Console.WriteLine("Context is: {0}", CallContext.LogicalGetData("WcfContext") ?? "(null)");

            Console.WriteLine("With context: Hello");
            CallContext.LogicalSetData("WcfContext", "Hello");
            Console.WriteLine("Response is: {0}", c.Test(new TestReq { Id = 1 }).Name);
            Console.WriteLine("Context is: {0}", CallContext.LogicalGetData("WcfContext") ?? "(null)");

            ((ICommunicationObject)c).Close();
            host.Close();
        }
    }

    [AttributeUsage(AttributeTargets.Class, Inherited = false, AllowMultiple = false)]
    public sealed class ContextInjectAttribute : Attribute, IServiceBehavior
    {
        public void AddBindingParameters(ServiceDescription serviceDescription, ServiceHostBase serviceHostBase, System.Collections.ObjectModel.Collection<ServiceEndpoint> endpoints, BindingParameterCollection bindingParameters) { }

        public void ApplyDispatchBehavior(ServiceDescription serviceDescription, ServiceHostBase serviceHostBase)
        {
            foreach (var ep in  serviceDescription.Endpoints)
                foreach (var op in ep.Contract.Operations)
                    op.OperationBehaviors.Add(new ContextInjectBehavior());
        }

        public void Validate(ServiceDescription serviceDescription, ServiceHostBase serviceHostBase) { }
    }

    public sealed class ContextInjectBehavior : IOperationBehavior
    {
        public void AddBindingParameters(OperationDescription operationDescription, BindingParameterCollection bindingParameters) { }

        public void ApplyClientBehavior(OperationDescription operationDescription, ClientOperation clientOperation)
        {
            clientOperation.Formatter = new MyClientMessageFormatter(clientOperation.Formatter);
        }

        public void ApplyDispatchBehavior(OperationDescription operationDescription, DispatchOperation dispatchOperation)
        {
            dispatchOperation.Formatter = new MyDispatchMessageFormatter(dispatchOperation.Formatter);
        }

        public void Validate(OperationDescription operationDescription) { }

    }

    public class MyDispatchMessageFormatter : IDispatchMessageFormatter
    {

        private readonly IDispatchMessageFormatter m_inner;

        public MyDispatchMessageFormatter(IDispatchMessageFormatter inner)
        {
            m_inner = inner;
        }

        public void DeserializeRequest(Message message, object[] parameters)
        {
            int index = message.Headers.FindHeader("context", "urn:context");
            if (index != -1)
            {
                var reader = message.Headers.GetReaderAtHeader(index);
                CallContext.LogicalSetData("WcfContext", reader.ReadString());
            }
            m_inner.DeserializeRequest(message, parameters);
        }

        public Message SerializeReply(MessageVersion messageVersion, object[] parameters, object result)
        {
            var msg = m_inner.SerializeReply(messageVersion, parameters, result);
            var context = CallContext.LogicalGetData("WcfContext") as string;
            if (context != null)
            {
                msg.Headers.Add(MessageHeader.CreateHeader("context", "urn:context", context, false));
                CallContext.FreeNamedDataSlot("WcfContext");
            }
            return msg;
        }
    }

    public class MyClientMessageFormatter : IClientMessageFormatter
    {
        private readonly IClientMessageFormatter m_inner;

        public MyClientMessageFormatter(IClientMessageFormatter inner)
        {
            m_inner = inner;
        }

        public object DeserializeReply(Message message, object[] parameters)
        {
            var result = m_inner.DeserializeReply(message, parameters);
            int index = message.Headers.FindHeader("context", "urn:context");
            if (index != -1)
            {
                var reader = message.Headers.GetReaderAtHeader(index);
                CallContext.LogicalSetData("WcfContext", reader.ReadString());
            }
            return result;
        }

        public Message SerializeRequest(MessageVersion messageVersion, object[] parameters)
        {
            var msg = m_inner.SerializeRequest(messageVersion, parameters);
            var context = CallContext.LogicalGetData("WcfContext") as string;
            if (context != null)
                msg.Headers.Add(MessageHeader.CreateHeader("context", "urn:context", context, false));
            return msg;
        }
    }

    [ServiceContract]
    public interface ITest
    {
        [OperationContract]
        TestResp Test(TestReq req);
    }

    [DataContract]
    public class TestReq
    {
        [DataMember]
        public int Id { get; set; }
    }

    [DataContract]
    public class TestResp
    {
        [DataMember]
        public string Name { get; set; }
    }

    [ContextInject]
    public class TestImpl
        : ITest
    {
        public TestResp Test(TestReq req)
        {
            var context = CallContext.LogicalGetData("WcfContext") as string;
            var result = new TestResp { Name = context + req.Id.ToString() };
            if (context != null)
            {
                context = "great!";
                CallContext.LogicalSetData("WcfContext", context);
            }
            return result;
        }
    }

}