feat: collect and display usage stats

ssh/direct_tcp_handler.go

@@ -77,7 +77,13 @@ func (s *Server) handleDirectTCP(srv *ssh.Server, conn *gossh.ServerConn, newCha
 		defer dconn.Close()
 		defer ch.Close()
 
-		if _, err := io.Copy(ch, dconn); err != nil {
+		reader := &instrumentedReader{
+			internal: dconn,
+			stats:    s.opts.Stats,
+			name:     StatTotalRxBytes,
+		}
+
+		if _, err := io.Copy(ch, reader); err != nil {
 			if errors.Is(err, net.ErrClosed) {
 				return
 			}
@@ -90,7 +96,13 @@ func (s *Server) handleDirectTCP(srv *ssh.Server, conn *gossh.ServerConn, newCha
 		defer dconn.Close()
 		defer ch.Close()
 
-		if _, err := io.Copy(dconn, ch); err != nil {
+		writer := &instrumentedWriter{
+			internal: dconn,
+			stats:    s.opts.Stats,
+			name:     StatTotalTxBytes,
+		}
+
+		if _, err := io.Copy(writer, ch); err != nil {
 			s.log("[ERROR] %+v", errors.WithStack(err))
 		}
 	}()

ssh/options.go

@@ -1,6 +1,10 @@
 package ssh
 
-import "log"
+import (
+	"log"
+
+	"forge.cadoles.com/wpetit/rebound/stat"
+)
 
 type Options struct {
 	Logger     func(message string, args ...any)
@@ -8,6 +12,7 @@ type Options struct {
 	PublicPort uint   `env:"PUBLIC_PORT"`
 	PublicHost string `env:"PUBLIC_HOST"`
 	HostKey    string `env:"HOST_KEY"`
+	Stats      *stat.Store
 }
 
 type OptionFunc func(*Options)
@@ -19,6 +24,7 @@ func DefaultOptions() *Options {
 		PublicPort: 2222,
 		PublicHost: "127.0.0.1",
 		HostKey:    "./host.key",
+		Stats:      stat.NewStore(),
 	}
 }
 
@@ -51,3 +57,9 @@ func WithLogger(logger func(message string, args ...any)) func(*Options) {
 		opts.Logger = logger
 	}
 }
+
+func WithStats(stats *stat.Store) func(*Options) {
+	return func(opts *Options) {
+		opts.Stats = stats
+	}
+}

ssh/request_handler.go

@@ -91,6 +91,8 @@ func (s *Server) handleRequest(ctx ssh.Context, srv *ssh.Server, req *gossh.Requ
 			return false, []byte{}
 		}
 
+		s.opts.Stats.Add(StatTotalOpenedTunnels, 1, 0)
+
 		destPort := 1
 
 		s.requestHandlerLock.Lock()

ssh/stats.go

@@ -0,0 +1,43 @@
+package ssh
+
+import (
+	"io"
+
+	"forge.cadoles.com/wpetit/rebound/stat"
+)
+
+const (
+	StatTotalOpenedTunnels = "total_opened_tunnels"
+	StatTotalTxBytes       = "total_tx_bytes"
+	StatTotalRxBytes       = "total_rx_bytes"
+)
+
+type instrumentedWriter struct {
+	name     string
+	stats    *stat.Store
+	internal io.Writer
+}
+
+// Write implements io.Writer.
+func (w *instrumentedWriter) Write(p []byte) (n int, err error) {
+	n, err = w.internal.Write(p)
+	w.stats.Add(w.name, float64(n), 0)
+	return n, err
+}
+
+var _ io.Writer = &instrumentedWriter{}
+
+type instrumentedReader struct {
+	name     string
+	stats    *stat.Store
+	internal io.Reader
+}
+
+// Read implements io.Reader.
+func (w *instrumentedReader) Read(p []byte) (n int, err error) {
+	n, err = w.internal.Read(p)
+	w.stats.Add(w.name, float64(n), 0)
+	return n, err
+}
+
+var _ io.Reader = &instrumentedReader{}