miaoShare/server.js

import express from "express";
import cors from "cors";
import Database from "better-sqlite3";
import fs from "fs";
import { WebSocketServer } from "ws";
import http from "http";
import { v4 as uuid4 } from "uuid";
import jwt from "jsonwebtoken";
import sha256 from "crypto-js/sha256.js";
import Config from './miaoShareConfig.json' assert { type: 'json' };
// function loadJSON(path){
// 	return JSON.parse(fs.readFileSync(path).toString());
// }
// const Config = loadJSON("./miaoShareConfig.json");

const HOST = "0.0.0.0";
const PORT = Config.PORT
// const PORT = 17057;
const fileDir = "./temp";
const dbPath = "./miaoShare.db";
const PATH = "./dist";

const JWT_SK = fs.readFileSync("JWT_SK.key").toString();

// 上传文件最大大小：200MB
const uploadMaxSize = 200 * 1024 * 1024;
// 一天的毫秒数
const millisecondsOfOneDay = parseInt(24 * 3600 * 1000);
// 文件3天过期
const LifeOfFile = millisecondsOfOneDay * 3;
// const LifeOfFile=0

// 数据库相关
const db = new Database(dbPath, { verbose: null });
db.exec(
	`CREATE TABLE IF NOT EXISTS id2fileName (
    id CHAR(8) PRIMARY KEY NOT NULL UNIQUE,
    fileName TEXT NOT NULL,
    timestamp INT NOT NULL,
    freeze tinyint(1) NOT NULL,
    downloadCredit INT NOT NULL,
	fileSize INT NOT NULL,
	sender_id CHAR(32) NOT NULL
);`
);
db.exec(`
CREATE TABLE IF NOT EXISTS USERS(
    user_id CHAR(32) PRIMARY KEY UNIQUE,
    password_hash CHAR(32) NOT NULL,
    user_group VARCHAR(32) NOT NULL,
	uploadFilesSizeCount BIGINT NOT NULL,
	email CHAR(64)
);`);
const stmt_getPermissionGroup = db.prepare("SELECT user_group FROM USERS WHERE user_id = ? AND password_hash = ?");
const stmt_getFileName = db.prepare(`SELECT fileName FROM id2fileName WHERE id = ?`);
const stmt_getUserUploadFilesSizeCount = db.prepare("SELECT uploadFilesSizeCount FROM USERS WHERE user_id = ?");
const stmt_getUserGroup = db.prepare("SELECT user_group FROM USERS WHERE user_id = ?");
const stmt_getTimestamp = db.prepare(`SELECT timestamp FROM id2fileName WHERE id = ?`);
const stmt_getFileSize = db.prepare(`SELECT fileSize FROM id2fileName WHERE id = ?`);
const stmt_getFreeze = db.prepare(`SELECT freeze FROM id2fileName WHERE id = ?`);
const stmt_getDownloadCredit = db.prepare(`SELECT downloadCredit FROM id2fileName WHERE id = ?`);
const stmt_getFreezenIDs = db.prepare(`SELECT id,fileSize,sender_id FROM id2fileName WHERE freeze = 1`);
const stmt_insertId2FileName = db.prepare(`INSERT INTO id2fileName (id, fileName, timestamp, freeze, downloadCredit, fileSize, sender_id) VALUES (?, ?, ?, ?, ?, ?, ?)`);
const stmt_deleteExpireData = db.prepare(`DELETE FROM id2fileName WHERE freeze=1`);
const stmt_freezeExpireData = db.prepare(`UPDATE id2fileName SET freeze=1 WHERE timestamp < ? OR downloadCredit < 1`);
const stmt_updateDownloadCredit = db.prepare(`UPDATE id2fileName SET downloadCredit=downloadCredit-1 WHERE id = ?`);
const stmt_updateUserUploadFilesSizeCount = db.prepare(`UPDATE USERS SET uploadFilesSizeCount=uploadFilesSizeCount+? WHERE user_id = ?`);
const stmt_updateUserUploadFilesSizeCount_d = db.prepare(`UPDATE USERS SET uploadFilesSizeCount=uploadFilesSizeCount-? WHERE user_id = ?`);

const db_insert = (id, fileName, downloadCredit, fileSize, senderId) => {
	stmt_insertId2FileName.run(id, fileName, parseInt(Date.now()), 0, downloadCredit, fileSize, senderId);
};
const db_getFileName = id => {
	try {
		if (stmt_getFreeze.get(id)["freeze"] == 1) return null;
		return stmt_getFileName.get(id)["fileName"];
	} catch (e) {
		return null;
	}
};
// 检查是否可下载
const db_downloadable = id => {
	try {
		// 懒
		if (stmt_getDownloadCredit.get(id)["downloadCredit"] < 1) return false;
		if (stmt_getFreeze.get(id)["freeze"] == 1) return false;
		return true;
	} catch (e) {
		return false;
	}
};
const db_updateDownloadCredit = id => {
	try {
		stmt_updateDownloadCredit.run(id);
		return true;
	} catch (e) {
		return false;
	}
};

