.NET Core 1.0 Compatible (Incomplete) Implementation of WebSocket

WebSocket.cs

using System;
using System.IO;
using System.Text;
using System.Net.Sockets;
using System.Net.Security;
using System.Threading.Tasks;

namespace Protty
{

  /// <summary>
  /// Incomplete websocket implementation for .NET Core
  /// </summary>
  public class WebSocket : IDisposable
  {

    private Uri uri;                // Connection URI info
    private TcpClient client;       // Connection object
    private SslStream sslStream;    // SSL Stream Wrapper
    private bool disposed = false;  // Object State

    // Max Read Buffer
    public static int ReadBufferSize = 4096;
    // Random Bytes Generator
    private static readonly Random Random = new Random();

    /// <summary>
    /// Websocket OpCodes in Parsing
    /// </summary>
    public enum OP
    {
      Continue = 0x00, // Packet is a continuation
      Text     = 0x01, // Packet is UTF-8 Text data
      Binary   = 0x02, // Packet is raw octet byte stream
      Close    = 0x80, // Packet is requesting close
      Ping     = 0x09, // Packet is a Ping
      Pong     = 0x0a  // Packet is a Pong
    };

    /// <summary>
    /// Websocket Frame Object
    /// </summary>
    public struct Frame
    {
      public OP OpCode;   // the websocket opcode
      public bool Fin;    // the websocket fin flag
      public bool Masked; // the websocket masked flag
      public byte[] Data; // the websocket contents
    };

    /// <summary>
    /// Create a new Websocket Connection Object
    /// </summary>
    public WebSocket()
    {
      sslStream = null;
      client = new TcpClient();
      client.NoDelay = true;
    }

    /// <summary>
    /// Generate Key for websocket handshake
    /// </summary>
    /// <returns>the generated key</returns>
    private string GenerateKey()
    {
      byte[] keyData = new byte[16];
      Random.NextBytes(keyData);
      return Convert.ToBase64String(keyData);
    }

    /// <summary>
    /// Get the internal Network stream to Read/Write from
    /// </summary>
    /// <returns>the internal network stream</returns>
    private Stream GetStream()
    {
      return sslStream != null ?
        (Stream)sslStream : client.GetStream();
    }

    /// <summary>
    /// Perform simple websocket handshake
    /// </summary>
    /// <returns></returns>
    private async Task HandShake()
    {
      // Send the Http request to upgrade to websocket
      await WriteAsync(String.Join("\r\n", new String[] {
        $"GET {uri.PathAndQuery} HTTP/1.1",
        $"Host: {uri.Host}:{uri.Port}",
        "Upgrade: WebSocket",
        "Connection: Upgrade",
        $"Sec-WebSocket-Key: {GenerateKey()}",
        "Sec-WebSocket-Version: 13",
        "\r\n"
      }));

      // Get response and test if Ok to upgrade
      string response = Encoding.UTF8.GetString(await ReadAsync());
      if (response.Split('\n')[0].Split(' ')[1] != "101")
        throw new Exception("Failed to perform handshake!");
    }

    /// <summary>
    /// Connect to the url provided
    /// </summary>
    /// <param name="url">the url to connect to</param>
    /// <returns>self instance</returns>
    public async Task<WebSocket> ConnectAsync(string url)
    {
      uri = new Uri(url);
      await client.ConnectAsync(uri.Host, uri.Port);
      if (uri.Scheme.EndsWith("s")) {
        sslStream = new SslStream(client.GetStream());
        await sslStream.AuthenticateAsClientAsync(uri.Host);
      }
      await HandShake();
      return this;
    }

    /// <summary>
    /// String wrapper around <see cref="WriteAsync(byte[])"/>
    /// </summary>
    /// <param name="data">the string data to write internally</param>
    /// <returns></returns>
    private async Task WriteAsync(string data) =>
      await WriteAsync(Encoding.UTF8.GetBytes(data));

    /// <summary>
    /// Write data to the internal network stream
    /// </summary>
    /// <param name="buffer">the data to write</param>
    /// <returns></returns>
    private async Task WriteAsync(byte[] buffer)
    {
      await GetStream().WriteAsync(buffer, 0, buffer.Length);
    }

