diff --git a/src/utils/hub.js b/src/utils/hub.js index 9361767..9cacb2e 100644 --- a/src/utils/hub.js +++ b/src/utils/hub.js @@ -151,23 +151,26 @@ class FileResponse { } /** - * Determines whether the given string is a valid HTTP or HTTPS URL. - * @param {string|URL} string The string to test for validity as an HTTP or HTTPS URL. + * Determines whether the given string is a valid URL. + * @param {string|URL} string The string to test for validity as an URL. + * @param {string[]} [protocols=null] A list of valid protocols. If specified, the protocol must be in this list. * @param {string[]} [validHosts=null] A list of valid hostnames. If specified, the URL's hostname must be in this list. - * @returns {boolean} True if the string is a valid HTTP or HTTPS URL, false otherwise. + * @returns {boolean} True if the string is a valid URL, false otherwise. */ -function isValidHttpUrl(string, validHosts = null) { - // https://stackoverflow.com/a/43467144 +function isValidUrl(string, protocols = null, validHosts = null) { let url; try { url = new URL(string); } catch (_) { return false; } + if (protocols && !protocols.includes(url.protocol)) { + return false; + } if (validHosts && !validHosts.includes(url.hostname)) { return false; } - return url.protocol === "http:" || url.protocol === "https:"; + return true; } /** @@ -178,7 +181,7 @@ function isValidHttpUrl(string, validHosts = null) { */ export async function getFile(urlOrPath) { - if (env.useFS && !isValidHttpUrl(urlOrPath)) { + if (env.useFS && !isValidUrl(urlOrPath, ['http:', 'https:', 'blob:'])) { return new FileResponse(urlOrPath); } else if (typeof process !== 'undefined' && process?.release?.name === 'node') { @@ -189,7 +192,7 @@ export async function getFile(urlOrPath) { headers.set('User-Agent', `transformers.js/${version}; is_ci/${IS_CI};`); // Check whether we are making a request to the Hugging Face Hub. - const isHFURL = isValidHttpUrl(urlOrPath, ['huggingface.co', 'hf.co']); + const isHFURL = isValidUrl(urlOrPath, ['http:', 'https:'], ['huggingface.co', 'hf.co']); if (isHFURL) { // If an access token is present in the environment variables, // we add it to the request headers. @@ -433,7 +436,7 @@ export async function getModelFile(path_or_repo_id, filename, fatal = true, opti if (env.allowLocalModels) { // Accessing local models is enabled, so we try to get the file locally. // If request is a valid HTTP URL, we skip the local file check. Otherwise, we try to get the file locally. - const isURL = isValidHttpUrl(requestURL); + const isURL = isValidUrl(requestURL, ['http:', 'https:']); if (!isURL) { try { response = await getFile(localPath); diff --git a/tests/utils.test.js b/tests/utils.test.js index 504d968..10c1bf2 100644 --- a/tests/utils.test.js +++ b/tests/utils.test.js @@ -1,6 +1,7 @@ import { AutoProcessor } from '../src/transformers.js'; import { mel_filter_bank } from '../src/utils/audio.js'; +import { getFile } from '../src/utils/hub.js'; import { MAX_TEST_EXECUTION_TIME } from './init.js'; @@ -42,4 +43,15 @@ describe('Utilities', () => { }, MAX_TEST_EXECUTION_TIME); }); + + describe('Hub utilities', () => { + + it('Read data from blob', async () => { + const blob = new Blob(['Hello, world!'], { type: 'text/plain' }); + const blobUrl = URL.createObjectURL(blob); + const data = await getFile(blobUrl); + expect(await data.text()).toBe('Hello, world!'); + }); + + }); });