Cutlery/calibre
1
0
Fork 0
mirror of https://github.com/kovidgoyal/calibre.git synced 2025-07-09 03:04:10 -04:00
Code Issues Packages Projects Releases Wiki Activity

Use RAII for get/put file operations

Browse Source
This commit is contained in:
Kovid Goyal 2021-04-23 09:34:19 +05:30
parent c9227ccd69
commit 481474aae2
No known key found for this signature in database
GPG Key ID: 06BC317B515ACE7C
2 changed files with 75 additions and 89 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

11
src/calibre/devices/mtp/driver.py
View File

@ -561,6 +561,7 @@ class MTP_DEVICE(BASE):
def main(): def main():
import io
dev = MTP_DEVICE(None) dev = MTP_DEVICE(None)
dev.startup() dev.startup()
try: try:
@ -574,7 +575,15 @@ def main():
dev.set_progress_reporter(prints) dev.set_progress_reporter(prints)
dev.open(cd, None) dev.open(cd, None)
dev.filesystem_cache.dump() dev.filesystem_cache.dump()
print('Prefix for main mem:', dev.prefix_for_location(None)) print('Prefix for main mem:', dev.prefix_for_location(None), flush=True)
raw = os.urandom(32 * 1024)
f = dev.put_file(dev.filesystem_cache.entries[0], 'developing-mtp-driver.bin', io.BytesIO(raw), len(raw))
print('Put file:', f, flush=True)
buf = io.BytesIO()
dev.get_file(f.mtp_id_path, buf)
if buf.getvalue() != raw:
raise ValueError('Getting previously put file did not return expected data')
print('Successfully got previously put file', flush=True)
finally: finally:
dev.shutdown() dev.shutdown()

153
src/calibre/devices/mtp/windows/content_enumeration.cpp
View File