    /// <summary>
    /// Read a single byte from the internal network stream
    /// </summary>
    /// <returns>a single byte or errors out if EOF</returns>
    private byte ReadByte()
    {
      byte read;
      read = (byte)GetStream().ReadByte();
      if (read < 0) throw new Exception("EOF Reached");
      return read;
    }

    /// <summary>
    /// Read Fixed or All data from Stream
    /// </summary>
    /// <param name="amount">the amount to read (-1 if all)</param>
    /// <returns>data collected from read</returns>
    private async Task<byte[]> ReadAsync(int amount = -1)
    {
      int received;  // the amount of data read
      byte[] buffer; // the buffer data is read into

      // read fixed amount of data
      if (amount > 0) {
        buffer = new byte[amount];
        received = await GetStream().ReadAsync(buffer, 0, buffer.Length);
        if (received < 0) throw new Exception("EOF Reached");
        if (received < buffer.Length) {
          byte[] actual = new byte[received];
          Array.Copy(buffer, actual, received);
          buffer = actual;
        }
        return buffer;
      }

      // Read until End of Stream
      using (MemoryStream reader = new MemoryStream())
      {
        buffer = new byte[ReadBufferSize];
        do
        {
          received = await GetStream().ReadAsync(buffer, 0, buffer.Length);
          if (received < 0) throw new Exception("EOF Reached");
          if (received > 0) reader.Write(buffer, 0, received);
          if (received < buffer.Length) break;
        } while (received > 0);
        return reader.ToArray();
      }
    }

    /// <summary>
    /// Send a websocket frame using parameters provided
    /// </summary>
    /// <param name="data">the raw byte data to send in payload</param>
    /// <param name="opcode">the opcode to use in frame</param>
    /// <param name="fin">if last packet or continuation</param>
    /// <param name="masked">if packet is masked</param>
    /// <returns></returns>
    public async Task SendAsync(byte[] data, OP opcode=OP.Text, bool fin=true, bool masked=true)
    {
      // Create payload to send
      int size = 2 + (masked ? 4 : 0) + data.Length;
      if (data.Length <= 125) { }
      else if (data.Length >= 126 && data.Length <= 65536) size += 2;
      else size += 8;
      int i, offset = 0;
      int length = data.Length;
      byte[] buffer = new byte[size];

      // set first header: FIN & Opcode
      buffer[offset++] = (byte)((fin ? 0x80 : 0) | (byte)opcode);

      // set basic payload size
      if (length <= 125)
        buffer[offset++] = (byte)((masked ? 0x80 : 0) | length);

      // set 32 bit payload size
      else if (length >= 126 && length <= 65536) {
        buffer[offset++] = (byte)((masked ? 0x80 : 0) | 0x7E);
        for (i = 8; i > -1; i -= 8)
          buffer[offset++] = (byte)((length >> i) & 0xFF);

      // set 64 bit payload size
      } else {
        buffer[offset++] = (byte)((masked ? 0x80 : 0) | 0x7F);
        for (i = 56; i > -1; i -= 8)
          buffer[offset++] = (byte)((length >> i) & 0xFF);
      }

      // create mask if needed
      byte[] mask = null;
      if (masked) {
        mask = new byte[4];
        Random.NextBytes(mask);
        for (i = 0; i < 4; i++)
          buffer[offset++] = mask[i];
      }

      // add the payload data & mask it if necessary
      for (i = 0; i < length; i++)
        buffer[offset++] = (!masked && mask == null)
          ? data[i] : (byte)(data[i] & mask[i % 4]);

      // Send data over connection
      await WriteAsync(buffer);
    }

    /// <summary>
    /// Received a full websocket frame from connection
    /// </summary>
    /// <returns>The built websocket frame</returns>
    public async Task<Frame> ReceiveAsync()
    {
      Frame frame = await ReadFrameAsync();
      if (frame.Fin) return frame;
      using (MemoryStream stream = new MemoryStream())
      {
        await stream.WriteAsync(frame.Data, 0, frame.Data.Length);
        while (!frame.Fin || frame.OpCode == OP.Continue)
        {
          frame = await ReadFrameAsync();
          await stream.WriteAsync(frame.Data, 0, frame.Data.Length);
        }
        frame.Data = stream.ToArray();
        // TODO: Implement Close Frame Code and Reason extraction
        if (frame.OpCode == OP.Close) {
          await SendAsync(frame.Data, OP.Close, true, false);
          client.Dispose();
          throw new Exception("Websocket closed");
        }
        return frame;
      }
    }

