hotfix: fixed tcp_proxy onData bug

modules/tcp_proxy/tcp_proxy_script.go

@@ -55,23 +55,67 @@ func (s *TcpProxyScript) OnData(from, to net.Addr, data []byte, callback func(ca
 			log.Error("error while executing onData callback: %s", err)
 			return nil
 		} else if ret != nil {
-			// thanks to @LucasParsy for his code and patience :)
-			if array, ok := ret.([]interface{}); ok {
-				result := make([]byte, len(array))
-				for i, v := range array {
-					if num, ok := v.(float64); ok && num >= 0 && num <= 255 {
-						result[i] = byte(num)
-					} else {
-						log.Error("array element at index %d is not a valid byte value %+v", i, v)
-						return nil
-					}
-				}
-
-				return result
-			} else {
-				log.Error("error while casting exported value to array of interface: value = %+v error = %+v", ret, err)
-			}
+			return toByteArray(ret)
 		}
 	}
 	return nil
 }
+
+func toByteArray(ret interface{}) []byte {
+	// Handle different array types that otto.Export() might return
+	switch v := ret.(type) {
+	case []interface{}:
+		// Mixed type array
+		result := make([]byte, len(v))
+		for i, elem := range v {
+			if num, ok := toNumber(elem); ok && num >= 0 && num <= 255 {
+				result[i] = byte(num)
+			} else {
+				log.Error("array element at index %d is not a valid byte value %+v", i, elem)
+				return nil
+			}
+		}
+		return result
+	case []int64:
+		// Array of integers
+		result := make([]byte, len(v))
+		for i, num := range v {
+			if num >= 0 && num <= 255 {
+				result[i] = byte(num)
+			} else {
+				log.Error("array element at index %d is not a valid byte value %d", i, num)
+				return nil
+			}
+		}
+		return result
+	case []float64:
+		// Array of floats
+		result := make([]byte, len(v))
+		for i, num := range v {
+			if num >= 0 && num <= 255 {
+				result[i] = byte(num)
+			} else {
+				log.Error("array element at index %d is not a valid byte value %f", i, num)
+				return nil
+			}
+		}
+		return result
+	default:
+		log.Error("unexpected array type returned from onData: %T, value = %+v", ret, ret)
+		return nil
+	}
+}
+
+// toNumber tries to convert an interface{} to a float64
+func toNumber(v interface{}) (float64, bool) {
+	switch n := v.(type) {
+	case float64:
+		return n, true
+	case int64:
+		return float64(n), true
+	case int:
+		return float64(n), true
+	default:
+		return 0, false
+	}
+}

modules/tcp_proxy/tcp_proxy_script_test.go

@@ -0,0 +1,118 @@
+package tcp_proxy
+
+import (
+	"net"
+	"testing"
+
+	"github.com/evilsocket/islazy/plugin"
+)
+
+func TestOnData_NoReturn(t *testing.T) {
+	jsCode := `
+		function onData(from, to, data, callback) {
+			// don't return anything
+		}
+	`
+
+	plug, err := plugin.Parse(jsCode)
+	if err != nil {
+		t.Fatalf("Failed to parse plugin: %v", err)
+	}
+
+	script := &TcpProxyScript{
+		Plugin:   plug,
+		doOnData: plug.HasFunc("onData"),
+	}
+
+	from := &net.TCPAddr{IP: net.ParseIP("192.168.1.1"), Port: 1234}
+	to := &net.TCPAddr{IP: net.ParseIP("192.168.1.2"), Port: 5678}
+	data := []byte("test data")
+
+	result := script.OnData(from, to, data, nil)
+	if result != nil {
+		t.Errorf("Expected nil result when callback returns nothing, got %v", result)
+	}
+}
+
+func TestOnData_ReturnsArrayOfIntegers(t *testing.T) {
+	jsCode := `
+		function onData(from, to, data, callback) {
+			// Return modified data as array of integers
+			return [72, 101, 108, 108, 111]; // "Hello" in ASCII
+		}
+	`
+
+	plug, err := plugin.Parse(jsCode)
+	if err != nil {
+		t.Fatalf("Failed to parse plugin: %v", err)
+	}
+
+	script := &TcpProxyScript{
+		Plugin:   plug,
+		doOnData: plug.HasFunc("onData"),
+	}
+
+	from := &net.TCPAddr{IP: net.ParseIP("192.168.1.1"), Port: 1234}
+	to := &net.TCPAddr{IP: net.ParseIP("192.168.1.2"), Port: 5678}
+	data := []byte("test data")
+
+	result := script.OnData(from, to, data, nil)
+	expected := []byte("Hello")
+
+	if result == nil {
+		t.Fatal("Expected non-nil result when callback returns array of integers")
+	}
+
+	if len(result) != len(expected) {
+		t.Fatalf("Expected result length %d, got %d", len(expected), len(result))
+	}
+
+	for i, b := range result {
+		if b != expected[i] {
+			t.Errorf("Expected byte at index %d to be %d, got %d", i, expected[i], b)
+		}
+	}
+}
+
+func TestOnData_ReturnsDynamicArray(t *testing.T) {
+	jsCode := `
+		function onData(from, to, data, callback) {
+			var result = [];
+			for (var i = 0; i < data.length; i++) {
+				result.push((data[i] + 1) % 256);
+			}
+			return result;
+		}
+	`
+
+	plug, err := plugin.Parse(jsCode)
+	if err != nil {
+		t.Fatalf("Failed to parse plugin: %v", err)
+	}
+
+	script := &TcpProxyScript{
+		Plugin:   plug,
+		doOnData: plug.HasFunc("onData"),
+	}
+
+	from := &net.TCPAddr{IP: net.ParseIP("192.168.1.1"), Port: 1234}
+	to := &net.TCPAddr{IP: net.ParseIP("192.168.1.2"), Port: 5678}
+	data := []byte{10, 20, 30, 40, 255}
+
+	result := script.OnData(from, to, data, nil)
+	expected := []byte{11, 21, 31, 41, 0} // 255 + 1 = 256 % 256 = 0
+
+	if result == nil {
+		t.Fatal("Expected non-nil result when callback returns array of integers")
+	}
+
+	if len(result) != len(expected) {
+		t.Fatalf("Expected result length %d, got %d", len(expected), len(result))
+	}
+
+	for i, b := range result {
+		if b != expected[i] {
+			t.Errorf("Expected byte at index %d to be %d, got %d", i, expected[i], b)
+		}
+	}
+}