diff --git a/modules/tcp_proxy/tcp_proxy_script.go b/modules/tcp_proxy/tcp_proxy_script.go index 60f103ea..50956ea0 100644 --- a/modules/tcp_proxy/tcp_proxy_script.go +++ b/modules/tcp_proxy/tcp_proxy_script.go @@ -1,6 +1,7 @@ package tcp_proxy import ( + "encoding/json" "net" "strings" @@ -62,60 +63,29 @@ func (s *TcpProxyScript) OnData(from, to net.Addr, data []byte, callback func(ca } 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 - } -} + // this approach is a bit hacky but it handles all cases -// 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 + // serialize ret to JSON + if jsonData, err := json.Marshal(ret); err == nil { + // attempt to deserialize as []float64 + var back2Array []float64 + if err := json.Unmarshal(jsonData, &back2Array); err == nil { + result := make([]byte, len(back2Array)) + for i, num := range back2Array { + 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 + } else { + log.Error("failed to deserialize %+v to []float64: %v", ret, err) + } + } else { + log.Error("failed to serialize %+v to JSON: %v", ret, err) } + + return nil } diff --git a/modules/tcp_proxy/tcp_proxy_script_test.go b/modules/tcp_proxy/tcp_proxy_script_test.go index cb09b052..27bdc099 100644 --- a/modules/tcp_proxy/tcp_proxy_script_test.go +++ b/modules/tcp_proxy/tcp_proxy_script_test.go @@ -116,3 +116,54 @@ func TestOnData_ReturnsDynamicArray(t *testing.T) { } } } + +func TestOnData_ReturnsMixedArray(t *testing.T) { + jsCode := ` + function charToInt(value) { + return value.charCodeAt() + } + + function onData(from, to, data) { + st_data = String.fromCharCode.apply(null, data) + if( st_data.indexOf("mysearch") != -1 ) { + payload = "mypayload"; + st_data = st_data.replace("mysearch", payload); + res_int_arr = st_data.split("").map(charToInt) // []uint16 + res_int_arr[0] = payload.length + 1; // first index is float64 and rest []uint16 + return res_int_arr; + } + return data; + } + ` + + 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.6"), Port: 5678} + data := []byte("Hello mysearch world") + + result := script.OnData(from, to, data, nil) + expected := []byte("\x0aello mypayload world") + + 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) + } + } +}