当前位置：首页 > article >正文

企业微信会话存档：大文件拉取、加密、上传

article 2024/11/28 2:36:04

承接之前的文章企业微信会话存档sdk报错：A fatal error has been detected by the Java Runtime Environment

在之前的那篇文章结尾，我说发现了系统另外一个隐含的bug：企业微信会话存档拉取媒体文件是分片拉取，我的处理方式有问题导致，保存的永远是最后一个分片。

这个问题之所以没有在测试环境复现，是因为测试环境都是拿的一些小文件测试，结果在线上运行一段时间后，就发现，下载下来的文件显然比真实文件小很多。仔细排查了一下才发现是这个问题。

怎么解决呢？既然企业微信的sdk是分片拉取媒体文件的，那么我们就必需对内容进行汇总。但是汇总之前我们得判断文件的大小，如果文件不大，那么直接用ByteArrayOutputStream接收即可；但是如果文件很大，直接用ByteArrayOutputStream可能导致OOM，这里需要使用临时文件接收。

我使用的是weixin-java-cp这个库，所以代码用这个库的方法演示。
对于小文件，使用ByteArrayOutputStream接收：

ByteArrayOutputStreamByteConsumer byteArrayOutputStreamByteConsumer = new ByteArrayOutputStreamByteConsumer(bos);
msgAuditService.getMediaFile(sdk, sdkFileId, null, null, 180L, byteArrayOutputStreamByteConsumer);

这里我们抽象了一个ByteArrayOutputStreamByteConsumer类，每拉一次分片都会调用这个类，往bos中写分片数据。

public class ByteArrayOutputStreamByteConsumer implements Consumer<byte[]> {
    private final ByteArrayOutputStream bos;

    public ByteArrayOutputStreamByteConsumer(ByteArrayOutputStream bos) {
        this.bos = bos;
    }

    @Override
    public void accept(byte[] bytes) {
        try {
            bos.write(bytes);
        } catch (IOException e) {
            throw new MsgFileException("写入ByteArrayOutputStream失败", e);
        }
    }
}

如果文件比较大，我们需要写入临时文件，使用hutool的工具类创建临时文件。

File tempFile = FileUtil.createTempFile("audit_media_", "." + StringUtils.defaultString(ext, "bin"), false);
msgAuditService.getMediaFile(sdk, sdkFileId, null, null, 180L, tempFile.getAbsolutePath());

至此，我们已经解决文件下载的问题，我们总结一下：

小文件：用ByteArrayOutputStream接收
大文件：用临时文件接收

下载问题只是第一个问题，还有两个问题等着我们。

第二个问题就是文件加密的问题。由于我们的项目媒体文件存在云存储，所以要求上传前必须加密，防止文件泄露。如果你们没有这个需要，那就忽略这个问题。

我之前的加密方式使用的是AES/ECB/PKCS5Padding这种模式，因为使用的是hutool的工具类，默认AES就是这个模式。
之前的做法是拿到文件的byte[]，然后对它加密。如果文件很大，比如几个G，那很容易就OOM，为啥之前没有OOM呢，因为第一个问题导致文件内容只有最后一个分片的大小，所以这个问题被隐藏了。

所以现在就不能这么干了，可以使用AES/CTR/NoPadding这种模式代替，这种模式非常适合大文件的加解密。不用一次性传入整个文件的内容，可以流式的进行处理。算法的原理可以自行搜索，这里就不赘述了。

jdk提供的CipherOutputStream可以简化操作：

try (is; CipherOutputStream cos = new CipherOutputStream(os, encryptCipher)) {
     while (true) {
         byte[] buffer = new byte[8192];

         int read = is.read(buffer, 0, 8192);
         if (read == -1) {
             break;
         }
         cos.write(buffer, 0, read); // 这里write的时候就会同时进行加密
     }
 } catch (IOException e) {
     throw new FileUploadException("文件加密过程失败", e);
 }

这里的encryptCipher是Cipher的实例：

private Cipher getCipher(String key, byte[] ivBytes) {
    SecureRandom secureRandom = new SecureRandom();
    secureRandom.nextBytes(ivBytes);
    IvParameterSpec ivParameterSpec = new IvParameterSpec(ivBytes);

    SecretKeySpec secretKeySpec = new SecretKeySpec(key.getBytes(StandardCharsets.UTF_8), "AES");
    Cipher encryptCipher;
    try {
        encryptCipher = Cipher.getInstance("AES/CTR/NoPadding");
        encryptCipher.init(Cipher.ENCRYPT_MODE, secretKeySpec, ivParameterSpec);
    } catch (Exception e) {
        throw new FileUploadException("加密套件初始化失败", e);
    }
    return encryptCipher;
}