PayloadAttribute.cs 4.92 KB
Newer Older
Maxim Bogdanov's avatar
Maxim Bogdanov committed
1 2 3 4
using System;
using System.Collections.Generic;
using System.Linq;
using System.Reflection;
5
using Stratis.Bitcoin.Utilities;
Maxim Bogdanov's avatar
Maxim Bogdanov committed
6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25

namespace Stratis.Bitcoin.P2P.Protocol.Payloads
{
    /// <summary>
    /// An attribute that enables mapping between command names and P2P netowrk types.
    /// </summary>
    [AttributeUsage(AttributeTargets.Class)]
    public class PayloadAttribute : Attribute
    {
        /// <summary>
        /// The command name.
        /// </summary>
        public string Name { get; set; }

        /// <summary>
        /// Initialize a new instance of the object.
        /// </summary>
        /// <param name="commandName"></param>
        public PayloadAttribute(string commandName)
        {
26 27 28
            if (commandName.Length > 12)
                throw new ArgumentException("Protocol violation: command name is limited to 12 characters.");
            
Maxim Bogdanov's avatar
Maxim Bogdanov committed
29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93
            this.Name = commandName;
        }
    }

    /// <summary>
    /// A provider that maps <see cref="PayloadAttribute"/> types with <see cref="Message.Command"/>.
    /// This is used by the P2P code to map and deserialize messages that are received from the tcp network to a concrete type.
    /// </summary>
    public class PayloadProvider
    {
        /// <summary>
        /// A mapping between the command name and the payload type.
        /// </summary>
        private readonly Dictionary<string, Type> nameToType;

        /// <summary>
        /// A mapping between the payload type and the command name.
        /// </summary>
        private readonly Dictionary<Type, string> typeToName;

        /// <summary>
        /// Initialize a new instance of the object.
        /// </summary>
        public PayloadProvider()
        {
            this.nameToType = new Dictionary<string, Type>();
            this.typeToName = new Dictionary<Type, string>();
        }

        /// <summary>
        /// Discover all payloads from the provided assembly, if no assembly is provided defaults to <see cref="PayloadAttribute"/>.
        /// </summary>
        /// <param name="assembly">The assembly to discover from or <see cref="PayloadAttribute"/> if <c>null</c>.</param>
        public PayloadProvider DiscoverPayloads(Assembly assembly = null)
        {
            assembly = assembly ?? typeof(PayloadAttribute).GetTypeInfo().Assembly;

            IEnumerable<TypeInfo> types = null;
                
            try
            {
                types = assembly.DefinedTypes;
            }
            catch (ReflectionTypeLoadException e)
            {
                types = e.Types.Where(t => t != null).Select(t => t.GetTypeInfo());
            }

            foreach (var pair in types
                .Where(t => t.Namespace == typeof(PayloadAttribute).Namespace)
                .Where(t => t.IsDefined(typeof(PayloadAttribute), true))
                .Select(t =>
                    new
                    {
                        Attr = t.GetCustomAttributes(typeof(PayloadAttribute), true).OfType<PayloadAttribute>().First(),
                        Type = t
                    }))
            {
                this.nameToType.Add(pair.Attr.Name, pair.Type.AsType());
                this.typeToName.Add(pair.Type.AsType(), pair.Attr.Name);
            }

            return this;
        }

94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110
        /// <summary>
        /// Add a payload to the Provider by specifying its type.
        /// </summary>
        /// <param name="type">The type to payload to add.  Must derive from <see cref="Payload"/>.</param>
        public void AddPayload(Type payloadType)
        {
            Guard.NotNull(payloadType, nameof(payloadType));
            Guard.Assert(payloadType.IsSubclassOf(typeof(Payload)));

            PayloadAttribute payloadAttribute = payloadType.GetCustomAttributes(typeof(PayloadAttribute), true)
                .OfType<PayloadAttribute>().First();
            Guard.Assert(payloadAttribute != null);

            this.nameToType.Add(payloadAttribute.Name, payloadType);
            this.typeToName.Add(payloadType, payloadAttribute.Name);
        }
        
Maxim Bogdanov's avatar
Maxim Bogdanov committed
111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133
        /// <summary>
        /// Get the <see cref="Payload"/> type associated with the command name.
        /// </summary>
        /// <param name="commandName">The command name.</param>
        /// <returns>The type of payload the command is associated with.</returns>
        public Type GetCommandType(string commandName)
        {
            if (!this.nameToType.TryGetValue(commandName, out Type result))
                return typeof(UnknowPayload);

            return result;
        }

        /// <summary>
        /// Check that a <see cref="Payload"/> type is allowed to be used in the P2P code.
        /// </summary>
        /// <param name="type">A type that represents a <see cref="Payload"/></param>
        /// <returns>True if the type is registered as a usable payload.</returns>
        public bool IsPayloadRegistered(Type type)
        {
            return this.typeToName.ContainsKey(type);
        }
    }
134
}