/**
 * Emulates stream.setTimeout() behaviour, but respects outgoing data
 * too.
 *
 * @param {Number} timeout Milliseconds
 */
exports.attach = function(stream, timeout) {
    var timer;
    var emitTimeout = function() {
        timer = undefined;
        stream.emit('timeout');
    };
    var updateTimer = function() {
        if (timer)
            clearTimeout(timer);
        timer = setTimeout(emitTimeout, timeout);
    };

    var oldWrite = stream.write;
    stream.write = function() {
        updateTimer();
        oldWrite.apply(this, arguments);
    };
    stream.addListener('data', updateTimer);
    stream.addListener('close', function() {
		if (timer)
		    clearTimeout(timer);
	    stream.write = oldWrite;
    });
    stream.addListener('end', function() {
    	if (timer)
    	    clearTimeout(timer);
        stream.write = oldWrite;
    });
};

@@ -427,16 +427,20 @@
 %0A    %7D;%0A
+    
 %0A    var
@@ -576,80 +576,21 @@
     
-stream.addListener('data', updateTimer);%0A    stream.addListener('close',
+var cleanup =
  fun
@@ -599,17 +599,20 @@
 ion() %7B%0A
-%09
+    
 %09if (tim
@@ -619,14 +619,14 @@
 er)%0A
-%09%09
     
+%09%09
 clea
@@ -646,173 +646,218 @@
 r);%0A
-%09    stream.write = oldWrite;%0A    %7D);%0A    stream.addListener('end', function() %7B%0A    %09if (timer)%0A    %09    clearTimeout(timer);%0A   
+    %09if ( oldWrite != stream.write) %7B%0A        %09oldWrite = stream.write;%0A    %09%7D%0A    %7D;%0A    stream.addListener('data', updateTimer);%0A    stream.addListener('close', cleanup);%0A
     
- 
 stream.
-write = oldWrite;%0A    %7D
+addListener('end', cleanup
 );%0A%7D

refactor the cleanup code