From 5a58b597dcb30cd54cc5a97ddb5be132fe4fec04 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Matej=20Va=C5=A1ek?= <matejvasek@gmail.com>
Date: Tue, 17 Sep 2024 16:38:22 +0200
Subject: [PATCH] Improve socat conn success detection
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

It appears that Go's MatchReader() returns true only if some new lines
are written to the writer after the searched word. It mostly all right
because socat actually writes several more lines to stderr. However this
new implementation is better it should signal immediately when searched
patter is written to the writer.

Signed-off-by: Matej Vašek <matejvasek@gmail.com>
---
 pkg/k8s/dialer.go           | 74 +++++++++++++++++++++++++++++++------
 pkg/k8s/dialer_unit_test.go | 24 ++++++++++++
 2 files changed, 87 insertions(+), 11 deletions(-)
 create mode 100644 pkg/k8s/dialer_unit_test.go

diff --git a/pkg/k8s/dialer.go b/pkg/k8s/dialer.go
index 0e42352c17..28463af28f 100644
--- a/pkg/k8s/dialer.go
+++ b/pkg/k8s/dialer.go
@@ -1,7 +1,6 @@
 package k8s
 
 import (
-	"bufio"
 	"bytes"
 	"context"
 	"errors"
@@ -107,19 +106,72 @@ func (c *contextDialer) DialContext(ctx context.Context, network string, addr st
 	}
 }
 
-var connSuccessfulRE = regexp.MustCompile("successfully connected")
-
 // Creates io.Writer which closes connectSuccess channel when string "successfully connected" is written to it.
 func detectConnSuccess(connectSuccess chan struct{}) io.Writer {
-	pr, pw := io.Pipe()
-	go func() {
-		ok := connSuccessfulRE.MatchReader(bufio.NewReader(pr))
-		if ok {
-			close(connectSuccess)
+	return newKMPWriter("successfully connected", connectSuccess)
+}
+
+// kmpWriter is a writer that search word w using KMP algorithm in the text written to the writer.
+// When searched word w appears in the input the channel ch is closed.
+// This can be used to detect if particular character sequence was written to the writer.
+type kmpWriter struct {
+	w     string
+	k     int
+	t     []int
+	ch    chan<- struct{}
+	found bool
+}
+
+func newKMPWriter(w string, ch chan<- struct{}) *kmpWriter {
+	// Building KMP table.
+	t := make([]int, len(w)+1)
+	t[0] = -1
+	pos := 1
+	cnd := 0
+	for pos < len(w) {
+		if w[pos] == w[cnd] {
+			t[pos] = t[cnd]
+		} else {
+			t[pos] = cnd
+			for cnd >= 0 && w[pos] != w[cnd] {
+				cnd = t[cnd]
+			}
 		}
-		_, _ = io.Copy(io.Discard, pr)
-	}()
-	return pw
+		pos++
+		cnd++
+	}
+	t[pos] = cnd
+
+	return &kmpWriter{
+		w:  w,
+		t:  t,
+		ch: ch,
+	}
+}
+
+func (d *kmpWriter) Write(s []byte) (n int, err error) {
+	if d.found {
+		return len(s), nil
+	}
+	j := 0
+	for j < len(s) {
+		if d.w[d.k] == s[j] {
+			j++
+			d.k++
+			if d.k == len(d.w) {
+				d.found = true
+				close(d.ch)
+				return len(s), nil
+			}
+		} else {
+			d.k = d.t[d.k]
+			if d.k < 0 {
+				j++
+				d.k++
+			}
+		}
+	}
+	return j, nil
 }
 
 var (
diff --git a/pkg/k8s/dialer_unit_test.go b/pkg/k8s/dialer_unit_test.go
new file mode 100644
index 0000000000..c5511e66d7
--- /dev/null
+++ b/pkg/k8s/dialer_unit_test.go
@@ -0,0 +1,24 @@
+package k8s
+
+import (
+	"testing"
+)
+
+func TestDetectSocatSuccess(t *testing.T) {
+	ch := make(chan struct{}, 1)
+	w := detectConnSuccess(ch)
+	_, err := w.Write([]byte("some data successucces"))
+	if err != nil {
+		t.Fatal(err)
+	}
+	_, err = w.Write([]byte("sfully connected"))
+	if err != nil {
+		t.Fatal(err)
+	}
+	select {
+	case <-ch:
+		t.Log("OK")
+	default:
+		t.Error("NOK")
+	}
+}