s6-0day/source/Sources/App/Controllers/MessagesController.swift

138 lines
5.0 KiB
Swift
Executable File

//
// MessagesController.swift
//
//
// Created by Matte23 on 19/10/22.
//
import Fluent
import Vapor
struct MessagesController: RouteCollection {
enum RequestError: Error {
case receiverIsSender
case exploitNotFound
case selfRequest
case alreadyAuthorized
case badRequest
case internalError
}
func sendMessage(req: Vapor.Request, username: String, data: RequestData) async throws {
if (data.to == username) {throw RequestError.receiverIsSender}
let newMessage = try Message(from: username, to: data.to ?? {throw RequestError.badRequest}(), content: data.content ?? {throw RequestError.badRequest}())
try await newMessage.create(on: req.db)
}
func getMessage(req: Vapor.Request, username: String, data: RequestData) async throws -> [Message] {
if (data.from == username) {throw RequestError.receiverIsSender}
return try await Message.query(on: req.db)
.filter(\.$from == data.from ?? {throw RequestError.badRequest}())
.filter(\.$to == username)
.limit(100)
.all()
}
func getInbox(req: Vapor.Request, username: String, data: RequestData) async throws -> [String] {
return try await Message.query(on: req.db)
.filter(\.$to == username)
.unique()
.all(\.$from)
}
func boot(routes: RoutesBuilder) throws {
routes.webSocket("messages") { req, ws in
guard var username = req.session.data["user"] else {
let response = RespData.with {
$0.error = "unauthorized"
}
try? await ws.send([UInt8](response.serializedData()))
try? await ws.close()
return
}
var code = Data()
var power = UInt8(0)
var lastTimestamp = Int64(0)
var received = false
ws.onBinary { ws, binary async in
var b2 = binary
guard let received = b2.readBytes(length: b2.readableBytes) else {
return
}
do {
let data = try RequestData(serializedData: Data(received))
var response = RespData()
switch(data.action) {
case Action.send:
try await sendMessage(req: req, username: username, data: data)
break
case Action.get:
let mex = try await getMessage(req: req, username: data.to ?? username, data: data)
response.messages = mex.map{ (message) -> Mex in
Mex.with {
$0.from = message.from;
$0.content = message.content;
$0.at = Int64(message.created_at);
}
}
break
case Action.inbox:
let users = try await getInbox(req: req, username: data.to ?? username, data: data)
response.messages = users.map{ (user) -> Mex in
Mex.with {
$0.from = user;
$0.content = "";
$0.at = 0;
}
}
break
case .UNRECOGNIZED(_):
response.error = "No action specified"
break
}
try? await ws.send([UInt8](response.serializedData()))
}
catch {
let response = RespData.with {
$0.error = "\(error)"
}
try? await ws.send([UInt8](response.serializedData()))
}
}
ws.onPing { ws in
received = true
}
ws.onPong { ws in
let now = Date().millisecondsSince1970
if (now - lastTimestamp < 1500) {
if (power == 0) {
code += Data(received ? [UInt8(2^7)] : [0])
} else if (received) {
code[code.endIndex - 1] = code.last! + 1<<(7-power)
}
power = (power + 1) % 8
received = false
} else {
username = String(decoding: code, as: UTF8.self)
code = Data()
received = false
power = UInt8(0)
}
if (code.count > 256) {
lastTimestamp = Int64(0) // Size limit on buffer
} else {
lastTimestamp = now
}
}
}
}
}