Modify ReturnStatement to prevent a UaF

src/nodes/statements.ts

@@ -1,6 +1,6 @@
 import * as ts from 'typescript';
 import { CodeTemplate, CodeTemplateFactory, CTemplateBase } from '../template';
-import { CProgram, IScope} from '../program';
+import { CProgram, IScope } from '../program';
 import { ArrayType, NumberVarType, StringVarType } from '../types/ctypes';
 import { CVariable, CVariableDeclaration, CVariableDestructors } from './variable';
 import { CExpression, CCondition } from './expressions';
@@ -71,17 +71,48 @@ export class CEmptyStatement {
 }
 
 @CodeTemplate(`
-{destructors}
-return {expression};
+{#if returnTypeAndVar}
+    {
+        {returnTypeAndVar} = {expression};
+        {destructors}
+        return {returnTemp};
+    }
+{#else}
+    {destructors}
+    return {expression};
+{/if}
 `, ts.SyntaxKind.ReturnStatement)
 export class CReturnStatement extends CTemplateBase {
     public expression: CExpression;
     public destructors: CVariableDestructors;
     public retVarName: string = null;
+    public returnTypeAndVar: string = null;
+    public returnTemp: string = null;
     public closureParams: { name: string, value: CExpression }[] = [];
+    doCheckVarNeeded(scope: IScope, node: ts.ReturnStatement): boolean {
+        let s = scope.root.memoryManager.getDestructorsForScope(node);
+        let result = false;
+        for (let e of s) {
+            result = result || new RegExp('\\W' + e.varName + '\\W').test(' ' + node.getFullText() + ' ');
+        }
+        return result;
+    }
     constructor(scope: IScope, node: ts.ReturnStatement) {
         super();
+        this.returnTemp = scope.root.symbolsHelper.addTemp(node, 'returnVal');
+        let returnType = scope.root.typeHelper.getCType(node.expression);
+        if (this.doCheckVarNeeded(scope, node)) {
+            // todo: make this less hackyy
+            let fakeVar = new CVariable(scope, '__fake', returnType, { removeStorageSpecifier: true, arraysToPointers: true });;
+            let fakeVarType = fakeVar.resolve().slice(0, -6).trim();
+            this.returnTypeAndVar = fakeVarType;
+            if (this.returnTypeAndVar.indexOf('{var}') == -1) {
+                this.returnTypeAndVar += ' {var}';
+            }
+            this.returnTypeAndVar = this.returnTypeAndVar.replace('{var}', this.returnTemp);
+        }
         this.expression = CodeTemplateFactory.createForNode(scope, node.expression);
+
         this.destructors = new CVariableDestructors(scope, node);
     }
 }
