fix: #112 #120

svrxjs · Oct 18, 2019 · b6113e1 · b6113e1
1 parent 91298f6
commit b6113e1
Show file tree

Hide file tree

Showing 2 changed files with 69 additions and 3 deletions.
diff --git a/packages/svrx/__tests__/spec/svrx.injector.js b/packages/svrx/__tests__/spec/svrx.injector.js
@@ -167,4 +167,63 @@ describe('Injector', () => {
         .end(done);
     });
   });
+
+  describe('Insert case', () => {
+    const svrx = new Svrx({
+      port: 8001,
+      middlewares: [
+        {
+          async onRoute(ctx, next) {
+            switch (ctx.url) {
+              case '/doc-write-body':
+                ctx.set('Content-Type', 'text/html');
+                ctx.body = `<head>
+                  <script>document.wirte('<body></body>')</script>
+                  </head>
+                  <body>mark</body>`;
+                break;
+              case '/only-body':
+                ctx.set('Content-Type', 'text/html');
+                ctx.body = '<body></body>';
+                break;
+              case '/only-body-stream':
+                ctx.set('Content-Type', 'text/html');
+                ctx.body = bufferToStream(Buffer.from('<body></body>'));
+                break;
+              case '/none-of-both':
+                ctx.set('Content-Type', 'text/html');
+                ctx.body = bufferToStream(Buffer.from('Hello,World'));
+                break;
+              default:
+                next();
+            }
+          },
+        },
+      ],
+    });
+
+
+    it('only last body will be inject', (done) => {
+      request(svrx.callback())
+        .get('/doc-write-body')
+        .expect(new RegExp('mark<script'))
+        .expect(new RegExp(`href="${svrx.config.get('urls.style')}"`))
+        .end(done);
+    });
+
+    it('missing head should inject style at body', (done) => {
+      request(svrx.callback())
+        .get('/only-body')
+        .expect(new RegExp(`src="${svrx.config.get('urls.script')}"`))
+        .expect(new RegExp(`href="${svrx.config.get('urls.style')}"`))
+        .end(done);
+    });
+    it('stream:missing head should inject style at body', (done) => {
+      request(svrx.callback())
+        .get('/only-body-stream')
+        .expect(new RegExp(`src="${svrx.config.get('urls.script')}"`))
+        .expect(new RegExp(`href="${svrx.config.get('urls.style')}"`))
+        .end(done);
+    });
+  });
 });
diff --git a/packages/svrx/lib/injector/index.js b/packages/svrx/lib/injector/index.js
@@ -137,15 +137,22 @@ module.exports = class Injector {
       `<script async src="${config.get('urls.script')}" charset="UTF-8"></script></body>`,
     ];
     const replaceStyle = [
-      '</head>',
-      `<link rel="stylesheet" type="text/css" href="${config.get('urls.style')}"/></head>`,
+      /<\/(head|body)>/,
+      `<link rel="stylesheet" type="text/css" href="${config.get('urls.style')}"/></$1>`,
     ];
 
     if (body instanceof Buffer) {
       body = body.toString('utf8');
     }
     if (typeof body === 'string') {
-      return this._replace(body.replace(...replaceScript).replace(...replaceStyle), 'string');
+      // fix: #120
+      const lastIndex = body.lastIndexOf(replaceScript[0]);
+      if (lastIndex > 0) {
+        body = body.slice(0, lastIndex)
+          + replaceScript[1]
+          + body.slice(replaceScript[0].length + lastIndex);
+      }
+      return this._replace(body.replace(...replaceStyle), 'string');
     }
     if (isReadableStream(body)) {
       return this._replace(