libtorch 模型加密

Posted 2020-12-12 yanghailin

模型部署到现场为了防止泄密，需要加密。加密一方面可以防止泄密，另一方面可以便于模型跟踪管理，防止混淆。
libtorch的加载模型的函数，torch::jit::load();我点开load可以看到函数。有两个：

TORCH_API std::shared_ptr<script::Module> load(const std::string& filename,
    c10::optional<c10::Device> device = c10::nullopt);


TORCH_API std::shared_ptr<script::Module> load(std::istream& in,
    c10::optional<c10::Device> device = c10::nullopt);

一般我们都是用上面这个，直接给模型路径就可以了，下面这个是流，没用过。然后我试试下面这个流的怎么用。百度了一下istream如何赋值，发现

        std::filebuf in;
        if (!in.open(path_pt, std::ios::in)) {
            std::cout << "fail to open file" << std::endl;
            return 1;
        }

      std::istream ins(&in);

用std::filebuf读进来，再给istream就可以，然后送到load里面模型可以跑！所以问题就变得简单了，我只要把流加密保存，读进来的时候再解密送到load就可以。然而，事情总是不是这么一帆风顺的。
所以我就开始尝试流加密。

#include <iostream>
#include <fstream>
#include <string.h>
using namespace std;
int main()
{
    string path_pt = "/data_2/everyday/0622/00000.pt";
    std::filebuf in;
    std::filebuf outbuf;
    outbuf.open("/data_2/everyday/0622/1/0000-en-xor",std::ios::out);
    if (!in.open(path_pt, std::ios::in)) {
        std::cout << "fail to open file" << std::endl;
        return 0;
    }

    FILE *in_file,*out_file;
    in_file=fopen(path_pt.c_str(),"rb");//以读的方式打开二进制文件
    char ch=fgetc(in_file);

        string key = "1923456789765021";//"A34B123RTAa1qwe3";
        int x = key.size();
        int i = 0;

    //这一断注释代码的功能是对文件的所有位加密，主要用于文本文件。
    while(!feof(in_file))
    {
        ch = ch^key[i>=x?i=0:i++];
        outbuf.sputc(ch);
        ch=fgetc(in_file);
    }
    outbuf.sputc(ch);
    outbuf.close();

}

这里通过异或加密了流并保存在本地。然后把加密的模型读到工程里面再解密

    string path_pt = "/data_2/everyday/0622/1/0000-en-xor";
    std::filebuf in, outbuf;
    if (!in.open(path_pt, std::ios::in)) {
        std::cout << "fail to open file" << std::endl;
        return 1;
    }

    string key = "1923456789765021";//"A34B123RTAa1qwe3";
    int x = key.size();
    int i = 0;

    do {
        char ch = in.sgetc();
        ch = ch^key[i>=x?i=0:i++];

        outbuf.sputc(ch);
        //        std::cout << (int)ch<<std::endl;
    } while ( in.snextc() != EOF );
    outbuf.sputc(in.sgetc());

    std::istream ins(&outbuf);

可是就是这里出问题了，这样整然后给libtorch的load函数，爆错。

Starting /data_2/everyday/0618/build-libtorch-refinenet-unknown-Default/example-app...
terminate called after throwing an instance of ‘c10::Error‘
  what():  [enforce fail at inline_container.cc:130] . PytorchStreamReader failed checking magic number.
frame #0: c10::ThrowEnforceNotMet(char const*, int, char const*, std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> > const&, void const*) + 0x76 (0x7f3df37c4a76 in /data_1/Yang/project/2019/chejian/3rdparty/libtorch/lib/libc10.so)
frame #1: torch::jit::PyTorchStreamReader::valid(char const*) + 0x107 (0x7f3e04c7a5c7 in /data_1/Yang/project/2019/chejian/3rdparty/libtorch/lib/libcaffe2.so)
frame #2: torch::jit::PyTorchStreamReader::PyTorchStreamReader(std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> >, std::istream*) + 0x1ae (0x7f3e04c7b0ce in /data_1/Yang/project/2019/chejian/3rdparty/libtorch/lib/libcaffe2.so)
frame #3: torch::jit::load(std::istream&, c10::optional<c10::Device>) + 0x2d1 (0x7f3e08a3fda1 in /data_1/Yang/project/2019/chejian/3rdparty/libtorch/lib/libtorch.so.1)
frame #4: main + 0x291 (0x456947 in /data_2/everyday/0618/build-libtorch-refinenet-unknown-Default/example-app)
frame #5: __libc_start_main + 0xf0 (0x7f3df0015830 in /lib/x86_64-linux-gnu/libc.so.6)
frame #6: _start + 0x29 (0x44ddd9 in /data_2/everyday/0618/build-libtorch-refinenet-unknown-Default/example-app)