    /// <summary>
    /// Read a frame internally and give it to ReceiveFrame processor
    /// </summary>
    /// <returns>The parsed websocket frame</returns>
    private async Task<Frame> ReadFrameAsync()
    {
      // Create frame for parsing
      Frame frame = new Frame();

      // Check reserved bytes
      byte buffer = ReadByte();
      if ((buffer & 0x40) != 0 || // rsv1
          (buffer & 0x20) != 0 || // rsv2
          (buffer & 0x10) != 0)   // rsv3
        throw new Exception("Invalid Frame: Reserved bytes are set");

      // Get FIN + OpCode
      frame.Fin = (buffer & 0x80) != 0;
      switch ((byte)(buffer & 0x0f))
      {
        case (int)OP.Continue: frame.OpCode = OP.Continue; break;
        case (int)OP.Text: frame.OpCode = OP.Text; break;
        case (int)OP.Binary: frame.OpCode = OP.Binary; break;
        case (int)OP.Close: frame.OpCode = OP.Close; break;
        case (int)OP.Ping: frame.OpCode = OP.Ping; break;
        case (int)OP.Pong: frame.OpCode = OP.Pong; break;
        default: frame.OpCode = OP.Text; break;
      }

      // Get Mask and payload length
      buffer = ReadByte();
      frame.Masked = (buffer & 0x80) != 0;
      int dataSize = (int)(0x7F & buffer);
      int dataCount = 0;

      // Get extended length range
      if (dataSize == 0x7F)
        dataCount = 8;
      else if (dataSize == 0x7E)
        dataCount = 2;
      while (--dataCount > 0)
        dataSize |= (ReadByte() & 0xFF) << (8 * dataCount);

      // Get masking key
      byte[] maskingKey = null;
      if (frame.Masked) {
        maskingKey = new byte[4];
        for (byte i = 0; i < maskingKey.Length; i++)
          maskingKey[i] = ReadByte();
      }

      // fetch data and demask if needed
      byte[] payload = await ReadAsync(dataSize);
      if (frame.Masked)
        for (uint i = 0; i < payload.Length; i++)
          payload[i] ^= maskingKey[i % 4];
      frame.Data = payload;

      // return the created frame
      return frame;
    }

    /// <summary>
    /// Blocking alternative to <see cref="ConnectAsync(string)"/>
    /// </summary>
    /// <param name="url">the url to connect to</param>
    /// <returns>self instance</returns>
    public WebSocket Connect(string url)
    {
      return ConnectAsync(url).Result;
    }

    /// <summary>
    /// Blocking alternative to <see cref="SendAsync(byte[], OP, bool, bool)"/>
    /// </summary>
    /// <param name="data">the data to send in frame</param>
    /// <param name="opcode">the frame opcode</param>
    /// <param name="fin">if frame end</param>
    /// <param name="masked">if frame masked</param>
    public void Send(byte[] data, OP opcode=OP.Text, bool fin=true, bool masked=true)
    {
      SendAsync(data, opcode).RunSynchronously();
    }

    /// <summary>
    /// Blocking Alternative to <see cref="ReceiveAsync"/>
    /// </summary>
    /// <returns>A parsed frame</returns>
    public Frame Receive()
    {
      return ReceiveAsync().Result;
    }

    /// <summary>
    /// Close the TCP Client object
    /// </summary>
    /// <param name="force">Whether or not to force close without sending op</param>
    /// <returns></returns>
    public async Task CloseAsync(bool force = false)
    {
      // TODO: Better close implementation
      if (!force) {
        byte[] dummyData = new byte[1] { (byte)'a' };
        await SendAsync(dummyData, OP.Close, true, false);
      }
      client.Dispose();
    }

    /// <summary>
    /// Blocking alternative to <see cref="CloseAsync(bool)"/>
    /// </summary>
    /// <param name="force">if force close connection</param>
    public void Close(bool force = false)
    {
      CloseAsync(force).RunSynchronously();
    }

    /* Custom dispose implementation */
    private void Dispose(bool disposing)
    {
      if (!disposed)
      {
        if (disposing) Close(true);
        disposed = true;
      }
    }

    /* Base Dispose Implementation */
    public void Dispose()
    {
      Dispose(true);
      GC.SuppressFinalize(this);
    }
  }
}