SocialPub/PrivaPub/Api/Mastodon/Controllers/StreamingController.cs

using System.Net.WebSockets;
using System.Text;
using System.Text.Json;
using System.Text.Json.Nodes;

using Microsoft.AspNetCore.Mvc;

using MongoDB.Entities;

using PrivaPub.Api.Mastodon.Infrastructure;
using PrivaPub.Api.Mastodon.Mappers;
using PrivaPub.Domain.Privacy;
using PrivaPub.Domain.Timelines;
using PrivaPub.Models.Social;
using PrivaPub.StaticServices;

namespace PrivaPub.Api.Mastodon.Controllers
{
	// Mastodon's streaming API: one WebSocket on /api/v1/streaming that subscribes to streams by message, or one stream
	// per request as server-sent events on /api/v1/streaming/{stream}. Every stream needs a signed-in persona (the token
	// comes as access_token, in the Authorization header, or as the WebSocket's protocol). What a stream tells is mapped
	// for that persona as it is sent: a post it may not see, or whose author it blocked or muted, is never sent.
	public class StreamingController : MastodonController
	{
		static readonly TimeSpan KeepAlive = TimeSpan.FromSeconds(15);

		readonly MastodonMapper _mapper;
		readonly DbEntities _dbEntities;

		public StreamingController(MastodonMapper mapper, DbEntities dbEntities)
		{
			_mapper = mapper;
			_dbEntities = dbEntities;
		}

		[HttpGet("/api/v1/streaming/health"), Microsoft.AspNetCore.Authorization.AllowAnonymous]
		public IActionResult Health() => Content("OK", "text/plain");

		// the stream a name and its parameter make, or why not: a hashtag needs its tag, a list must be the persona's own
		async Task<(Streams.Key Key, string Error)> Resolve(string name, string tag, string list, CancellationToken token)
		{
			switch (name)
			{
				case Streams.User or Streams.UserNotification or Streams.Public or Streams.PublicLocal or Streams.PublicRemote:
					return (new Streams.Key(name), default);
				case Streams.Hashtag or Streams.HashtagLocal:
					var normal = TagsController.Normalise(tag);
					return normal.Length == 0 ? (default, "Missing tag name parameter") : (new Streams.Key(name, normal), default);
				case Streams.List:
					return !string.IsNullOrEmpty(list) && await DB.Default.Find<PersonaList>().Match(l => l.ID == list && l.AvatarId == MyId).ExecuteAnyAsync(token)
						? (new Streams.Key(name, list), default)
						: (default, "Unknown list");
				default:
					return (default, "Unknown stream type");
			}
		}

		// what the event says, for this persona: a status or a notification as JSON, a deleted id as itself; null when it
		// is not the persona's to see
		async Task<string> Payload(Streams.Event e, CancellationToken token)
		{
			if (e.Name == "delete")
				return e.Id;
			if (e.Name == "notification")
			{
				var notification = await _dbEntities.Notifications.Match(n => n.ID == e.Id && n.AvatarId == MyId).ExecuteFirstAsync(token);
				var mapped = notification == default ? default : (await _mapper.Notifications(new List<Notification> { notification }, MyId, token)).FirstOrDefault();
				return mapped == default ? default : JsonSerializer.Serialize(mapped, MastodonJson.Options);
			}
			var post = await _dbEntities.Posts.MatchID(e.Id).ExecuteFirstAsync(token);
			if (post == default || !await VisibilityPolicy.CanSee(post, MyId, token))
				return default;
			var status = (await _mapper.Statuses(new[] { post }, MyId, token)).FirstOrDefault();
			return status == default ? default : JsonSerializer.Serialize(status, MastodonJson.Options);
		}