错误当然是看不懂，应该就是流不对呗。然后我就取出outbuf一两个看看：

do {
            char ch = in.sgetc();
            ch = ch^key[i>=x?i=0:i++];

            outbuf.sputc(ch);
            //        std::cout << (int)ch<<std::endl;
        } while ( in.snextc() != EOF );
        outbuf.sputc(in.sgetc());
        //    outbuf.sputc(EOF);
       
            char ch1 = outbuf.sgetc();
            std::cout<<(int)ch1<<std::endl;

            outbuf.snextc();
            ch1 = outbuf.sgetc();
            std::cout<<(int)ch1<<std::endl;

发现打印出来的全是-1.。。。。。咋回事呢？
是不是流你压到最后指针就在文件的最后，要把指针挪到最前面就可以了呢？然后各种查找资料，这个std::filebuf人家用的比较少，只有个c++官网介绍，http://www.cplusplus.com/reference/fstream/filebuf/。然后，找到了把指针移动到文件头的函数。

            outbuf.pubseekpos(0,std::ios::in);
            outbuf.pubsync();

可是还是不好使啊！难道函数的问题吗？

  string path_pt = "/data_2/everyday/0622/1/0000-en-xor";


        //    std::filebuf in,outbuf;

        std::filebuf in, outbuf;
        if (!in.open(path_pt, std::ios::in)) {
            std::cout << "fail to open file" << std::endl;
            return 1;
        }

            long filesize1 = static_cast<long>((in.pubseekoff (0,std::ios::end,std::ios::in)));
           in.pubseekpos(0,std::ios::in); //

           char ch12 = in.sgetc();
           std::cout<<(int)ch12<<std::endl;

            in.snextc();
            ch12 = in.sgetc();
           std::cout<<(int)ch12<<std::endl;

我验证了直接读取本地的std::filebuf in可以，filesize1可以记录大小，当不进行in.pubseekpos(0,std::ios::in);再读就是-1.说明函数有效的。
然后我又实验，

string path_pt = "/data_2/everyday/0622/1/0000-en-xor";


        //    std::filebuf in,outbuf;

        std::filebuf in, outbuf;
        if (!in.open(path_pt, std::ios::in)) {
            std::cout << "fail to open file" << std::endl;
            return 1;
        }
        
        outbuf.open("/data_2/everyday/0622/1/0000-jiemi", std::ios::out);

        string key = "1923456789765021";//"A34B123RTAa1qwe3";
        int x = key.size();
        int i = 0;

        do {
            char ch = in.sgetc();
            ch = ch^key[i>=x?i=0:i++];

            outbuf.sputc(ch);
            //        std::cout << (int)ch<<std::endl;
        } while ( in.snextc() != EOF );
        outbuf.sputc(in.sgetc());
         outbuf.close();

我把解密之后的模型保存本地，然后再读取进来经过：
std::istream ins(&in);
std::shared_ptrtorch::jit::script::Module m_model = torch::jit::load(ins);
这两步是可以的！！！！！可是问题出现在哪里呢？主要问题在于直接解密的时候：
outbuf.snextc();
ch1 = outbuf.sgetc();
std::cout<<(int)ch1<<std::endl;
这里输出死活都是-1,。不知道为啥。。。
然后还有一个检查方法，就是std::istream ins(&in); 检测ins里面的东西。

 std::istream ins(&in);


 vector<char> v_ch;
 char ch_tmp[1];
 int cnt_ = 0;
 while (ins.read(ch_tmp, 1))
  {
      cnt_ ++;
    v_ch.push_back(ch_tmp[0]);
     std::cout<<(int)ch_tmp[0] <<std::endl;
 }
 while(1);

