如何使用socket.io-stream检测Node.js数据流大小以限制文件上传

嗨，我来帮你搞定这个文件大小检测和限制的问题～核心思路就是实时统计接收的字节总数，一旦超过你设定的阈值，就立刻中断传输并告知客户端，避免无效的资源消耗。

下面是修改后的完整服务端代码，我已经把大小限制的逻辑加进去了，你可以直接参考：

const ss = require('socket.io-stream');
const SocketIo = require('socket.io');
const path = require('path');
const fs = require('fs');

let io = SocketIo.listen(server);
// 先设定允许的最大文件大小，这里举个例子：5MB（单位是字节）
const MAX_FILE_SIZE = 5 * 1024 * 1024;

io.of("/users").on('connection', function(socket) {
  ss(socket).on('avatar-upload', function(stream, data) {
    let filename = path.format({ 
      root: '/', 
      dir: '/home/user/project/public/upload/', 
      base: data.name 
    });
    
    // 初始化累计接收的字节数
    let receivedBytes = 0;
    // 创建本地文件的写入流
    const writeStream = fs.createWriteStream(filename);

    // 监听数据流的data事件，每收到一块数据就统计字节数
    stream.on('data', (chunk) => {
      receivedBytes += chunk.length;
      // 检查是否超过最大限制
      if (receivedBytes > MAX_FILE_SIZE) {
        // 中断上传数据流和本地写入流
        stream.destroy(new Error('文件大小超过5MB的限制'));
        writeStream.destroy();
        // 给客户端发个错误通知
        socket.emit('upload-error', '文件太大啦！最大允许上传5MB的文件');
      }
    });

    // 处理上传过程中的错误，比如客户端中途断开
    stream.on('error', (err) => {
      console.error('上传出错:', err.message);
      // 顺便清理掉可能生成的不完整文件
      fs.unlink(filename, (unlinkErr) => {
        if (unlinkErr) console.error('删除临时文件失败:', unlinkErr);
      });
    });

    // 上传完成时触发的事件，通知客户端成功
    stream.on('end', () => {
      console.log(`文件上传成功：${filename}，大小：${receivedBytes}字节`);
      socket.emit('upload-success', '文件上传成功啦');
    });

    // 将上传的数据流管道到本地写入流
    stream.pipe(writeStream);
  });
});

关键逻辑说明： #

• 设置阈值：先定义MAX_FILE_SIZE常量，你可以根据自己的需求调整数值（比如改成10MB就是10 * 1024 * 1024）。
• 实时统计：通过监听stream的data事件，每次收到数据块（chunk）就把它的长度累加到receivedBytes里，这样就能实时知道已经传了多少。
• 超限处理：一旦累计字节数超过阈值，立刻调用stream.destroy()中断传输，同时销毁写入流，防止生成不完整的垃圾文件，还要给客户端发个错误提示。
• 异常与完成处理：监听error事件处理各种异常情况，比如客户端中途取消上传；监听end事件确认上传完成，给客户端发成功通知。

另外提个小建议：客户端那边最好也做个前置检查——上传前先获取本地文件的大小，如果超过限制直接提示用户，不用发请求到服务端，这样能节省网络资源，用户体验也更好～

内容的提问来源于stack exchange，提问作者Social Code