		[HttpGet("/api/v1/streaming"), Scope("read")]
		public async Task<IActionResult> Socket(CancellationToken token)
		{
			if (!HttpContext.WebSockets.IsWebSocketRequest)
				return Error(StatusCodes.Status400BadRequest, "This is a WebSocket endpoint; server-sent events are at /api/v1/streaming/{stream}");
			var protocol = HttpContext.WebSockets.WebSocketRequestedProtocols.FirstOrDefault();
			using var socket = await HttpContext.WebSockets.AcceptWebSocketAsync(protocol);
			var listener = Streams.Open(MyId);
			var sending = new SemaphoreSlim(1, 1);
			async Task Send(JsonObject message)
			{
				var bytes = Encoding.UTF8.GetBytes(message.ToJsonString());
				await sending.WaitAsync(token);
				try
				{
					await socket.SendAsync(bytes, WebSocketMessageType.Text, true, token);
				}
				finally
				{
					sending.Release();
				}
			}
			async Task Subscribe(string type, string name, string tag, string list)
			{
				var (key, error) = await Resolve(name, tag, list, token);
				if (error != default)
				{
					await Send(new JsonObject { ["error"] = error, ["status"] = 400 });
					return;
				}
				if (type == "unsubscribe")
					listener.Stop(key);
				else
					listener.Listen(key);
			}
			try
			{
				if (Params.Get("stream") is { Length: > 0 } first)
					await Subscribe("subscribe", first, Params.Get("tag"), Params.Get("list"));
				using var done = CancellationTokenSource.CreateLinkedTokenSource(token);
				var reading = Task.Run(async () =>
				{
					var buffer = new byte[4096];
					while (socket.State == WebSocketState.Open && !done.IsCancellationRequested)
					{
						var text = new StringBuilder();
						WebSocketReceiveResult received;
						do
						{
							received = await socket.ReceiveAsync(buffer, done.Token);
							if (received.MessageType == WebSocketMessageType.Close)
							{
								done.Cancel();
								return;
							}
							text.Append(Encoding.UTF8.GetString(buffer, 0, received.Count));
						}
						while (!received.EndOfMessage && text.Length < 16384);
						JsonNode message;
						try
						{
							message = JsonNode.Parse(text.ToString());
						}
						catch (JsonException)
						{
							continue;
						}
						if (message?["type"]?.GetValue<string>() is "subscribe" or "unsubscribe" && message["stream"]?.GetValue<string>() is { } stream)
							await Subscribe(message["type"]!.GetValue<string>(), stream, message["tag"]?.GetValue<string>(), message["list"]?.GetValue<string>());
					}
				}, done.Token);
				await foreach (var (key, e) in listener.Queue.Reader.ReadAllAsync(done.Token))
				{
					if (!listener.Hears(key) || await Payload(e, done.Token) is not { } payload)
						continue;
					await Send(new JsonObject { ["stream"] = new JsonArray(key.Wire.Select(w => (JsonNode)w).ToArray()), ["event"] = e.Name, ["payload"] = payload });
				}
				await reading;
			}
			catch (Exception ex) when (ex is OperationCanceledException or WebSocketException)
			{
			}
			finally
			{
				Streams.Close(listener);
				if (socket.State == WebSocketState.Open)
					await socket.CloseAsync(WebSocketCloseStatus.NormalClosure, default, CancellationToken.None);
			}
			return new EmptyResult();
		}

		[HttpGet("/api/v1/streaming/{*stream}"), Scope("read")]
		public async Task<IActionResult> Events(string stream, CancellationToken token)
		{
			var name = stream.Replace('/', ':');
			var (key, error) = await Resolve(name, Params.Get("tag"), Params.Get("list"), token);
			if (error != default)
				return Error(StatusCodes.Status400BadRequest, error);
			Response.ContentType = "text/event-stream";
			Response.Headers.CacheControl = "no-cache";
			Response.Headers["X-Accel-Buffering"] = "no";
			var listener = Streams.Open(MyId);
			listener.Listen(key);
			var writing = new SemaphoreSlim(1, 1);
			async Task Write(string text)
			{
				await writing.WaitAsync(token);
				try
				{
					await Response.WriteAsync(text, token);
					await Response.Body.FlushAsync(token);
				}
				finally
				{
					writing.Release();
				}
			}
			try
			{
				await Write(":)\n\n");
				var beating = Task.Run(async () =>
				{
					while (!token.IsCancellationRequested)
					{
						await Task.Delay(KeepAlive, token);
						await Write(":thump\n\n");
					}
				}, token);
				await foreach (var (_, e) in listener.Queue.Reader.ReadAllAsync(token))
				{
					if (await Payload(e, token) is not { } payload)
						continue;
					await Write($"event: {e.Name}\ndata: {payload}\n\n");
				}
				await beating;
			}
			catch (OperationCanceledException)
			{
			}
			finally
			{
				Streams.Close(listener);
			}
			return new EmptyResult();
		}
	}
}