From 6b1637d1d805627a482c8d81bd2ddfbb06e1aa87 Mon Sep 17 00:00:00 2001 From: William Petit Date: Sun, 24 Sep 2023 12:21:44 -0600 Subject: [PATCH] feat: collect and display usage stats --- .gitignore | 3 +- .goreleaser.yaml | 2 +- cmd/server/main.go | 6 ++ http/options.go | 11 +++ http/server.go | 43 +++++++- http/stats.go | 5 + http/templates/footer.html | 2 +- http/templates/index.html | 48 +++++++-- misc/packaging/openrc/rebound.conf | 1 + misc/packaging/systemd/rebound.env | 1 + options.go | 33 +++++-- server.go | 35 ++++++- ssh/direct_tcp_handler.go | 16 ++- ssh/options.go | 14 ++- ssh/request_handler.go | 2 + ssh/stats.go | 43 ++++++++ stat/store.go | 154 +++++++++++++++++++++++++++++ 17 files changed, 394 insertions(+), 25 deletions(-) create mode 100644 http/stats.go create mode 100644 ssh/stats.go create mode 100644 stat/store.go diff --git a/.gitignore b/.gitignore index 2e0d29d..54bb977 100644 --- a/.gitignore +++ b/.gitignore @@ -8,4 +8,5 @@ /dist tools/ /CHANGELOG.md -/.chglog \ No newline at end of file +/.chglog +/stats.json \ No newline at end of file diff --git a/.goreleaser.yaml b/.goreleaser.yaml index 2802be9..e6df99b 100644 --- a/.goreleaser.yaml +++ b/.goreleaser.yaml @@ -9,7 +9,7 @@ builds: ldflags: - -s - -w - - -X 'main.Version=${MKT_PROJECT_VERSION}' + - -X "main.Version={{ .Env.GORELEASER_CURRENT_TAG }}" gcflags: - -trimpath="${PWD}" asmflags: diff --git a/cmd/server/main.go b/cmd/server/main.go index 4631d77..fb7eca4 100644 --- a/cmd/server/main.go +++ b/cmd/server/main.go @@ -11,6 +11,8 @@ import ( "github.com/pkg/errors" ) +var Version string = "unknown" + func main() { opts := rebound.DefaultOptions() @@ -18,8 +20,12 @@ func main() { log.Fatalf("[ERROR] %+v", errors.WithStack(err)) } + opts.HTTP.TemplateData.Version = Version + server := rebound.NewServer( rebound.WithAddress(opts.Address), + rebound.WithStatsFile(opts.StatsFile), + rebound.WithStatsFileSaveInterval(opts.StatsFileSaveInterval), rebound.WithSSHOption( ssh.WithSockDir(opts.SSH.SockDir), ssh.WithPublicHost(opts.SSH.PublicHost), diff --git a/http/options.go b/http/options.go index b09cc1f..06ffc2e 100644 --- a/http/options.go +++ b/http/options.go @@ -2,16 +2,20 @@ package http import ( "log" + + "forge.cadoles.com/wpetit/rebound/stat" ) type Options struct { Logger func(message string, args ...any) CustomDir string `env:"CUSTOM_DIR"` TemplateData *TemplateData `envPrefix:"TEMPLATE_DATA_"` + Stats *stat.Store } type TemplateData struct { Title string `env:"TITLE"` + Version string SSHPublicHost string `env:"SSH_PUBLIC_HOST"` SSHPublicPort int `env:"SSH_PUBLIC_PORT"` } @@ -27,6 +31,7 @@ func DefaultOptions() *Options { SSHPublicHost: "127.0.0.1", SSHPublicPort: 2222, }, + Stats: stat.NewStore(), } } @@ -47,3 +52,9 @@ func WithTemplateData(templateData *TemplateData) func(*Options) { opts.TemplateData = templateData } } + +func WithStats(stats *stat.Store) func(*Options) { + return func(opts *Options) { + opts.Stats = stats + } +} diff --git a/http/server.go b/http/server.go index 9708cc8..876327c 100644 --- a/http/server.go +++ b/http/server.go @@ -3,9 +3,11 @@ package http import ( "bytes" "embed" + "fmt" "html/template" "io" "io/fs" + "log/slog" "net" "net/http" "os" @@ -29,8 +31,45 @@ type Server struct { templates template.Template } +var templateFuncs = template.FuncMap{ + "humanSize": func(b float64) string { + const unit = 1000 + if b < unit { + return fmt.Sprintf("%d B", int64(b)) + } + + div, exp := int64(unit), 0 + + for n := b / unit; n >= unit; n /= unit { + div *= unit + exp++ + } + + return fmt.Sprintf("%.1f %cB", + float64(b)/float64(div), "kMGTPE"[exp]) + }, +} + func (s *Server) serveHomepage(w http.ResponseWriter, r *http.Request) { - s.renderTemplate(w, "index", s.opts.TemplateData) + stats, err := s.opts.Stats.Snapshot() + if err != nil { + slog.Error("could not make stats snapshot", slog.Any("error", errors.WithStack(err))) + http.Error(w, http.StatusText(http.StatusInternalServerError), http.StatusInternalServerError) + + return + } + + data := struct { + TemplateData + Stats map[string]float64 + }{ + TemplateData: *s.opts.TemplateData, + Stats: stats, + } + + s.opts.Stats.Add(StatTotalPageView, 1, 0) + + s.renderTemplate(w, "index", data) } func (s *Server) Serve(l net.Listener) error { @@ -67,7 +106,7 @@ func (s *Server) Serve(l net.Listener) error { } func (s *Server) parseTemplates(fs fs.FS) error { - templates, err := template.ParseFS(fs, "templates/*.html") + templates, err := template.New("").Funcs(templateFuncs).ParseFS(fs, "templates/*.html") if err != nil { return errors.WithStack(err) } diff --git a/http/stats.go b/http/stats.go new file mode 100644 index 0000000..e2b0082 --- /dev/null +++ b/http/stats.go @@ -0,0 +1,5 @@ +package http + +const ( + StatTotalPageView = "total_page_view" +) diff --git a/http/templates/footer.html b/http/templates/footer.html index 6ea7bbd..bec5685 100644 --- a/http/templates/footer.html +++ b/http/templates/footer.html @@ -2,7 +2,7 @@ diff --git a/http/templates/index.html b/http/templates/index.html index 835c702..b8378ab 100644 --- a/http/templates/index.html +++ b/http/templates/index.html @@ -13,14 +13,46 @@

Bienvenue sur Rebound!

-
-

Rebound est un serveur SSH permettant de créer des tunnels TCP/IP éphémères et privés entre 2 machines positionnées - derrière un NAT.

-

Pour l'utiliser un simple client SSH suffit !

- 
ssh -R 0:127.0.0.1:<port> rebound@{{ .SSHPublicHost }} -p {{ .SSHPublicPort }}
-

Où <port> est à remplacer par le port du service - s'exécutant sur votre machine en local. -

Une fois connecté, suivez les instructions. 😉

+
+
+

Rebound est un serveur SSH permettant de créer des tunnels TCP/IP éphémères et privés entre 2 machines positionnées + derrière un NAT.

+

Pour l'utiliser un simple client SSH suffit !

+ 
ssh -R 0:127.0.0.1:<port> rebound@{{ .SSHPublicHost }} -p {{ .SSHPublicPort }}
+

Où <port> est à remplacer par le port du service + s'exécutant sur votre machine en local. +

Une fois connecté, suivez les instructions. 😉

+
+
+
+
+
+
+

En savoir plus

+
+ À venir... +
+
+
+

Statistiques

+ + + + + + + + + + + + + + + +
Total tunnels ouverts{{ index .Stats "total_opened_tunnels" }}
Total données entrantes{{ humanSize ( index .Stats "total_rx_bytes" ) }}
Total données sortantes{{ humanSize ( index .Stats "total_tx_bytes" ) }}
+
+
diff --git a/misc/packaging/openrc/rebound.conf b/misc/packaging/openrc/rebound.conf index dd658af..b7fa691 100644 --- a/misc/packaging/openrc/rebound.conf +++ b/misc/packaging/openrc/rebound.conf @@ -4,6 +4,7 @@ export REBOUND_SSH_PUBLIC_HOST=rebound export REBOUND_SSH_PUBLIC_PORT=2222 export REBOUND_SSH_SOCK_DIR=/var/lib/rebound/socks export REBOUND_SSH_HOST_KEY=/etc/rebound/host.key +export REBOUND_STATS_FILE=/var/lib/rebound/stats.json export REBOUND_HTTP_TEMPLATE_DATA_TITLE=Rebound export REBOUND_HTTP_TEMPLATE_DATA_SSH_PUBLIC_HOST=127.0.0.1 export REBOUND_HTTP_TEMPLATE_DATA_SSH_PUBLIC_PORT=8080 \ No newline at end of file diff --git a/misc/packaging/systemd/rebound.env b/misc/packaging/systemd/rebound.env index d1d362d..8bc306d 100644 --- a/misc/packaging/systemd/rebound.env +++ b/misc/packaging/systemd/rebound.env @@ -4,6 +4,7 @@ REBOUND_SSH_PUBLIC_HOST=rebound REBOUND_SSH_PUBLIC_PORT=8080 REBOUND_SSH_SOCK_DIR=/var/lib/rebound/socks REBOUND_SSH_HOST_KEY=/var/lib/rebound/host.key +REBOUND_STATS_FILE=/var/lib/rebound/stats.json REBOUND_HTTP_TEMPLATE_DATA_TITLE=Rebound REBOUND_HTTP_TEMPLATE_DATA_SSH_PUBLIC_HOST=127.0.0.1 REBOUND_HTTP_TEMPLATE_DATA_SSH_PUBLIC_PORT=8080 \ No newline at end of file diff --git a/options.go b/options.go index d184ad1..289ad87 100644 --- a/options.go +++ b/options.go @@ -2,6 +2,7 @@ package rebound import ( "log" + "time" "forge.cadoles.com/wpetit/rebound/http" "forge.cadoles.com/wpetit/rebound/ssh" @@ -10,10 +11,12 @@ import ( ) type Options struct { - Address string `env:"REBOUND_ADDRESS"` - Logger func(message string, args ...any) - SSH *ssh.Options `envPrefix:"REBOUND_SSH_"` - HTTP *http.Options `envPrefix:"REBOUND_HTTP_"` + Address string `env:"REBOUND_ADDRESS"` + StatsFile string `env:"REBOUND_STATS_FILE"` + StatsFileSaveInterval time.Duration `env:"REBOUND_STATS_FILE_SAVE_INTERVAL"` + Logger func(message string, args ...any) + SSH *ssh.Options `envPrefix:"REBOUND_SSH_"` + HTTP *http.Options `envPrefix:"REBOUND_HTTP_"` } func (o *Options) ParseEnv() error { @@ -28,10 +31,12 @@ type OptionFunc func(*Options) func DefaultOptions() *Options { return &Options{ - Address: "127.0.0.1:2222", - Logger: log.Printf, - SSH: ssh.DefaultOptions(), - HTTP: http.DefaultOptions(), + Address: "127.0.0.1:2222", + StatsFile: "stats.json", + StatsFileSaveInterval: 30 * time.Second, + Logger: log.Printf, + SSH: ssh.DefaultOptions(), + HTTP: http.DefaultOptions(), } } @@ -49,6 +54,18 @@ func WithLogger(logger func(message string, args ...any)) func(*Options) { } } +func WithStatsFile(path string) func(*Options) { + return func(o *Options) { + o.StatsFile = path + } +} + +func WithStatsFileSaveInterval(interval time.Duration) func(*Options) { + return func(o *Options) { + o.StatsFileSaveInterval = interval + } +} + func WithSSHOption(funcs ...ssh.OptionFunc) func(*Options) { return func(o *Options) { for _, fn := range funcs { diff --git a/server.go b/server.go index f522bc3..6046caa 100644 --- a/server.go +++ b/server.go @@ -1,21 +1,34 @@ package rebound import ( + "log/slog" "net" + "os" + "time" "forge.cadoles.com/wpetit/rebound/http" "forge.cadoles.com/wpetit/rebound/ssh" + "forge.cadoles.com/wpetit/rebound/stat" "github.com/pkg/errors" ) type Server struct { listener net.Listener opts *Options + stats *stat.Store } func (s *Server) Start() error { s.log("[INFO] listening on %s", s.opts.Address) + if err := s.stats.Load(s.opts.StatsFile); err != nil { + if errors.Is(err, os.ErrNotExist) { + s.log("[INFO] stats file does not exist. ignoring.") + } else { + return errors.WithStack(err) + } + } + listener, err := net.Listen("tcp", s.opts.Address) if err != nil { return errors.WithStack(err) @@ -34,6 +47,7 @@ func (s *Server) Start() error { ssh.WithPublicPort(s.opts.SSH.PublicPort), ssh.WithSockDir(s.opts.SSH.SockDir), ssh.WithLogger(s.opts.SSH.Logger), + ssh.WithStats(s.stats), ) if err := server.Serve(sshListener); err != nil { @@ -49,6 +63,7 @@ func (s *Server) Start() error { http.WithCustomDir(s.opts.HTTP.CustomDir), http.WithTemplateData(s.opts.HTTP.TemplateData), http.WithLogger(s.opts.HTTP.Logger), + http.WithStats(s.stats), ) if err := server.Serve(httpListener); err != nil { @@ -57,6 +72,23 @@ func (s *Server) Start() error { } }() + go func() { + defer listener.Close() + + ticker := time.NewTicker(s.opts.StatsFileSaveInterval) + + for { + <-ticker.C + + slog.Info("saving stats", slog.String("file", s.opts.StatsFile), slog.Duration("interval", s.opts.StatsFileSaveInterval)) + if err := s.stats.Save(s.opts.StatsFile); err != nil { + slog.Error("could not save stat file", slog.Any("error", errors.WithStack(err))) + + return + } + } + }() + return nil } @@ -85,6 +117,7 @@ func NewServer(funcs ...OptionFunc) *Server { } return &Server{ - opts: opts, + opts: opts, + stats: stat.NewStore(), } } diff --git a/ssh/direct_tcp_handler.go b/ssh/direct_tcp_handler.go index d59a2b9..3d3bed1 100644 --- a/ssh/direct_tcp_handler.go +++ b/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)) } }() diff --git a/ssh/options.go b/ssh/options.go index dd9856d..c33137e 100644 --- a/ssh/options.go +++ b/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 + } +} diff --git a/ssh/request_handler.go b/ssh/request_handler.go index 81d9dfb..eb00428 100644 --- a/ssh/request_handler.go +++ b/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() diff --git a/ssh/stats.go b/ssh/stats.go new file mode 100644 index 0000000..3e753b3 --- /dev/null +++ b/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{} diff --git a/stat/store.go b/stat/store.go new file mode 100644 index 0000000..27fd04b --- /dev/null +++ b/stat/store.go @@ -0,0 +1,154 @@ +package stat + +import ( + "encoding/json" + "os" + "path/filepath" + "sync" + + "log/slog" + + "github.com/pkg/errors" +) + +type Store struct { + data sync.Map + loadSaveLock sync.Mutex +} + +func (s *Store) Load(path string) error { + s.loadSaveLock.Lock() + defer s.loadSaveLock.Unlock() + + file, err := os.OpenFile(path, os.O_RDONLY, os.ModePerm) + if err != nil { + return errors.WithStack(err) + } + + decoder := json.NewDecoder(file) + data := map[string]any{} + + if err := decoder.Decode(&data); err != nil { + return errors.WithStack(err) + } + + s.data.Range(func(key, value any) bool { + s.data.Delete(key) + return true + }) + + for k, v := range data { + s.data.Store(k, v) + } + + return nil +} + +func (s *Store) Save(path string) error { + s.loadSaveLock.Lock() + defer s.loadSaveLock.Unlock() + + data, err := s.Snapshot() + if err != nil { + return errors.WithStack(err) + } + + dir := filepath.Dir(path) + filename := filepath.Base(path) + + temp, err := os.CreateTemp(dir, filename+".new*") + if err != nil { + return errors.WithStack(err) + } + + defer func() { + if err := os.Remove(temp.Name()); err != nil && !errors.Is(err, os.ErrNotExist) { + slog.Error("could not remove temporary file", + slog.String("file", temp.Name()), + slog.Any("error", errors.WithStack(err)), + ) + } + }() + + encoder := json.NewEncoder(temp) + if err := encoder.Encode(data); err != nil { + return errors.WithStack(err) + } + + if err := os.Rename(temp.Name(), path); err != nil { + return errors.WithStack(err) + } + + return nil +} + +func (s *Store) Snapshot() (map[string]float64, error) { + data := map[string]float64{} + + var err error + s.data.Range(func(rawKey, rawValue any) bool { + key, ok := rawKey.(string) + if !ok { + err = errors.Errorf("unexpected stat key of '%v'", rawKey) + + return false + } + + value, ok := rawValue.(float64) + if !ok { + err = errors.Errorf("unexpected stat value of '%v'", rawValue) + + return false + } + + data[key] = value + + return true + }) + if err != nil { + return nil, errors.WithStack(err) + } + + return data, nil +} + +func (s *Store) Add(name string, added float64, defaultValue float64) float64 { + for { + value := s.Get(name, defaultValue) + if value == defaultValue { + s.data.Store(name, defaultValue) + } + + sum := value + added + if s.data.CompareAndSwap(name, value, value+added) { + return sum + } + } +} + +func (s *Store) Set(name string, value float64) float64 { + s.data.Store(name, value) + + return value +} + +func (s *Store) Get(name string, defaultValue float64) float64 { + rawValue, ok := s.data.Load(name) + if !ok { + return defaultValue + } + + value, ok := rawValue.(float64) + if !ok { + return defaultValue + } + + return value +} + +func NewStore() *Store { + return &Store{ + data: sync.Map{}, + loadSaveLock: sync.Mutex{}, + } +}