repos / pico

pico services mono repo
git clone https://github.com/picosh/pico.git

pico / pkg / pssh
Eric Bower  ·  2025-12-17

logger.go

  1package pssh
  2
  3import (
  4	"log/slog"
  5	"time"
  6
  7	"github.com/picosh/pico/pkg/db"
  8)
  9
 10type ctxLoggerKey struct{}
 11type ctxUserKey struct{}
 12
 13type FindUserInterface interface {
 14	FindUser(string) (*db.User, error)
 15	FindUserByPubkey(string) (*db.User, error)
 16}
 17
 18type GetLoggerInterface interface {
 19	GetLogger(s *SSHServerConnSession) *slog.Logger
 20}
 21
 22func LogMiddleware(getLogger GetLoggerInterface, database FindUserInterface) SSHServerMiddleware {
 23	return func(sshHandler SSHServerHandler) SSHServerHandler {
 24		return func(s *SSHServerConnSession) error {
 25			ct := time.Now()
 26
 27			logger := GetLogger(s)
 28			if logger == slog.Default() || logger == s.Logger {
 29				logger = getLogger.GetLogger(s)
 30
 31				user := GetUser(s)
 32				if user == nil {
 33					_, impersonated := s.Permissions().Extensions["imp_id"]
 34
 35					var user *db.User
 36					var err error
 37					var found bool
 38
 39					if !impersonated {
 40						pubKey, ok := s.Permissions().Extensions["pubkey"]
 41						if ok {
 42							user, err = database.FindUserByPubkey(pubKey)
 43							found = true
 44						}
 45					} else {
 46						userID, ok := s.Permissions().Extensions["user_id"]
 47						if ok {
 48							user, err = database.FindUser(userID)
 49							found = true
 50						}
 51					}
 52
 53					if found {
 54						// identity provided by ssh-cert
 55						identity := s.Permissions().Extensions["identity"]
 56						if err == nil && user != nil {
 57							logger = logger.With(
 58								"user", user.Name,
 59								"userId", user.ID,
 60								"ip", s.RemoteAddr().String(),
 61								"identity", identity,
 62							)
 63
 64							SetUser(s, user)
 65						} else {
 66							logger.Error("`user` not set in permissions", "err", err)
 67						}
 68					}
 69				}
 70
 71				SetLogger(s, logger)
 72			}
 73
 74			pty, _, ok := s.Pty()
 75
 76			width, height := 0, 0
 77			term := ""
 78			if pty != nil {
 79				term = pty.Term
 80				width = pty.Window.Width
 81				height = pty.Window.Height
 82			}
 83
 84			logger.Info(
 85				"connect",
 86				"sshUser", s.User(),
 87				"pty", ok,
 88				"term", term,
 89				"windowWidth", width,
 90				"windowHeight", height,
 91			)
 92
 93			err := sshHandler(s)
 94			if err != nil {
 95				logger.Error("error", "err", err)
 96			}
 97
 98			if pty != nil {
 99				term = pty.Term
100				width = pty.Window.Width
101				height = pty.Window.Height
102			}
103
104			logger.Info(
105				"disconnect",
106				"sshUser", s.User(),
107				"pty", ok,
108				"term", term,
109				"windowWidth", width,
110				"windowHeight", height,
111				"duration", time.Since(ct),
112				"err", err,
113			)
114
115			return err
116		}
117	}
118}
119
120func GetLogger(s *SSHServerConnSession) *slog.Logger {
121	logger := slog.Default()
122	if s == nil {
123		return logger
124	}
125
126	logger = s.Logger
127
128	if v, ok := s.Context().Value(ctxLoggerKey{}).(*slog.Logger); ok {
129		return v
130	}
131
132	return logger
133}
134
135func SetLogger(s *SSHServerConnSession, logger *slog.Logger) {
136	if s == nil {
137		return
138	}
139
140	s.SetValue(ctxLoggerKey{}, logger)
141}
142
143func GetUser(s *SSHServerConnSession) *db.User {
144	if v, ok := s.Context().Value(ctxUserKey{}).(*db.User); ok {
145		return v
146	}
147
148	return nil
149}
150
151func SetUser(s *SSHServerConnSession, user *db.User) {
152	if s == nil {
153		return
154	}
155
156	s.SetValue(ctxUserKey{}, user)
157}