正常可以推理的时候v_ch里面有2千万的数据，而直接在线解密的时候是0！。。折腾了好久无解，总感觉哪里还差一步就可以解决这个问题，也重定向了，无解。。。。

然后找来jiaming，他百度用了其他方法实验，反正只要把流塞到std::istream ins(&in);这里就可以了。

                char* buf = nullptr;
                string path_pt = "/data_2/everyday/0623/nice/jiami2";
                string key = "QO1##gX@@3";

                FILE* file = fopen(path_pt.c_str(), "r");
                fseek(file, 0, SEEK_END);
                unsigned size = ftell(file);
                fseek(file, 0, SEEK_SET);
                buf = (char*)malloc(size);
                memset(buf,0,size);
                unsigned int i=0,j=0;
                while(!feof(file))
                {
                    char ch = (fgetc(file))^(key[i>=key.length()?i=0:i++]);
                    buf[j++]=ch;
                }

                strstreambuf sbuf(buf, sizeof(buf));
//                free(buf);
//                buf = nullptr;
                std::istream ins(&sbuf);

这样子报错：

terminate called after throwing an instance of ‘c10::Error‘
  what():  [enforce fail at inline_container.cc:127] . PytorchStreamReader failed reading zip archive: not a ZIP archive
frame #0: c10::ThrowEnforceNotMet(char const*, int, char const*, std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> > const&, void const*) + 0x76 (0x7f362057fa76 in /data_1/Yang/project/2019/chejian/3rdparty/libtorch/lib/libc10.so)
frame #1: torch::jit::PyTorchStreamReader::valid(char const*) + 0xa6 (0x7f3631a35566 in /data_1/Yang/project/2019/chejian/3rdparty/libtorch/lib/libcaffe2.so)
frame #2: torch::jit::PyTorchStreamReader::PyTorchStreamReader(std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> >, std::istream*) + 0x1fb (0x7f3631a3611b in /data_1/Yang/project/2019/chejian/3rdparty/libtorch/lib/libcaffe2.so)
frame #3: torch::jit::load(std::istream&, c10::optional<c10::Device>) + 0x2d1 (0x7f36357fada1 in /data_1/Yang/project/2019/chejian/3rdparty/libtorch/lib/libtorch.so.1)
frame #4: main + 0x296 (0x45665c in /data_2/everyday/0618/build-libtorch-refinenet-unknown-Default/example-app)
frame #5: __libc_start_main + 0xf0 (0x7f361cdd0830 in /lib/x86_64-linux-gnu/libc.so.6)
frame #6: _start + 0x29 (0x44dae9 in /data_2/everyday/0618/build-libtorch-refinenet-unknown-Default/example-app)

然后，他说你把正常的和这个不正常的二进制流打印出来比较就是了。接在上面函数之后这么写的：

             char* buf = nullptr;
                string path_pt = "/data_2/everyday/0623/nice/jiami2";
                string key = "QO1##gX@@3";

                FILE* file = fopen(path_pt.c_str(), "r");
                fseek(file, 0, SEEK_END);
                unsigned size = ftell(file);
                fseek(file, 0, SEEK_SET);
                buf = (char*)malloc(size);
                memset(buf,0,size);
                unsigned int i=0,j=0;
                while(!feof(file))
                {
                    char ch = (fgetc(file))^(key[i>=key.length()?i=0:i++]);
                    buf[j++]=ch;
                }

                strstreambuf sbuf(buf, sizeof(buf));
               std::istream ins(&sbuf);
                vector<char> v_ch;
                char ch_tmp[1];
                while (ins.read(ch_tmp, 1))
                {
                    v_ch.push_back(ch_tmp[0]);
                    std::cout<<(int)ch_tmp[0] <<std::endl;
                }
               while(1);