const db_getTimestamp = id => {
	return stmt_getTimestamp.get(id)["timestamp"];
};

const db_getFileSize = id => {
	return stmt_getFileSize.get(id)["fileSize"];
};
const db_getUserUploadFilesSizeCount = id => {
	return stmt_getUserUploadFilesSizeCount.get(id)["uploadFilesSizeCount"];
};

function db_updateUserUploadFilesSizeCount(id,size){
	return stmt_updateUserUploadFilesSizeCount.run(size,id)
}

// 减
function db_updateUserUploadFilesSizeCount_d(id,size){
	return stmt_updateUserUploadFilesSizeCount_d.run(size,id)
}

// 更新存储的文件，删除已经被冻结的文件以及相关记录，然后冻结过期文件
const updateDBAndFiles = () => {
	try {
		let rows = stmt_getFreezenIDs.all();
		stmt_deleteExpireData.run();
		for (let row of rows) {
			let fileName = row["id"];
			let sid = row['sender_id']
			let fsize = row['fileSize']
			db_updateUserUploadFilesSizeCount_d(sid,fsize)
			let path = fileDir + "/" + fileName;
			try {
				fs.unlink(path, e => {
					if (e) throw e;
				});
			} catch (e) {
				console.log(e);
				continue;
			}
		}
		stmt_freezeExpireData.run(parseInt(Date.now()) - LifeOfFile);
	} catch (e) {
		console.log(e);
		return;
	}
};
updateDBAndFiles();
// 3小时更新一次
setInterval(updateDBAndFiles, 1000 * 3600 * 3);
// // 测试环境一分钟更新一次，记得注释掉
// setInterval(updateDBAndFiles, 1000 * 60);

const app = express();

app.use(cors());
app.options("*", cors());
// var tempMap = {}

if (!fs.existsSync(fileDir)) {
	fs.mkdirSync(fileDir);
}

function db_getPermissionGroup(id, pwdH) {
	let res = stmt_getPermissionGroup.get(id, pwdH);
	// console.log(id,pwdH,res)
	if (res && res.user_group) {
		return res.user_group;
	}
	return null;
}

function generateId(isGID) {
	var returnStr = "",
		range = 8,
		charStr = "ABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789";

	for (var i = 0; i < range; i++) {
		if (isGID === true && i == 0) {
			var index = parseInt(Math.random() * 26);
		} else {
			var index = parseInt(Math.random() * 36);
		}
		returnStr += charStr.substring(index, index + 1);
	}
	return returnStr;
}

// 上传文件接口
function uploadFile(req, res, fileName, downloadCredit, fileSize) {
	let id = generateId();
	const fileStream = fs.createWriteStream(fileDir + "/" + id);
	db_insert(id, fileName, downloadCredit, fileSize, req.sender_id);
	db_updateUserUploadFilesSizeCount(req.sender_id,fileSize)
	// tempMap[id] = fileName;
	// console.log(tempMap);
	req.pipe(fileStream);
	res.on("data", c => {
		console.log(c);
	});
	req.on("end", function () {
		fileStream.end();
		res.statusCode = 200;
		res.end(id);
	});
	req.on("error", err => {
		fileStream.end();
		res.statusCode = 500;
		res.end("File save failed: " + err.message);
	});
}

function downloadFile(req, res, fileId) {
	try {
		if (db_downloadable(fileId) == false) {
			res.writeHead(404);
			res.end(null);
			return;
		}
		db_updateDownloadCredit(fileId);
		var fileName = db_getFileName(fileId);
		if (fileName == null) {
			res.writeHead(404);
			res.end(null);
			return;
		}
		let data = fs.readFileSync(fileDir + "/" + fileId);
		res.writeHead(200, {
			"Content-Type": "application/octet-stream ",
			"File-Name": fileName,
			"Content-Length": data.length
		});
		res.end(data);
	} catch (err) {
		console.log(err);
		res.writeHead(404);
		res.end(null);
	}
}

const getUserGroup = (function(){
	const cache = {}
	return function(userId){
		if(!cache[userId]){
			cache[userId]= stmt_getUserGroup.get(userId)['user_group']
		}
		return cache[userId]
	}
})()

// 验证token
const verifyToken = (req, res, next) => {
	try {
		const token = req.headers.authorization.split(" ")[1]; // 从Authorization头中提取JWT
		console.log(req.headers.authorization,token)
		let verifyData = jwt.verify(token, JWT_SK);
		if (verifyData.exp > parseInt(Date.now())) {
			throw "token has expired";
		}
		req.sender_id = verifyData.id
		next();
	} catch (err) {
		console.log(err)
		return res.status(401).json({ message: "Invalid token", err });
	}
};