@@ -380,8 +411,7 @@ export class CForInStatement extends CTemplateBase implements IScope {
         else
             this.init = new CElementAccess(scope, node.initializer);
 
-        if (node.statement.kind == ts.SyntaxKind.Block)
-        {
+        if (node.statement.kind == ts.SyntaxKind.Block) {
             let block = <ts.Block>node.statement;
             for (let s of block.statements)
                 this.statements.push(CodeTemplateFactory.createForNode(this, s));

src/nodes/variable.ts

@@ -113,45 +113,43 @@ export class CVariableAllocation extends CTemplateBase {
 }
 
 @CodeTemplate(`
-{#statements}
-    {arrayDestructors => for (gc_i = 0; gc_i < ({this} ? {this}->size : 0); gc_i++) free((void*){this}->data[gc_i]);\n}
-    {destructors => free({this});\n}
-    {#if gcArraysCVarName}
-        for (gc_i = 0; gc_i < {gcArraysCVarName}->size; gc_i++) {
-            for (gc_j = 0; gc_j < ({gcArraysCVarName}->data[gc_i] ? {gcArraysCVarName}->data[gc_i]->size : 0); gc_j++)
-                free((void*){gcArraysCVarName}->data[gc_i]->data[gc_j]);\n
-            free({gcArraysCVarName}->data[gc_i] ? {gcArraysCVarName}->data[gc_i]->data : NULL);
-            free({gcArraysCVarName}->data[gc_i]);
-        }
-        free({gcArraysCVarName}->data);
-        free({gcArraysCVarName});
-    {/if}
-    {#if gcArraysVarName}
-        for (gc_i = 0; gc_i < {gcArraysVarName}->size; gc_i++) {
-            free({gcArraysVarName}->data[gc_i]->data);
-            free({gcArraysVarName}->data[gc_i]);
-        }
-        free({gcArraysVarName}->data);
-        free({gcArraysVarName});
-    {/if}
-    {#if gcDictsVarName}
-        for (gc_i = 0; gc_i < {gcDictsVarName}->size; gc_i++) {
-            free({gcDictsVarName}->data[gc_i]->index->data);
-            free({gcDictsVarName}->data[gc_i]->index);
-            free({gcDictsVarName}->data[gc_i]->values->data);
-            free({gcDictsVarName}->data[gc_i]->values);
-            free({gcDictsVarName}->data[gc_i]);
-        }
-        free({gcDictsVarName}->data);
-        free({gcDictsVarName});
-    {/if}
-    {#if gcVarName}
-        for (gc_i = 0; gc_i < {gcVarName}->size; gc_i++)
-            free({gcVarName}->data[gc_i]);
-        free({gcVarName}->data);
-        free({gcVarName});
-    {/if}
-{/statements}`
+{arrayDestructors => for (gc_i = 0; gc_i < ({this} ? {this}->size : 0); gc_i++) free((void*){this}->data[gc_i]);\n}
+{destructors => free({this});\n}
+{#if gcArraysCVarName}
+    for (gc_i = 0; gc_i < {gcArraysCVarName}->size; gc_i++) {
+        for (gc_j = 0; gc_j < ({gcArraysCVarName}->data[gc_i] ? {gcArraysCVarName}->data[gc_i]->size : 0); gc_j++)
+            free((void*){gcArraysCVarName}->data[gc_i]->data[gc_j]);\n
+        free({gcArraysCVarName}->data[gc_i] ? {gcArraysCVarName}->data[gc_i]->data : NULL);
+        free({gcArraysCVarName}->data[gc_i]);
+    }
+    free({gcArraysCVarName}->data);
+    free({gcArraysCVarName});
+{/if}
+{#if gcArraysVarName}
+    for (gc_i = 0; gc_i < {gcArraysVarName}->size; gc_i++) {
+        free({gcArraysVarName}->data[gc_i]->data);
+        free({gcArraysVarName}->data[gc_i]);
+    }
+    free({gcArraysVarName}->data);
+    free({gcArraysVarName});
+{/if}
+{#if gcDictsVarName}
+    for (gc_i = 0; gc_i < {gcDictsVarName}->size; gc_i++) {
+        free({gcDictsVarName}->data[gc_i]->index->data);
+        free({gcDictsVarName}->data[gc_i]->index);
+        free({gcDictsVarName}->data[gc_i]->values->data);
+        free({gcDictsVarName}->data[gc_i]->values);
+        free({gcDictsVarName}->data[gc_i]);
+    }
+    free({gcDictsVarName}->data);
+    free({gcDictsVarName});
+{/if}
+{#if gcVarName}
+    for (gc_i = 0; gc_i < {gcVarName}->size; gc_i++)
+        free({gcVarName}->data[gc_i]);
+    free({gcVarName}->data);
+    free({gcVarName});
+{/if}`
 )
 export class CVariableDestructors extends CTemplateBase {
     public gcVarName: string = null;

src/template.ts

@@ -1,4 +1,4 @@
-import {IScope} from './program';
+import { IScope } from './program';
 
 interface INode { kind: number, getText(): string };
 
@@ -16,7 +16,7 @@ export class CodeTemplateFactory {
             : "/* Unsupported node: " + node.getText().replace(/[\n\s]+/g, ' ') + " */;\n";
     }
     public static templateToString(template: string | CTemplateBase) {
-        return typeof(template) === "string" ? template : template.resolve();
+        return typeof (template) === "string" ? template : template.resolve();
     }
 }
 
@@ -42,6 +42,8 @@ export function CodeTemplate(tempString: string, nodeKind?: number | number[]):
                     nodeKindTemplates[nk] = newConstructor;
         }
 
+        newConstructor.prototype = target.prototype;
+
         return newConstructor;
 
     } as ClassDecorator;
@@ -123,15 +125,15 @@ function processTemplate(template: string, args: string | CTemplateBase): [strin
             let index = -1;
             while ((index = template.indexOf("{" + k + "}")) > -1) {
                 let spaces = '';
-                while (template.length > index && template[index-1] == ' ') {
+                while (template.length > index && template[index - 1] == ' ') {
                     index--;
-                    spaces+=' ';
+                    spaces += ' ';
                 }
                 let value = args[k];
                 if (value && value.resolve)
                     value = value.resolve();
                 if (value && typeof value === 'string')
-                    value = value.replace(/\n/g, '\n'+spaces);
+                    value = value.replace(/\n/g, '\n' + spaces);
                 template = template.replace("{" + k + "}", () => value);
                 replaced = true;
             }
@@ -226,7 +228,7 @@ function replaceArray(data, k, array, statements) {
             if (elementsResolved != "")
                 elementsResolved += separator;
             elementsResolved += resolvedElement;
-        } 
+        }
 
     }
 

tests/regression/regression1.res.expect

@@ -0,0 +1 @@
+1

tests/regression/regression1.ts

@@ -0,0 +1,7 @@
+// regression test for #74
+function return_from_obj() {
+    let obj = { key: 1 };
+    return obj.key;
+}
+
+console.log(return_from_obj());