diff --git a/modules/tcp_proxy/tcp_proxy_script.go b/modules/tcp_proxy/tcp_proxy_script.go index 7121a1a4..60f103ea 100644 --- a/modules/tcp_proxy/tcp_proxy_script.go +++ b/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 + } +} diff --git a/modules/tcp_proxy/tcp_proxy_script_test.go b/modules/tcp_proxy/tcp_proxy_script_test.go new file mode 100644 index 00000000..cb09b052 --- /dev/null +++ b/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) + } + } +}