app.use(express.static(PATH));
app.post("/user/login", express.json(),function (req, res) {
	res.header("Access-Control-Allow-Origin", "*");
	res.header("Access-Control-Allow-Headers", "*");
	res.header("Access-Control-Allow-Methods", "POST");
	// console.log(req.body)
	let { id, pwd } = req.body;
	console.log(req.body)
	let pwdH = sha256(id + pwd).toString();
	let userGroup = db_getPermissionGroup(id, pwdH);
	if (userGroup !== null) {
		res.writeHead(200);
		let token = jwt.sign({ id, pwd }, JWT_SK, { expiresIn: "168h" });
		res.end(JSON.stringify({ status: "Success", userGroup, token }));
	} else {
		res.end(JSON.stringify({ status: "Error" }));
	}
	// console.log(req)
});
// app.all()
app.post("/file/upload", verifyToken, function (req, res) {
	res.header("Access-Control-Allow-Origin", "*");
	res.header("Access-Control-Allow-Headers", "*");
	res.header("Access-Control-Allow-Methods", "POST");
	res.header("Access-Control-Expose-Headers", "File-Name");
	var fileSize = parseInt(req.get("Content-Length"));
	var fileName = req.query.fileName;
	var downloadCredit = parseInt(req.query.downloadCredit);
	if (downloadCredit == -1) downloadCredit = 2147483647;
	if (!fileName) {
		res.statusCode = 500;
		res.end("FileName Not Found");
	} else if (fileSize > uploadMaxSize || db_getUserUploadFilesSizeCount(req.sender_id) + fileSize > Config.userGroupSetting[getUserGroup(req.sender_id)].cloudCapacity) {
		res.statusCode = 500;
		res.end("Maximum size limit exceeded");
	} else {
		uploadFile(req, res, fileName, downloadCredit, fileSize);
	}
});

app.get("/file/download", verifyToken, function (req, res) {
	res.header("Access-Control-Allow-Origin", "*");
	res.header("Access-Control-Allow-Headers", "*");
	res.header("Access-Control-Allow-Methods", "GET");
	res.header("Access-Control-Expose-Headers", "File-Name");
	var fileId = req.query.fileName;
	if (!fileId) {
		res.statusCode = 500;
		res.end("FileName Not Found");
	} else {
		downloadFile(req, res, fileId);
	}
});

// // 启动服务器
// app.listen(PORT, HOST, function(){
//     console.log(app)
//     // console.log(`miaoShare server listening on ${}`)
// });

const server = http.createServer(app);
const wss = new WebSocketServer({ server });
// 先实现最低限度功能，暂时不考虑身份认证
const groupsMap = {};
// 一天不活跃的群组将被清除
const LifeOfGroup = millisecondsOfOneDay * 1;
// const LifeOfGroup = 0;
// 根据uuid查询加入的群组
const uuid2Groups = {};

db.exec(
	`CREATE TABLE IF NOT EXISTS CHATGROUPS (
	gID CHAR(8) PRIMARY KEY UNIQUE NOT NULL,
	verify TEXT NOT NULL,
	activeTimestamp INT NOT NULL
);`
);

const stmt_getExpiredGroups = db.prepare(`SELECT gID FROM CHATGROUPS WHERE activeTimestamp < ? `);

const stmt_getGroupVerify = db.prepare(`SELECT verify FROM CHATGROUPS WHERE gID = ? `);

const stmt_insertCHATGROUPS = db.prepare(`INSERT INTO CHATGROUPS (gID,verify, activeTimestamp) VALUES (?, ?, ?)`);

const stmt_updateCHATGROUPS_ActiveTimestamp = db.prepare(`UPDATE CHATGROUPS SET activeTimestamp=? WHERE gID=?`);

const stmt_deleteGroup = db.prepare(`DELETE FROM CHATGROUPS WHERE gID=?`);

const generateStmt_insertMessage_map = {};

const updateCHATGROUPS = () => {
	// console.log(1)
	try {
		let gIDs = stmt_getExpiredGroups.all(parseInt(Date.now()) - LifeOfGroup);
		for (let row of gIDs) {
			let gID = row.gID;
			db.exec(`DROP TABLE IF EXISTS ${gID}`);
			stmt_deleteGroup.run(gID);
		}
	} catch (e) {
		console.log(e);
		return;
	}
};
updateCHATGROUPS();
setInterval(updateCHATGROUPS, 1000 * 3600 * 3);

