diff --git a/packages/react-router-server/src/server/index.tsx b/packages/react-router-server/src/server/index.tsx index c27c9a4c448..455756c6ec8 100644 --- a/packages/react-router-server/src/server/index.tsx +++ b/packages/react-router-server/src/server/index.tsx @@ -1,2 +1,5 @@ export { StartServer } from './StartServer' -export { transformStreamWithRouter } from './transformStreamWithRouter' +export { + transformStreamWithRouter, + transformReadableStreamWithRouter, +} from './transformStreamWithRouter' diff --git a/packages/react-router-server/src/server/transformStreamWithRouter.tsx b/packages/react-router-server/src/server/transformStreamWithRouter.tsx index 9fa338a4309..9e3584d4cd0 100644 --- a/packages/react-router-server/src/server/transformStreamWithRouter.tsx +++ b/packages/react-router-server/src/server/transformStreamWithRouter.tsx @@ -2,83 +2,120 @@ import { Transform } from 'stream' import type { AnyRouter } from '@tanstack/react-router' export function transformStreamWithRouter(router: AnyRouter) { - return transformStreamHtmlCallback(async () => { + const callbacks = transformHtmlCallbacks(injectorFromRouter(router)) + return new Transform({ + transform(chunk, _encoding, callback) { + callbacks + .transform(chunk, this.push.bind(this)) + .then(() => callback()) + .catch((err) => callback(err)) + }, + flush(callback) { + callbacks + .flush(this.push.bind(this)) + .then(() => callback()) + .catch((err) => callback(err)) + }, + }) +} + +export function transformReadableStreamWithRouter(router: AnyRouter) { + const callbacks = transformHtmlCallbacks(injectorFromRouter(router)) + return new TransformStream({ + transform(chunk, controller) { + return callbacks.transform(chunk, (chunkToPush) => { + controller.enqueue(chunkToPush) + return true + }) + }, + flush(controller) { + return callbacks.flush((chunkToPush) => { + controller.enqueue(chunkToPush) + return true + }) + }, + }) +} + +function injectorFromRouter(router: AnyRouter) { + return async () => { const injectorPromises = router.injectedHtml.map((d) => typeof d === 'function' ? d() : d, ) const injectors = await Promise.all(injectorPromises) router.injectedHtml = [] return injectors.join('') - }) + } } -function transformStreamHtmlCallback(injector: () => Promise) { + +// regex pattern for matching closing body and html tags +const patternBody = /(<\/body>)/ +const patternHtml = /(<\/html>)/ + +// regex pattern for matching closing tags +const pattern = /(<\/[a-zA-Z][\w:.-]*?>)/g + +const textDecoder = new TextDecoder() + +function transformHtmlCallbacks(injector: () => Promise) { let leftover = '' // If a closing tag is split across chunks, store the HTML to add after it // This expects that all the HTML that's added is closed properly let leftoverHtml = '' - return new Transform({ - transform(chunk, encoding, callback) { - const chunkString = leftover + chunk.toString() - - // regex pattern for matching closing body and html tags - const patternBody = /(<\/body>)/ - const patternHtml = /(<\/html>)/ + return { + async transform(chunk: any, push: (chunkToPush: string) => boolean) { + const chunkString = leftover + textDecoder.decode(chunk) const bodyMatch = chunkString.match(patternBody) const htmlMatch = chunkString.match(patternHtml) - injector() - .then((html) => { - // If a sequence was found - if (bodyMatch && htmlMatch && bodyMatch.index! < htmlMatch.index!) { - const bodyIndex = bodyMatch.index! + bodyMatch[0].length - const htmlIndex = htmlMatch.index! + htmlMatch[0].length - - // Add the arbitrary HTML before the closing body tag - const processed = - chunkString.slice(0, bodyIndex) + - html + - chunkString.slice(bodyIndex, htmlIndex) + - chunkString.slice(htmlIndex) + try { + const html = await injector() + // If a sequence was found + if (bodyMatch && htmlMatch && bodyMatch.index! < htmlMatch.index!) { + const bodyIndex = bodyMatch.index! + bodyMatch[0].length + const htmlIndex = htmlMatch.index! + htmlMatch[0].length - this.push(processed) - leftover = '' - } else { - // For all other closing tags, add the arbitrary HTML after them - const pattern = /(<\/[a-zA-Z][\w:.-]*?>)/g - let result - let lastIndex = 0 + // Add the arbitrary HTML before the closing body tag + const processed = + chunkString.slice(0, bodyIndex) + + html + + chunkString.slice(bodyIndex, htmlIndex) + + chunkString.slice(htmlIndex) - while ((result = pattern.exec(chunkString)) !== null) { - lastIndex = result.index + result[0].length - } + push(processed) + leftover = '' + } else { + // For all other closing tags, add the arbitrary HTML after them + let result + let lastIndex = 0 - // If a closing tag was found, add the arbitrary HTML and send it through - if (lastIndex > 0) { - const processed = - chunkString.slice(0, lastIndex) + html + leftoverHtml - this.push(processed) - leftover = chunkString.slice(lastIndex) - } else { - // If no closing tag was found, store the chunk to process with the next one - leftover = chunkString - leftoverHtml += html - } + while ((result = pattern.exec(chunkString)) !== null) { + lastIndex = result.index + result[0].length } - callback() - }) - .catch((err) => { - console.error(err) - callback(err) - }) + // If a closing tag was found, add the arbitrary HTML and send it through + if (lastIndex > 0) { + const processed = + chunkString.slice(0, lastIndex) + html + leftoverHtml + push(processed) + leftover = chunkString.slice(lastIndex) + } else { + // If no closing tag was found, store the chunk to process with the next one + leftover = chunkString + leftoverHtml += html + } + } + } catch (err) { + console.error(err) + throw err + } }, - flush(callback) { + async flush(push: (chunkToPush: string) => boolean) { if (leftover) { - this.push(leftover) + push(leftover) } - callback() }, - }) + } }