发现v_ch只要8位！！！为啥？？？前几位的文件流char如下：
80
75
3
4
0
0
8
8
0
0
59
124
-51
80
0
如果说遇到0就算截止，可是前8个本身就有0啊为啥不在之前截止。。无解。
可能是sizeof的问题，然后改成如下：

 char* buf = nullptr;
                string path_pt = "/data_2/everyday/0623/nice/jiami2";
                string key = "QO1##gX@@3";

                FILE* file = fopen(path_pt.c_str(), "r");
                fseek(file, 0, SEEK_END);
                unsigned size = ftell(file);
                fseek(file, 0, SEEK_SET);
                buf = (char*)malloc(size);
                memset(buf,0,size);
                unsigned int i=0,j=0;
                while(!feof(file))
                {
                    char ch = (fgetc(file))^(key[i>=key.length()?i=0:i++]);
                    buf[j++]=ch;
                }

                strstreambuf sbuf(buf, j);
               std::istream ins(&sbuf);
                vector<char> v_ch;
                char ch_tmp[1];
                while (ins.read(ch_tmp, 1))
                {
                    v_ch.push_back(ch_tmp[0]);
                    std::cout<<(int)ch_tmp[0] <<std::endl;
                }
               while(1);

这回v_ch也有2千万的数据了，大小也是一样的，可以torch还是报错！！！然后jiamin说比较输出来的哪里不一样。因为cout打印终端，但是可以直接运行可执行文件让把输出的内容保存在本地。比如我这个工程构建会生成example的可执行文件。
然后，直接运行可执行文件：

./example &> 1.txt

就可以把打印的内容直接保存在文本。
这样，正常的和不正常的都保存，然后比较两个文本，发现不正常的最后比正常的多了一个数字。然后我把上面的代码改成：
strstreambuf sbuf(buf, j-1);
就可以了！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！
这个小功能花费了一天半的时间，其实我一开始的半天就搞通可以流读到load里面，然后都是在解决上面的问题，std::filebuf的那个问题还不知道哪里有问题！

==============================================================================================
下面给出加密的代码：

#include <iostream>
#include <fstream>
#include<memory.h>

void ShowUsage()
{
    std::cout << "Usage for encryption" << std::endl;
    std::cout << "path_pt" << std::endl;
    std::cout << "path_save_encryption" << std::endl;
    //    std::cout << "length" << std::endl;
    std::cout << "example:
 ./encryption /data_2/small.pt /data_2/small_en" << std::endl;
    return;
}

std::string rand_str(const int len)
{
    srand( (unsigned)time( NULL ) );
    std::string str;
    char symbol[20] = {‘!‘,‘@‘,‘#‘,‘$‘,‘%‘,‘^‘,‘&‘,‘*‘,‘(‘,‘)‘,‘?‘,‘.‘,‘<‘,‘>‘,‘~‘,‘-‘,‘+‘,‘{‘,‘]‘,‘[‘};
    int i;
    char ch;
    for(i=0;i<len;++i)
    {
        switch((rand()%4))
        {
        case 1:
            ch=‘A‘+rand()%26;
            break;
        case 2:
            ch=‘a‘+rand()%26;
            break;
        case 3:
            ch=symbol[rand()%20];
            break;
        default:
            ch=‘0‘+rand()%10;
            break;
        }
        str += ch;
    }

    return str;
}



int main(int argc, char *argv[])
{
    if(argc < 3)
    {
        ShowUsage();
        return -1;
    }

    std::string str_len = "16";
    std::string path_pt = argv[1];
    std::string path_save_jiami = argv[2];
    if(argc >= 4)
    {
        str_len = argv[3];
    }

    //    std::string path_pt = "/data_2/everyday/0622/00000.pt";
    //    std::string path_save_jiami = "/data_2/everyday/0623/nice/jiami2";
    //    std::string str_len = "10";

    int len = std::stoi(str_len);
    std::string key = rand_str(len);

    std::filebuf in;
    std::filebuf outbuf;
    outbuf.open(path_save_jiami,std::ios::out);
    if (!in.open(path_pt, std::ios::in)) {
        std::cout << "fail to open model pt" << std::endl;
        std::cout << "please check path: " << path_pt << std::endl;
        return 0;
    }

    FILE *in_file;
    in_file=fopen(path_pt.c_str(),"rb");//以读的方式打开二进制文件
    char ch=fgetc(in_file);

    int i = 0;
    while(!feof(in_file))
    {
        ch = ch^key[i>=key.size()?i=0:i++];
        outbuf.sputc(ch);
        ch=fgetc(in_file);
    }
    outbuf.sputc(ch);
    outbuf.close();

    std::cout<<"
success crerate encryption model!" << std::endl;
    std::cout<<"key=
"<< key << std::endl;

    return 0;
}

解密的话上面零零散散的也有。