refactor socket

This commit is contained in:
wancat committed 2024-02-19 19:24:18 +08:00
1 parent ad66191d88
commit 2d05149cea
1 file changed
+34 -16
@@ -1,43 +1,56 @@
import { Server } from 'socket.io' import { Server } from 'socket.io'
import type { Socket } from 'socket.io'
import validator from 'validator' import validator from 'validator'
import type { Server as HttpServer } from 'http'
import { Session } from '@linux-odyssey/models' import { Session } from '@linux-odyssey/models'
import { getAndStartContainer, attachContainer } from '../containers/docker.js' import { getAndStartContainer, attachContainer } from '../containers/docker.js'
import SessionMiddleware from '../middleware/session.js' import SessionMiddleware from '../middleware/session.js'
import { genJWT } from '../utils/auth.js' import { genJWT } from '../utils/auth.js'
import logger from '../utils/logger.js' import logger from '../utils/logger.js'
const sessions = new Map() // eslint-disable-next-line no-unused-vars
type SocketCallback = (event: string, ...args: any[]) => void
function listenToSession(sessionId, callback) { const sessions = new Map<string, SocketCallback[]>()
function listenToSession(sessionId: string, callback: SocketCallback) {
const callbacks = sessions.get(sessionId) || [] const callbacks = sessions.get(sessionId) || []
sessions.set(sessionId, callbacks.concat(callback)) sessions.set(sessionId, callbacks.concat(callback))
} }
function removeFromSession(sessionId, event, callback) { function removeFromSession(sessionId: string, callback: SocketCallback) {
const callbacks = sessions.get((sessionId, event)) || [] const callbacks = sessions.get(sessionId) || []
sessions.set( sessions.set(
(sessionId, event), sessionId,
callbacks.filter((cb) => cb !== callback) callbacks.filter((cb) => cb !== callback)
) )
} }
export function pushToSession(sessionId, event, ...args) { export function pushToSession(
sessionId: string,
event: string,
...args: any[]
) {
const callbacks = sessions.get(sessionId) const callbacks = sessions.get(sessionId)
if (callbacks) { if (callbacks) {
callbacks.forEach((callback) => callback(event, ...args)) callbacks.forEach((callback) => callback(event, ...args))
} }
} }
async function connectContainer(socket, next) { // eslint-disable-next-line no-unused-vars
const user = socket.request.session?.passport?.user async function connectContainer(socket: Socket, next: (err?: Error) => void) {
// console.log(user) const user = (socket.request as any).session?.passport?.user
if (!user) { if (!user) {
next(new Error('User not found.')) next(new Error('User not found.'))
return return
} }
// console.log(socket.handshake) // console.log(socket.handshake)
const { sessionId } = socket.handshake.query const { sessionId } = socket.handshake.query
if (!(sessionId && validator.isMongoId(sessionId))) { if (
!sessionId ||
typeof sessionId !== 'string' ||
!validator.isMongoId(sessionId)
) {
logger.warn('Invalid Session ID.', user.username, sessionId) logger.warn('Invalid Session ID.', user.username, sessionId)
next(new Error('Invalid Session ID.')) next(new Error('Invalid Session ID.'))
return return
@@ -53,10 +66,15 @@ async function connectContainer(socket, next) {
sessionId: session.id, sessionId: session.id,
}) })
if (!session.containerId) {
next(new Error('Session has no container.'))
return
}
try { try {
const container = await getAndStartContainer(session.containerId) const container = await getAndStartContainer(session.containerId)
const stream = await attachContainer(container, { token }) const stream = await attachContainer(container, { token })
socket.context = { socket.data.context = {
session, session,
stream, stream,
} }
@@ -66,9 +84,9 @@ async function connectContainer(socket, next) {
} }
} }
function onConnect(socket) { function onConnect(socket: Socket) {
const { session, stream } = socket.context const { session, stream } = socket.data.context
stream.socket.on('data', (chunk) => { stream.socket.on('data', (chunk: Buffer) => {
socket.emit('terminal', chunk.toString()) socket.emit('terminal', chunk.toString())
}) })
@@ -76,7 +94,7 @@ function onConnect(socket) {
stream.socket.write(message) stream.socket.write(message)
}) })
const socketCallback = (event, ...data) => { const socketCallback = (event: string, ...data: any[]) => {
socket.emit(event, ...data) socket.emit(event, ...data)
} }
@@ -89,7 +107,7 @@ function onConnect(socket) {
}) })
} }
export default (server) => { export default (server: HttpServer) => {
const io = new Server(server) const io = new Server(server)
io.engine.use(SessionMiddleware) io.engine.use(SessionMiddleware)