// 创建群组的数据库
function createGroupTableDB(gID, verify) {
	db.exec(
		`CREATE TABLE IF NOT EXISTS ${gID} (
		message TEXT NOT NULL,
		type TEXT NOT NULL,
		sender TEXT NOT NULL,
		timestamp INT NOT NULL
	);`
	);
	stmt_insertCHATGROUPS.run(gID, verify, parseInt(Date.now()));
}

// 插入信息
function insertMessageByGID(gID, msg, from, timestamp) {
	if (!generateStmt_insertMessage_map[gID]) {
		generateStmt_insertMessage_map[gID] = db.prepare(`INSERT INTO ${gID} (message,type,sender, timestamp) VALUES (?, ?, ?, ?)`);
	}
	generateStmt_insertMessage_map[gID].run(msg, "gMsg", from, timestamp);
}

function updateActiveTimestamp(gID) {
	stmt_updateCHATGROUPS_ActiveTimestamp.run(parseInt(Date.now()), gID);
}

function createGroup(verify) {
	// 暂时不考虑id随机生成时重复的问题
	let gID = generateId(true);
	groupsMap[gID] = { verify, member: [], msgHistory: [] };
	createGroupTableDB(gID, verify);
	return gID;
}

function getGroupVerify(gID) {
	initMsgHistory(gID);
	return groupsMap[gID].verify;
}

// 正式加入群组
function joinGroup(socket, gID, callback) {
	initMsgHistory(gID);
	groupsMap[gID].member.push(socket);
	updateActiveTimestamp(gID);
	callback();
}

// 暂时这么写
function getHistory(gID) {
	initMsgHistory(gID);
	return groupsMap[gID].msgHistory;
}

function broadcastMsg(gID, msg, from) {
	const timestamp = parseInt(Date.now());
	initMsgHistory(gID);
	updateActiveTimestamp(gID);
	insertMessageByGID(gID, msg, from, timestamp);

	let bmsg = JSON.stringify({ type: "gMsg", gID, msg, from, timestamp });
	groupsMap[gID].msgHistory.push({ type: "gMsg", msg, from, timestamp });
	for (let socket of groupsMap[gID].member) {
		socket.send(bmsg);
	}
}

function initMsgHistory(gID) {
	if (!groupsMap[gID]) {
		let verify = stmt_getGroupVerify.get(gID).verify;
		// 暂时这么写
		let msgHistory = db
			.prepare(`SELECT message,type,sender,timestamp FROM ${gID}`)
			.all()
			.map(v => {
				return { type: v.type, msg: v.message, from: v.sender, timestamp: v.timestamp };
			});
		// console.log(msgHistory)
		groupsMap[gID] = { verify, member: [], msgHistory };
	}
}

// 灵车
wss.on("connection", socket => {
	console.log("新的连接已建立");
	const no = wss.clients.size;
	const joinedGroup = [];
	var uuid = uuid4();
	function setUuid(new_uuid) {
		uuid = new_uuid;
	}
	socket.on("message", message => {
		try {
			console.log(message.toString());
			let msg = JSON.parse(message.toString());
			if (msg.type) {
				let type = msg.type;
				// 创建群组
				if (type === "create") {
					socket.send(JSON.stringify({ type: "sMsg", type1: "create", gID: createGroup(msg.verify) }));
				} else if (type === "join") {
					// 加入群组——客端验证key
					if (msg.type1 === "verify") {
						// 已经加入则return
						if (joinedGroup.includes(msg.gID)) {
							return;
						}
						socket.send(JSON.stringify({ type: "sMsg", type1: "joinVerify", verify: getGroupVerify(msg.gID), gID: msg.gID }));
					} else if (msg.type1 === "success") {
						if (joinedGroup.includes(msg.gID)) {
							return;
						}
						// 成功加入
						joinGroup(socket, msg.gID, () => {
							joinedGroup.push(msg.gID);
							// 不考虑失败属于是
							socket.send(JSON.stringify({ type: "sMsg", gID: msg.gID, type1: "history", history: JSON.stringify(getHistory(msg.gID)) }));
						});
					}
				} else if (type === "message") {
					let gID = msg.gID;
					broadcastMsg(gID, msg.msg, uuid);
				} else if (type === "heartbeat") {
					socket.send(JSON.stringify({ type: "sMsg", type1: "heartbeat" }));
				} else if (type === "uuid") {
					if (msg.type1 === "get") {
						socket.send(JSON.stringify({ type: "sMsg", type1: "uuid", uuid }));
					} else if (msg.type1 === "set") {
						uuid = msg.uuid;
					}
				}
			}
		} catch (e) {
			console.log(e);
		}
	});

	socket.on("close", () => {
		for (let gID of joinedGroup) {
			groupsMap[gID].member = groupsMap[gID].member.filter(v => v !== socket);
		}
		console.log(`No.${no} exit.`);
		console.log(groupsMap);
	});
});

server.listen(PORT, HOST, () => {});