add structured logging with Loki integration

go.mod

@@ -14,14 +14,58 @@ require (
 
 require (
 	github.com/anmitsu/go-shlex v0.0.0-20200514113438-38f4b401e2be // indirect
+	github.com/beorn7/perks v1.0.1 // indirect
+	github.com/cespare/xxhash/v2 v2.3.0 // indirect
+	github.com/davecgh/go-spew v1.1.1 // indirect
+	github.com/dennwc/varint v1.0.0 // indirect
+	github.com/go-kit/kit v0.10.0 // indirect
+	github.com/go-kit/log v0.2.1 // indirect
+	github.com/go-logfmt/logfmt v0.5.1 // indirect
+	github.com/gogo/protobuf v1.3.2 // indirect
+	github.com/golang/protobuf v1.5.3 // indirect
+	github.com/golang/snappy v0.0.4 // indirect
+	github.com/grafana/loki-client-go v0.0.0-20240913122146-e119d400c3a5 // indirect
+	github.com/grafana/loki/pkg/push v0.0.0-20240912152814-63e84b476a9a // indirect
+	github.com/grafana/regexp v0.0.0-20220304095617-2e8d9baf4ac2 // indirect
 	github.com/jinzhu/inflection v1.0.0 // indirect
+	github.com/jpillora/backoff v1.0.0 // indirect
+	github.com/json-iterator/go v1.1.12 // indirect
+	github.com/kylelemons/godebug v1.1.0 // indirect
+	github.com/mitchellh/mapstructure v1.4.3 // indirect
+	github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd // indirect
+	github.com/modern-go/reflect2 v1.0.2 // indirect
+	github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822 // indirect
+	github.com/mwitkow/go-conntrack v0.0.0-20190716064945-2f068394615f // indirect
 	github.com/oschwald/maxminddb-golang v1.13.0 // indirect
+	github.com/pkg/errors v0.9.1 // indirect
+	github.com/pmezard/go-difflib v1.0.0 // indirect
+	github.com/prometheus/client_golang v1.20.4 // indirect
+	github.com/prometheus/client_model v0.6.1 // indirect
+	github.com/prometheus/common v0.63.0 // indirect
+	github.com/prometheus/procfs v0.15.1 // indirect
+	github.com/prometheus/prometheus v0.35.0 // indirect
 	github.com/puzpuzpuz/xsync/v3 v3.5.1 // indirect
+	github.com/samber/lo v1.51.0 // indirect
+	github.com/samber/slog-common v0.19.0 // indirect
+	github.com/samber/slog-loki/v3 v3.6.0 // indirect
+	github.com/samber/slog-multi v1.5.0 // indirect
+	github.com/stretchr/testify v1.10.0 // indirect
 	github.com/tmthrgd/go-hex v0.0.0-20190904060850-447a3041c3bc // indirect
 	github.com/vmihailenco/msgpack/v5 v5.4.1 // indirect
 	github.com/vmihailenco/tagparser/v2 v2.0.0 // indirect
 	go.opentelemetry.io/otel v1.37.0 // indirect
 	go.opentelemetry.io/otel/trace v1.37.0 // indirect
+	go.uber.org/atomic v1.9.0 // indirect
+	go.uber.org/goleak v1.3.0 // indirect
+	golang.org/x/exp v0.0.0-20240325151524-a685a6edb6d8 // indirect
+	golang.org/x/net v0.45.0 // indirect
+	golang.org/x/oauth2 v0.25.0 // indirect
 	golang.org/x/sys v0.37.0 // indirect
+	golang.org/x/text v0.30.0 // indirect
+	google.golang.org/genproto v0.0.0-20230410155749-daa745c078e1 // indirect
+	google.golang.org/grpc v1.56.3 // indirect
+	google.golang.org/protobuf v1.36.5 // indirect
+	gopkg.in/yaml.v2 v2.4.0 // indirect
+	gopkg.in/yaml.v3 v3.0.1 // indirect
 	mellium.im/sasl v0.3.2 // indirect
 )

inc.env

@@ -2,3 +2,4 @@ PORT=2222
 GEOIP_DB=GeoLite2-City.mmdb
 DB_DSN="postgrtesql://postgres:password@localhost/postgres"
 ROOT_PATH="./"
+LOKI_HOST=""

main.go

@@ -3,10 +3,16 @@ package main
 import (
 	"fmt"
 	"log"
+	"log/slog"
 	"net"
+	"os"
 	"slices"
 	"strings"
 
+	"github.com/grafana/loki-client-go/loki"
+	slogloki "github.com/samber/slog-loki/v3"
+	slogmulti "github.com/samber/slog-multi"
+
 	"sshchat/db"
 	"sshchat/utils"
 
@@ -17,7 +23,7 @@ import (
 
 var config = utils.GetConfig()
 
-func sessionHandler(s ssh.Session, geoip *geoip2.Reader, pgDb *bun.DB) {
+func sessionHandler(s ssh.Session, geoip *geoip2.Reader, pgDb *bun.DB, logger *slog.Logger) {
 	ptyReq, _, isPty := s.Pty()
 	if !isPty {
 		_, _ = fmt.Fprintln(s, "Err: PTY requires. Reconnect with -t option.")
@@ -35,25 +41,25 @@ func sessionHandler(s ssh.Session, geoip *geoip2.Reader, pgDb *bun.DB) {
 
 	geoStatus := utils.GetIPInfo(remote, geoip)
 	if geoStatus == nil {
-		log.Printf("[sshchat] %s connected. %s / UNK [FORCE DISCONNECT]", username, remote)
+		logger.Info("[sshchat] connected", "user", username, "remote", remote, "country", "UNK", "status", "FORCE DISCONNECT")
 		_, _ = fmt.Fprintf(s, "[system] Your access country is blacklisted. UNK")
 		_ = s.Close()
 		return
 	} else {
-		log.Printf("[sshchat] %s connected. %s / %s", username, remote, geoStatus.Country)
+		logger.Info("[sshchat] connected", "user", username, "remote", remote, "country", geoStatus.Country)
 	}
 
 	if slices.Contains(config.CountryBlacklist, geoStatus.Country) {
-		log.Printf("[sshchat] %s country blacklisted. %s", username, remote)
+		logger.Info("[sshchat] country blacklisted", "user", username, "remote", remote)
 		_, _ = fmt.Fprintf(s, "[system] Your access country is blacklisted. %s\n", geoStatus.Country)
 		_ = s.Close()
 	}
 
 	if geoStatus.Country == "ZZ" {
 		if strings.HasPrefix(remote, "127") || strings.HasPrefix(remote, "::1") {
-			log.Printf("[sshchat] %s is localhost whitelisted.", username)
+			logger.Info("[sshchat] localhost whitelisted", "user", username)
 		} else {
-			log.Printf("[sshchat] unknown country blacklisted. %s", username)
+			logger.Info("[sshchat] unknown country blacklisted", "user", username)
 			_, _ = fmt.Fprintf(s, "[system] Unknown country is blacklisted. %s\n", geoStatus.Country)
 			_ = s.Close()
 		}
@@ -63,16 +69,52 @@ func sessionHandler(s ssh.Session, geoip *geoip2.Reader, pgDb *bun.DB) {
 
 	defer func() {
 		client.Close()
-		log.Printf("[sshchat] %s disconnected. %s / %s", username, remote, geoStatus.Country)
+		logger.Info("[sshchat] disconnected", "user", username, "remote", remote, "country", geoStatus.Country)
 	}()
 
 	client.EventLoop()
 }
 
+func getLogger(lokiHost string) (*slog.Logger, error) {
+	if lokiHost == "" {
+		logger := slog.New(slog.NewTextHandler(os.Stdout, nil))
+		logger.Info("Loki host is not set. Logging to stdout")
+
+		return logger, nil
+	}
+
+	config, _ := loki.NewDefaultConfig(lokiHost)
+	config.TenantID = "sshchat"
+	client, err := loki.New(config)
+	if err != nil {
+		slog.Error("Failed to create Loki client", "error", err)
+		return nil, err
+	}
+	defer client.Stop()
+
+	logger := slog.New(
+		slogmulti.Fanout(
+			slog.NewTextHandler(os.Stdout, nil),
+			slogloki.Option{Level: slog.LevelDebug, Client: client}.NewLokiHandler(),
+		),
+	)
+	logger = logger.With(slog.String("app", "sshchat"))
+	logger.Info("Logging to Loki", "host", lokiHost)
+
+	return logger, nil
+}
+
 func main() {
+	logger, err := getLogger(config.LokiHost)
+	if err != nil {
+		logger.Error("Failed to create logger", "error", err)
+		return
+	}
+
 	geoip, err := utils.GetDB(config.RootPath + "/" + config.Geoip)
 	if err != nil {
-		log.Fatalf("Geoip db is error: %v", err)
+		logger.Error("Geoip db is error", "error", err)
+		return
 	}
 
 	pgDb, err := db.GetDB(config.PgDsn)
@@ -84,10 +126,11 @@ func main() {
 
 	keys, err := utils.CheckHostKey(config.RootPath)
 	if err != nil {
-		log.Print("Failed to check SSH keys: generate one.\n", err)
+		logger.Error("Failed to check SSH keys: generate one", "error", err)
 		err = utils.GenerateHostKey(config.RootPath)
 		if err != nil {
-			log.Fatal(err)
+			logger.Error("Fatal error", "error", err)
+			return
 		}
 
 		keys, err = utils.CheckHostKey(config.RootPath)
@@ -99,7 +142,7 @@ func main() {
 	s := &ssh.Server{
 		Addr: ":" + port,
 		Handler: func(s ssh.Session) {
-			sessionHandler(s, geoip, pgDb)
+			sessionHandler(s, geoip, pgDb, logger)
 		},
 	}
 	for _, key := range keys {
@@ -110,6 +153,8 @@ func main() {
 		_ = pgDb.Close()
 	}()
 
-	log.Print("Listening on :" + port)
+	logger.Info("Starting server", "port", port)
-	log.Fatal(s.ListenAndServe())
+	if err := s.ListenAndServe(); err != nil {
+		logger.Error("Server failed", "error", err)
+	}
 }

utils/config.go

@@ -13,6 +13,7 @@ type Config struct {
 	CountryBlacklist []string
 	PgDsn            string
 	RootPath         string
+	LokiHost         string
 }
 
 func GetConfig() *Config {
@@ -23,6 +24,7 @@ func GetConfig() *Config {
 	countryBlacklist := os.Getenv("COUNTRY_BLACKLIST")
 	pgDsn := os.Getenv("DB_DSN")
 	rootPath := os.Getenv("ROOT_PATH")
+	lokiHost := os.Getenv("LOKI_HOST")
 
 	return &Config{
 		Port:             port,
@@ -30,5 +32,6 @@ func GetConfig() *Config {
 		CountryBlacklist: strings.Split(countryBlacklist, ","),
 		PgDsn:            pgDsn,
 		RootPath:         rootPath,
+		LokiHost:         lokiHost,
 	}
 }