From ae53a1d5fdad9fc0d622c55a9c4e36a352bc4bc8 Mon Sep 17 00:00:00 2001
From: Ryan McGuire <rmcguire@libretechconsulting.com>
Date: Wed, 29 Jan 2025 19:37:23 -0500
Subject: [PATCH] Invalid AWN Request Fixer

---
 main.go                                |  2 +-
 pkg/ambient/ambienthttp/ambienthttp.go | 25 ++++++++++++++++++-------
 2 files changed, 19 insertions(+), 8 deletions(-)

diff --git a/main.go b/main.go
index 8bd3392..cc00c8d 100644
--- a/main.go
+++ b/main.go
@@ -47,7 +47,7 @@ func main() {
 					HandlerFunc: aw.GetAWNHandlerFunc(ctx),
 				},
 			},
-			CustomListener: ambienthttp.NewLFStrippingListener(ctx,
+			CustomListener: ambienthttp.NewAWNMutatingListener(ctx,
 				awConfig.HTTP.Listen), // Necessary to fix certain bad AWN firmware
 			HealthChecks: []srv.HealthCheckFunc{
 				// TODO: Implement
diff --git a/pkg/ambient/ambienthttp/ambienthttp.go b/pkg/ambient/ambienthttp/ambienthttp.go
index 6517e54..4c38a8c 100644
--- a/pkg/ambient/ambienthttp/ambienthttp.go
+++ b/pkg/ambient/ambienthttp/ambienthttp.go
@@ -32,8 +32,9 @@ type LFStrippingListener struct {
 }
 
 type LFStrippingConn struct {
-	net.Conn
+	ctx    context.Context
 	reader io.Reader
+	net.Conn
 }
 
 func (l *LFStrippingListener) WrapConn(conn net.Conn) net.Conn {
@@ -49,34 +50,47 @@ func (l *LFStrippingListener) WrapConn(conn net.Conn) net.Conn {
 		// Only restore newline if not a bad request
 		if !badReqURI.Match(line) {
 			newData = append(newData, '\n')
+		} else {
+			zerolog.Ctx(l.ctx).Trace().Bytes("line", line).
+				Msg("malformed request found, stripped 0x0a")
+		}
+
+		if len(line) == 0 {
+			break
 		}
 	}
 	if scanner.Err() != nil {
 		zerolog.Ctx(l.ctx).Err(scanner.Err()).Send()
 	}
 
+	zerolog.Ctx(l.ctx).Trace().
+		Int("numBytes", len(newData)).
+		Bytes("request", newData).
+		Msg("stripping conn complete")
+
 	// Use a multi-reader to prepend the modified request
 	finalReader := io.MultiReader(bytes.NewReader(newData), conn)
 
 	return &LFStrippingConn{
 		Conn:   conn,
+		ctx:    l.ctx,
 		reader: finalReader,
 	}
 }
 
-func NewLFStrippingListener(ctx context.Context, listen string) net.Listener {
-	// Create the underlying TCP listener
+func NewAWNMutatingListener(ctx context.Context, listen string) net.Listener {
 	rawListener, err := net.Listen("tcp", listen)
 	if err != nil {
 		panic(err)
 	}
+
+	// Encapsulate the raw listener with ours
 	return &LFStrippingListener{
 		Listener: rawListener,
 		ctx:      ctx,
 	}
 }
 
-// Accept waits for and returns the next connection to the listener.
 func (l *LFStrippingListener) Accept() (net.Conn, error) {
 	conn, err := l.Listener.Accept()
 	if err != nil {
@@ -85,13 +99,10 @@ func (l *LFStrippingListener) Accept() (net.Conn, error) {
 	return l.WrapConn(conn), nil
 }
 
-// Close closes the listener.
-// Any blocked Accept operations will be unblocked and return errors.
 func (l *LFStrippingListener) Close() error {
 	return l.Listener.Close()
 }
 
-// Addr returns the listener's network address.
 func (l *LFStrippingListener) Addr() net.Addr {
 	return l.Listener.Addr()
 }