add buf hook

1 files changed, 51 insertions, 0 deletions

diff --git a/pkg/server/hook/buf.go b/pkg/server/hook/buf.go
new file mode 100644
index 0000000..0306c73
--- /dev/null
+++ b/pkg/server/hook/buf.go
@@ -0,0 +1,51 @@
+package hook
+
+import (
+	"bytes"
+	"time"
+	"tunnel/pkg/server/queue"
+)
+
+const bufDelay = 1 * time.Millisecond
+
+func buffering(rq, wq queue.Q) error {
+	for b := range rq {
+		var w bytes.Buffer
+
+		t := time.NewTimer(bufDelay)
+		stop := false
+
+		for !stop {
+			select {
+			case x := <-rq:
+				if len(x) == 0 {
+					stop = true
+					break
+				}
+
+				if w.Len() == 0 {
+					w.Write(b)
+				}
+
+				w.Write(x)
+
+			case <-t.C:
+				stop = true
+			}
+		}
+
+		if w.Len() > 0 {
+			wq <- w.Bytes()
+		} else {
+			wq <- b
+		}
+
+		t.Stop()
+	}
+
+	return nil
+}
+
+func init() {
+	registerFunc("buf", buffering)
+}