1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
#include <poll.h>
#include <memory>
#include <smb2/smb2.h>
#include <smb2/libsmb2.h>

// 统一的同步回调数据结构
struct WaitCbData {
  bool is_finished{false};
  int status{0};
  void* ptr{nullptr};
};

// 公用的同步等待轮询函数
auto WaitForReply(struct smb2_context* smb2, 
                  std::shared_ptr<WaitCbData> cb_data) -> int {
  while (!cb_data->is_finished) {
    struct pollfd pfd{};
    pfd.fd = smb2_get_fd(smb2);
    pfd.events = smb2_which_events(smb2);

    if (poll(&pfd, 1, 1000) < 0) {
      return -1;
    }
    if (pfd.revents == 0) {
      continue;
    }
    if (smb2_service(smb2, pfd.revents) < 0) {
      return -1;
    }
  }
  return 0;
}

// =========================================================
// 1. Close 接口封装 (解决 UAF 与 Double-Free)
// =========================================================

auto CloseFileCb(struct smb2_context* smb2, int status, 
                 void* command_data, void* private_data) -> void {
  auto* sp = static_cast<std::shared_ptr<WaitCbData>*>(private_data);
  if (sp == nullptr) return;

  if (status != SMB2_STATUS_SHUTDOWN && *sp) {
    (*sp)->is_finished = true;
    (*sp)->status = status;
  }
  delete sp; 
}

auto Smb2CloseSync(struct smb2_context* smb2, struct smb2fh* fh) -> int {
  auto cb_data = std::make_shared<WaitCbData>();
  auto* cb_ptr = new std::shared_ptr<WaitCbData>(cb_data);

  if (smb2_close_async(smb2, fh, CloseFileCb, cb_ptr) < 0) {
    delete cb_ptr;
    return -1;
  }

  if (WaitForReply(smb2, cb_data) < 0) {
    return -1; // 异常退出，保留 cb_ptr，由 CloseFileCb 安全回收
  }
  return cb_data->status;
}

// =========================================================
// 2. OpenDir 接口封装 (解决 UAF 与所有权转移)
// =========================================================

auto OpenDirCb(struct smb2_context* smb2, int status, 
               void* command_data, void* private_data) -> void {
  auto* sp = static_cast<std::shared_ptr<WaitCbData>*>(private_data);
  if (sp == nullptr) return;

  if (status != SMB2_STATUS_SHUTDOWN && *sp) {
    (*sp)->is_finished = true;
    (*sp)->status = status;
    (*sp)->ptr = command_data;
  }

  // 成功获取目录指针时，禁止释放 sp，必须将生命周期委托给 smb2dir
  if (status == 0 && command_data != nullptr) {
    return; 
  }
  delete sp; 
}

auto Smb2OpenDirSync(struct smb2_context* smb2, 
                     const char* path) -> struct smb2dir* {
  auto cb_data = std::make_shared<WaitCbData>();
  auto* cb_ptr = new std::shared_ptr<WaitCbData>(cb_data);

  auto* pdu = smb2_opendir_async_pdu(smb2, path, OpenDirCb, cb_ptr, nullptr);
  if (pdu == nullptr) {
    delete cb_ptr;
    return nullptr;
  }

  if (WaitForReply(smb2, cb_data) < 0) {
    smb2_free_pdu(smb2, pdu);
    return nullptr;
  }

  auto* dir = static_cast<struct smb2dir*>(cb_data->ptr);
  if (dir != nullptr) {
    // 利用 C++17 无捕获 Lambda 隐式转换为 C 函数指针，安全处理资源释放
    dir->free_cb_data = [](void* ptr) {
      delete static_cast<std::shared_ptr<WaitCbData>*>(ptr);
    };
  }
  
  smb2_free_pdu(smb2, pdu);
  return dir;
}

// =========================================================
// 3. ShareEnum 接口封装 (解决超时内存泄漏)
// =========================================================

auto ShareEnumCb(struct smb2_context* smb2, int status, 
                 void* command_data, void* private_data) -> void {
  auto* sp = static_cast<std::shared_ptr<WaitCbData>*>(private_data);
  if (sp == nullptr) return;

  if (status != SMB2_STATUS_SHUTDOWN && *sp) {
    (*sp)->is_finished = true;
    (*sp)->status = status;
    (*sp)->ptr = command_data;
  }
  delete sp; 
}

auto Smb2ShareEnumSync(struct smb2_context* smb2, 
                       enum SHARE_INFO_enum level) -> struct srvsvc_NetrShareEnum_rep* {
  auto cb_data = std::make_shared<WaitCbData>();
  auto* cb_ptr = new std::shared_ptr<WaitCbData>(cb_data);

  if (smb2_share_enum_async(smb2, level, ShareEnumCb, cb_ptr) < 0) {
    delete cb_ptr;
    return nullptr;
  }

  // 原生代码在此处出错时直接 return NULL 导致失控。
  // 此处即使 return，cb_ptr 也必然会被迟滞触发的 ShareEnumCb 回收，完美避免泄露。
  if (WaitForReply(smb2, cb_data) < 0) {
    return nullptr; 
  }

  return static_cast<struct srvsvc_NetrShareEnum_rep*>(cb_data->ptr);
}