feat: adding WATER transport field

getlantern · Jul 2, 2024 · 87041eb · 87041eb
1 parent 2b0018f
commit 87041eb
Show file tree

Hide file tree

Showing 5 changed files with 24 additions and 14 deletions.
diff --git a/http-proxy/main.go b/http-proxy/main.go
@@ -186,8 +186,9 @@ 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 = flag.String("water-wasm", "", "Base64 encoded WASM for WATER")
+	waterAddr      = flag.String("water-addr", "", "Address at which to listen for WATER connections.")
+	waterWASM      = flag.String("water-wasm", "", "Base64 encoded WASM for WATER")
+	waterTransport = flag.String("water-transport", "", "WATER based transport name")
 
 	track = flag.String("track", "", "The track this proxy is running on")
 )
@@ -474,6 +475,7 @@ func main() {
 		AlgenevaAddr:                       *algenevaAddr,
 		WaterAddr:                          *waterAddr,
 		WaterWASM:                          *waterWASM,
+		WaterTransport:                     *waterTransport,
 	}
 	if *maxmindLicenseKey != "" {
 		log.Debug("Will use Maxmind for geolocating clients")

diff --git a/http_proxy.go b/http_proxy.go
@@ -185,8 +185,9 @@ type Proxy struct {
 
 	AlgenevaAddr string
 
-	WaterAddr string
-	WaterWASM string
+	WaterAddr      string
+	WaterWASM      string
+	WaterTransport string
 
 	throttleConfig throttle.Config
 	instrument     instrument.Instrument

diff --git a/water/listener.go b/water/listener.go
@@ -15,7 +15,7 @@ var log = golog.LoggerFor("water")
 
 // NewWATERListener creates a WATER listener
 // Currently water doesn't support customized TCP connections and we need to listen and receive requests directly from the WATER listener
-func NewWATERListener(ctx context.Context, address, wasm string) (net.Listener, error) {
+func NewWATERListener(ctx context.Context, transport, address, wasm string) (net.Listener, error) {
 	decodedWASM, err := base64.StdEncoding.DecodeString(wasm)
 	if err != nil {
 		log.Errorf("failed to decode WASM base64: %v", err)
@@ -25,7 +25,7 @@ func NewWATERListener(ctx context.Context, address, wasm string) (net.Listener,
 	cfg := &water.Config{
 		TransportModuleBin: decodedWASM,
 		//NetworkListener:     baseListener,
-		OverrideLogger: slog.New(newLogHandler(log)),
+		OverrideLogger: slog.New(newLogHandler(log, transport)),
 	}
 
 	waterListener, err := cfg.ListenContext(ctx, "tcp", address)

diff --git a/water/listener_test.go b/water/listener_test.go
@@ -33,7 +33,7 @@ func TestWATERListener(t *testing.T) {
 		TransportModuleBin: wasm,
 	}
 
-	ll, err := NewWATERListener(ctx, "127.0.0.1:3000", b64WASM)
+	ll, err := NewWATERListener(ctx, "reverse_v0", "127.0.0.1:3000", b64WASM)
 	require.Nil(t, err)
 
 	messageRequest := "hello"

diff --git a/water/logger.go b/water/logger.go
@@ -12,15 +12,16 @@ import (
 // slogHandler is a Handler that implements the slog.Handler interface
 // and writes log records to a golog.Logger.
 type slogHandler struct {
-	logger   golog.Logger
-	minLevel slog.Level
-	opts     slog.HandlerOptions
-	attrs    string
-	groups   []string
+	logger    golog.Logger
+	transport string
+	minLevel  slog.Level
+	opts      slog.HandlerOptions
+	attrs     string
+	groups    []string
 }
 
-func newLogHandler(logger golog.Logger) *slogHandler {
-	return &slogHandler{logger: logger}
+func newLogHandler(logger golog.Logger, transport string) *slogHandler {
+	return &slogHandler{logger: logger, transport: transport}
 }
 
 // Enabled reports whether the handler handles records at the given level.
@@ -63,6 +64,12 @@ func (h *slogHandler) Handle(ctx context.Context, record slog.Record) error {
 	}
 
 	messageBuilder := new(strings.Builder)
+	messageBuilder.WriteString(record.Time.String())
+	messageBuilder.WriteString("-")
+	messageBuilder.WriteString(record.Level.String())
+	messageBuilder.WriteString(" ")
+	messageBuilder.WriteString(h.transport)
+	messageBuilder.WriteString(": ")
 	messageBuilder.WriteString(record.Message)
 	messageBuilder.WriteString(" ")
 	record.Attrs(func(attr slog.Attr) bool {