@ -503,49 +503,46 @@ wpd::get_filesystem(IPortableDevice *device, const wchar_t *storage_id, IPortabl
PyObject* PyObject*
wpd::get_file(IPortableDevice *device, const wchar_t *object_id, PyObject *dest, PyObject *callback) { // {{{ wpd::get_file(IPortableDevice *device, const wchar_t *object_id, PyObject *dest, PyObject *callback) { // {{{
IPortableDeviceContent *content = NULL; CComPtr<IPortableDeviceContent> content;
IPortableDeviceResources *resources = NULL; CComPtr<IPortableDeviceResources> resources;
IPortableDeviceProperties *devprops = NULL; CComPtr<IPortableDeviceProperties> devprops;
IPortableDeviceValues *values = NULL; CComPtr<IPortableDeviceValues> values;
IPortableDeviceKeyCollection *properties = NULL; CComPtr<IPortableDeviceKeyCollection> properties;
IStream *stream = NULL; CComPtr<IStream> stream;
HRESULT hr; HRESULT hr;
DWORD bufsize = 4096; DWORD bufsize = 4096;
char *buf = NULL;
ULONG bytes_read = 0, total_read = 0; ULONG bytes_read = 0, total_read = 0;
BOOL ok = FALSE;
PyObject *res = NULL;
ULONGLONG filesize = 0; ULONGLONG filesize = 0;
Py_BEGIN_ALLOW_THREADS; Py_BEGIN_ALLOW_THREADS;
hr = device->Content(&content); hr = device->Content(&content);
Py_END_ALLOW_THREADS; Py_END_ALLOW_THREADS;
if (FAILED(hr)) { hresult_set_exc("Failed to create content interface", hr); goto end; } if (FAILED(hr)) { hresult_set_exc("Failed to create content interface", hr); return NULL; }
Py_BEGIN_ALLOW_THREADS; Py_BEGIN_ALLOW_THREADS;
hr = content->Properties(&devprops); hr = content->Properties(&devprops);
Py_END_ALLOW_THREADS; Py_END_ALLOW_THREADS;
if (FAILED(hr)) { hresult_set_exc("Failed to get IPortableDeviceProperties interface", hr); goto end; } if (FAILED(hr)) { hresult_set_exc("Failed to get IPortableDeviceProperties interface", hr); return NULL; }
Py_BEGIN_ALLOW_THREADS; Py_BEGIN_ALLOW_THREADS;
hr = CoCreateInstance(CLSID_PortableDeviceKeyCollection, NULL, hr = properties.CoCreateInstance(CLSID_PortableDeviceKeyCollection, NULL, CLSCTX_INPROC_SERVER);
CLSCTX_INPROC_SERVER, IID_PPV_ARGS(&properties));
Py_END_ALLOW_THREADS; Py_END_ALLOW_THREADS;
if (FAILED(hr)) { hresult_set_exc("Failed to create filesystem properties collection", hr); goto end; } if (FAILED(hr)) { hresult_set_exc("Failed to create filesystem properties collection", hr); return NULL; }
hr = properties->Add(WPD_OBJECT_SIZE); hr = properties->Add(WPD_OBJECT_SIZE);
if (FAILED(hr)) { hresult_set_exc("Failed to add filesize property to properties collection", hr); goto end; } if (FAILED(hr)) { hresult_set_exc("Failed to add filesize property to properties collection", hr); return NULL; }
Py_BEGIN_ALLOW_THREADS; Py_BEGIN_ALLOW_THREADS;
hr = devprops->GetValues(object_id, properties, &values); hr = devprops->GetValues(object_id, properties, &values);
Py_END_ALLOW_THREADS; Py_END_ALLOW_THREADS;
if (FAILED(hr)) { hresult_set_exc("Failed to get filesize for object", hr); goto end; } if (FAILED(hr)) { hresult_set_exc("Failed to get filesize for object", hr); return NULL; }
hr = values->GetUnsignedLargeIntegerValue(WPD_OBJECT_SIZE, &filesize); hr = values->GetUnsignedLargeIntegerValue(WPD_OBJECT_SIZE, &filesize);
if (FAILED(hr)) { hresult_set_exc("Failed to get filesize from values collection", hr); goto end; } if (FAILED(hr)) { hresult_set_exc("Failed to get filesize from values collection", hr); return NULL; }
Py_BEGIN_ALLOW_THREADS; Py_BEGIN_ALLOW_THREADS;
hr = content->Transfer(&resources); hr = content->Transfer(&resources);
Py_END_ALLOW_THREADS; Py_END_ALLOW_THREADS;
if (FAILED(hr)) { hresult_set_exc("Failed to create resources interface", hr); goto end; } if (FAILED(hr)) { hresult_set_exc("Failed to create resources interface", hr); return NULL; }
Py_BEGIN_ALLOW_THREADS; Py_BEGIN_ALLOW_THREADS;
hr = resources->GetStream(object_id, WPD_RESOURCE_DEFAULT, STGM_READ, &bufsize, &stream); hr = resources->GetStream(object_id, WPD_RESOURCE_DEFAULT, STGM_READ, &bufsize, &stream);
@ -554,50 +551,41 @@ wpd::get_file(IPortableDevice *device, const wchar_t *object_id, PyObject *dest,
if (HRESULT_FROM_WIN32(ERROR_BUSY) == hr) { if (HRESULT_FROM_WIN32(ERROR_BUSY) == hr) {
PyErr_SetString(WPDFileBusy, "Object is in use"); PyErr_SetString(WPDFileBusy, "Object is in use");
} else hresult_set_exc("Failed to create stream interface to read from object", hr); } else hresult_set_exc("Failed to create stream interface to read from object", hr);
goto end; return NULL;
} }
buf = (char *)calloc(bufsize+10, 1); generic_raii<char*, PyMem_Free> buf(reinterpret_cast<char*>(PyMem_Malloc(bufsize)));
if (buf == NULL) { PyErr_NoMemory(); goto end; } if (!buf) return PyErr_NoMemory();
while (TRUE) { while (total_read < filesize) {
bytes_read = 0; bytes_read = 0;
Py_BEGIN_ALLOW_THREADS; Py_BEGIN_ALLOW_THREADS;
hr = stream->Read(buf, bufsize, &bytes_read); hr = stream->Read(buf.ptr(), bufsize, &bytes_read);
Py_END_ALLOW_THREADS; Py_END_ALLOW_THREADS;
total_read = total_read + bytes_read;
if (hr == STG_E_ACCESSDENIED) { if (hr == STG_E_ACCESSDENIED) {
PyErr_SetString(PyExc_IOError, "Read access is denied to this object"); break; PyErr_SetFromWindowsErr(ERROR_ACCESS_DENIED);
return NULL;
} else if (SUCCEEDED(hr)) { } else if (SUCCEEDED(hr)) {
if (bytes_read > 0) { if (bytes_read > 0) {
res = PyObject_CallMethod(dest, "write", "y#", buf, bytes_read); total_read = total_read + bytes_read;
if (res == NULL) break; pyobject_raii res(PyObject_CallMethod(dest, "write", "y#", buf.ptr(), bytes_read));
Py_DECREF(res); res = NULL; if (!res) { return NULL; }
if (callback != NULL) Py_XDECREF(PyObject_CallFunction(callback, "kK", total_read, filesize)); if (callback != NULL) {
pyobject_raii r(PyObject_CallFunction(callback, "kK", total_read, filesize));
} }
} else { hresult_set_exc("Failed to read file from device", hr); break; } }
} else { hresult_set_exc("Failed to read file from device", hr); return NULL; }
if (bytes_read == 0) { if (bytes_read == 0) {
ok = TRUE; pyobject_raii r(PyObject_CallMethod(dest, "flush", NULL));
Py_XDECREF(PyObject_CallMethod(dest, "flush", NULL));
break; break;
} }
} }
if (ok && total_read != filesize) { if (total_read < filesize) {
ok = FALSE;
PyErr_SetString(WPDError, "Failed to read all data from file"); PyErr_SetString(WPDError, "Failed to read all data from file");
return NULL;
} }
end:
if (content != NULL) content->Release();
if (devprops != NULL) devprops->Release();
if (resources != NULL) resources->Release();
if (stream != NULL) stream->Release();
if (values != NULL) values->Release();
if (properties != NULL) properties->Release();
if (buf != NULL) free(buf);
if (!ok) return NULL;
Py_RETURN_NONE; Py_RETURN_NONE;
} // }}} } // }}}
@ -684,34 +672,30 @@ end:
PyObject* PyObject*
wpd::put_file(IPortableDevice *device, const wchar_t *parent_id, const wchar_t *name, PyObject *src, unsigned PY_LONG_LONG size, PyObject *callback) { // {{{ wpd::put_file(IPortableDevice *device, const wchar_t *parent_id, const wchar_t *name, PyObject *src, unsigned PY_LONG_LONG size, PyObject *callback) { // {{{
IPortableDeviceContent *content = NULL; CComPtr<IPortableDeviceContent> content;
IPortableDeviceValues *values = NULL; CComPtr<IPortableDeviceValues> values;
IPortableDeviceProperties *devprops = NULL; CComPtr<IPortableDeviceProperties> devprops;
IPortableDeviceKeyCollection *properties = NULL; CComPtr<IPortableDeviceKeyCollection> properties;
IStream *temp = NULL; CComPtr<IStream> temp;
IPortableDeviceDataStream *dest = NULL; CComPtr<IPortableDeviceDataStream> dest;
char *buf = NULL;
wchar_t *newid = NULL;
PyObject *ans = NULL, *raw;
HRESULT hr; HRESULT hr;
DWORD bufsize = 0; DWORD bufsize = 0;
BOOL ok = FALSE;
Py_ssize_t bytes_read = 0; Py_ssize_t bytes_read = 0;
ULONG bytes_written = 0, total_written = 0; ULONG bytes_written = 0, total_written = 0;
values = create_object_properties(parent_id, name, WPD_CONTENT_TYPE_GENERIC_FILE, size); values = create_object_properties(parent_id, name, WPD_CONTENT_TYPE_GENERIC_FILE, size);
if (values == NULL) goto end; if (!values) return NULL;
Py_BEGIN_ALLOW_THREADS; Py_BEGIN_ALLOW_THREADS;
hr = device->Content(&content); hr = device->Content(&content);
Py_END_ALLOW_THREADS; Py_END_ALLOW_THREADS;
if (FAILED(hr)) { hresult_set_exc("Failed to create content interface", hr); goto end; } if (FAILED(hr)) { hresult_set_exc("Failed to create content interface", hr); return NULL; }
hr = content->Properties(&devprops); hr = content->Properties(&devprops);
if (FAILED(hr)) { hresult_set_exc("Failed to get IPortableDeviceProperties interface", hr); goto end; } if (FAILED(hr)) { hresult_set_exc("Failed to get IPortableDeviceProperties interface", hr); return NULL; }
properties = create_filesystem_properties_collection(); properties = create_filesystem_properties_collection();
if (properties == NULL) goto end; if (!properties) return NULL;
Py_BEGIN_ALLOW_THREADS; Py_BEGIN_ALLOW_THREADS;
hr = content->CreateObjectWithPropertiesAndData(values, &temp, &bufsize, NULL); hr = content->CreateObjectWithPropertiesAndData(values, &temp, &bufsize, NULL);
@ -720,53 +704,46 @@ wpd::put_file(IPortableDevice *device, const wchar_t *parent_id, const wchar_t *
if (HRESULT_FROM_WIN32(ERROR_BUSY) == hr) { if (HRESULT_FROM_WIN32(ERROR_BUSY) == hr) {
PyErr_SetString(WPDFileBusy, "Object is in use"); PyErr_SetString(WPDFileBusy, "Object is in use");
} else hresult_set_exc("Failed to create stream interface to write to object", hr); } else hresult_set_exc("Failed to create stream interface to write to object", hr);
goto end; return NULL;
} }
hr = temp->QueryInterface(IID_PPV_ARGS(&dest)); hr = temp->QueryInterface(IID_PPV_ARGS(&dest));
if (FAILED(hr)) { hresult_set_exc("Failed to create IPortableDeviceStream", hr); goto end; } if (FAILED(hr)) { hresult_set_exc("Failed to create IPortableDeviceStream", hr); return NULL; }
while(TRUE) { while(true) {
raw = PyObject_CallMethod(src, "read", "k", bufsize); #define ABORT { dest->Revert(); return NULL; }
if (raw == NULL) break; pyobject_raii raw(PyObject_CallMethod(src, "read", "k", bufsize));
PyBytes_AsStringAndSize(raw, &buf, &bytes_read); if (!raw) ABORT;
char *buffer;
if (PyBytes_AsStringAndSize(raw.ptr(), &buffer, &bytes_read) == -1) ABORT;
if (bytes_read > 0) { if (bytes_read > 0) {
Py_BEGIN_ALLOW_THREADS; Py_BEGIN_ALLOW_THREADS;
hr = dest->Write(buf, (ULONG)bytes_read, &bytes_written); hr = dest->Write(buffer, (ULONG)bytes_read, &bytes_written);
Py_END_ALLOW_THREADS; Py_END_ALLOW_THREADS;
Py_DECREF(raw); if (hr == STG_E_MEDIUMFULL) { PyErr_SetString(WPDError, "Cannot write to device as it is full"); ABORT; }
if (hr == STG_E_MEDIUMFULL) { PyErr_SetString(WPDError, "Cannot write to device as it is full"); break; } if (hr == STG_E_ACCESSDENIED) { PyErr_SetExcFromWindowsErr(WPDError, ERROR_ACCESS_DENIED); ABORT; }
if (hr == STG_E_ACCESSDENIED) { PyErr_SetString(WPDError, "Cannot write to file as access is denied"); break; } if (hr == STG_E_WRITEFAULT) { PyErr_SetString(WPDError, "Cannot write to file as there was a disk I/O error"); ABORT; }
if (hr == STG_E_WRITEFAULT) { PyErr_SetString(WPDError, "Cannot write to file as there was a disk I/O error"); break; } if (FAILED(hr)) { hresult_set_exc("Cannot write to file", hr); ABORT; }
if (FAILED(hr)) { hresult_set_exc("Cannot write to file", hr); break; } if (bytes_written != bytes_read) { PyErr_SetString(WPDError, "Writing to file failed, not all bytes were written"); ABORT; }
if (bytes_written != bytes_read) { PyErr_SetString(WPDError, "Writing to file failed, not all bytes were written"); break; }
total_written += bytes_written; total_written += bytes_written;
if (callback != NULL) Py_XDECREF(PyObject_CallFunction(callback, "kK", total_written, size)); if (callback != NULL) { pyobject_raii r(PyObject_CallFunction(callback, "kK", total_written, size)); }
} else Py_DECREF(raw); }
if (bytes_read == 0) { ok = TRUE; break; } if (bytes_read == 0) { break; }
#undef ABORT
} }
if (!ok) {dest->Revert(); goto end;}
Py_BEGIN_ALLOW_THREADS; Py_BEGIN_ALLOW_THREADS;
hr = dest->Commit(STGC_DEFAULT); hr = dest->Commit(STGC_DEFAULT);
Py_END_ALLOW_THREADS; Py_END_ALLOW_THREADS;
if (FAILED(hr)) { hresult_set_exc("Failed to write data to file, commit failed", hr); goto end; } if (FAILED(hr)) { hresult_set_exc("Failed to write data to file, commit failed", hr); return NULL; }
if (callback != NULL) Py_XDECREF(PyObject_CallFunction(callback, "kK", total_written, size)); if (callback != NULL) Py_XDECREF(PyObject_CallFunction(callback, "kK", total_written, size));
com_wchar_raii newid;
Py_BEGIN_ALLOW_THREADS; Py_BEGIN_ALLOW_THREADS;
hr = dest->GetObjectID(&newid); hr = dest->GetObjectID(newid.unsafe_address());
Py_END_ALLOW_THREADS; Py_END_ALLOW_THREADS;
if (FAILED(hr)) { hresult_set_exc("Failed to get id of newly created file", hr); goto end; } if (FAILED(hr)) { hresult_set_exc("Failed to get id of newly created file", hr); return NULL; }
ans = get_object_properties(devprops, properties, newid); return get_object_properties(devprops, properties, newid.ptr());
end:
if (content != NULL) content->Release();
if (values != NULL) values->Release();
if (devprops != NULL) devprops->Release();
if (properties != NULL) properties->Release();
if (temp != NULL) temp->Release();
if (dest != NULL) dest->Release();
if (newid != NULL) CoTaskMemFree(newid);
return ans;
} // }}} } // }}}