From 5d88f5ce45cc1e45fc515f5d24c80afd22921c88 Mon Sep 17 00:00:00 2001
From: WendelHime <6754291+WendelHime@users.noreply.github.com>
Date: Mon, 24 Jun 2024 12:15:03 -0300
Subject: [PATCH] fix: receiving WASM as base64 and decoding before using it

---
 http-proxy/main.go    | 13 ++-----------
 http_proxy.go         |  2 +-
 water/reverse.go      | 11 +++++++++--
 water/reverse_test.go |  5 ++++-
 4 files changed, 16 insertions(+), 15 deletions(-)

diff --git a/http-proxy/main.go b/http-proxy/main.go
index c6268763..90c4b768 100644
--- a/http-proxy/main.go
+++ b/http-proxy/main.go
@@ -187,7 +187,7 @@ var (
 	algenevaAddr = flag.String("algeneva-addr", "", "Address at which to listen for algenAddr connections.")
 
 	waterAddr = flag.String("water-addr", "", "Address at which to listen for WATER connections.")
-	waterWASM []byte
+	waterWASM = flag.String("water-wasm", "", "Base64 encoded WASM for WATER")
 
 	track = flag.String("track", "", "The track this proxy is running on")
 )
@@ -197,15 +197,6 @@ const (
 )
 
 func main() {
-	flag.Func("water-wasm", "WATER WASM bytes.", func(s string) error {
-		if s == "" {
-			return nil
-		}
-
-		waterWASM = []byte(s)
-		return nil
-	})
-
 	iniflags.SetAllowUnknownFlags(true)
 	iniflags.Parse()
 	if *version {
@@ -482,7 +473,7 @@ func main() {
 		BroflakeKey:                        os.Getenv("BROFLAKE_KEY"),
 		AlgenevaAddr:                       *algenevaAddr,
 		WaterAddr:                          *waterAddr,
-		WaterWASM:                          waterWASM,
+		WaterWASM:                          *waterWASM,
 	}
 	if *maxmindLicenseKey != "" {
 		log.Debug("Will use Maxmind for geolocating clients")
diff --git a/http_proxy.go b/http_proxy.go
index df6bec6a..e88c5bb7 100644
--- a/http_proxy.go
+++ b/http_proxy.go
@@ -186,7 +186,7 @@ type Proxy struct {
 	AlgenevaAddr string
 
 	WaterAddr string
-	WaterWASM []byte
+	WaterWASM string
 
 	throttleConfig throttle.Config
 	instrument     instrument.Instrument
diff --git a/water/reverse.go b/water/reverse.go
index 43601779..2b11240b 100644
--- a/water/reverse.go
+++ b/water/reverse.go
@@ -2,6 +2,7 @@ package water
 
 import (
 	"context"
+	"encoding/base64"
 	"net"
 	"sync"
 
@@ -21,9 +22,15 @@ type llistener struct {
 	closeError   error
 }
 
-func NewReverseListener(ctx context.Context, address string, wasm []byte) (*llistener, error) {
+func NewReverseListener(ctx context.Context, address string, wasm string) (*llistener, error) {
+	decodedWASM, err := base64.StdEncoding.DecodeString(wasm)
+	if err != nil {
+		log.Errorf("failed to decode WASM base64: %v", err)
+		return nil, err
+	}
+
 	cfg := &water.Config{
-		TransportModuleBin: wasm,
+		TransportModuleBin: decodedWASM,
 	}
 
 	waterListener, err := cfg.ListenContext(ctx, "tcp", address)
diff --git a/water/reverse_test.go b/water/reverse_test.go
index b0cc9233..b20dd465 100644
--- a/water/reverse_test.go
+++ b/water/reverse_test.go
@@ -3,6 +3,7 @@ package water
 import (
 	"bytes"
 	"context"
+	"encoding/base64"
 	"net"
 	"os"
 	"testing"
@@ -20,13 +21,15 @@ func TestReverseListener(t *testing.T) {
 	wasm, err := os.ReadFile(wasmPath)
 	require.Nil(t, err)
 
+	b64WASM := base64.StdEncoding.EncodeToString(wasm)
+
 	ctx := context.Background()
 
 	cfg := &water.Config{
 		TransportModuleBin: wasm,
 	}
 
-	ll, err := NewReverseListener(ctx, addr, wasm)
+	ll, err := NewReverseListener(ctx, addr, b64WASM)
 	require.Nil(t, err)
 
 	messageRequest := "hello"