repos / pico

ssh.go

  1package shared
  2
  3import (
  4	"fmt"
  5	"log/slog"
  6
  7	"github.com/charmbracelet/ssh"
  8	"github.com/picosh/pico/db"
  9	"github.com/picosh/utils"
 10)
 11
 12type ctxUserKey struct{}
 13type ctxFeatureFlagKey struct{}
 14
 15func GetUser(ctx ssh.Context) (*db.User, error) {
 16	user, ok := ctx.Value(ctxUserKey{}).(*db.User)
 17	if !ok {
 18		return user, fmt.Errorf("user not set on `ssh.Context()` for connection")
 19	}
 20	return user, nil
 21}
 22
 23func SetUser(ctx ssh.Context, user *db.User) {
 24	ctx.SetValue(ctxUserKey{}, user)
 25}
 26
 27func GetFeatureFlag(ctx ssh.Context) (*db.FeatureFlag, error) {
 28	ff, ok := ctx.Value(ctxFeatureFlagKey{}).(*db.FeatureFlag)
 29	if !ok || ff.Name == "" {
 30		return ff, fmt.Errorf("feature flag not set on `ssh.Context()` for connection")
 31	}
 32	return ff, nil
 33}
 34
 35func SetFeatureFlag(ctx ssh.Context, ff *db.FeatureFlag) {
 36	ctx.SetValue(ctxFeatureFlagKey{}, ff)
 37}
 38
 39type ctxPublicKey struct{}
 40
 41func GetPublicKey(ctx ssh.Context) (ssh.PublicKey, error) {
 42	pk, ok := ctx.Value(ctxPublicKey{}).(ssh.PublicKey)
 43	if !ok {
 44		return nil, fmt.Errorf("public key not set on `ssh.Context()` for connection")
 45	}
 46	return pk, nil
 47}
 48
 49func SetPublicKey(ctx ssh.Context, pk ssh.PublicKey) {
 50	ctx.SetValue(ctxPublicKey{}, pk)
 51}
 52
 53type SshAuthHandler struct {
 54	DBPool db.DB
 55	Logger *slog.Logger
 56	Cfg    *ConfigSite
 57}
 58
 59func NewSshAuthHandler(dbpool db.DB, logger *slog.Logger, cfg *ConfigSite) *SshAuthHandler {
 60	return &SshAuthHandler{
 61		DBPool: dbpool,
 62		Logger: logger,
 63		Cfg:    cfg,
 64	}
 65}
 66
 67func (r *SshAuthHandler) PubkeyAuthHandler(ctx ssh.Context, key ssh.PublicKey) bool {
 68	SetPublicKey(ctx, key)
 69
 70	pubkey := utils.KeyForKeyText(key)
 71
 72	user, err := r.DBPool.FindUserForKey(ctx.User(), pubkey)
 73	if err != nil {
 74		r.Logger.Error(
 75			"could not find user for key",
 76			"key", key,
 77			"err", err,
 78		)
 79		return false
 80	}
 81
 82	if user.Name == "" {
 83		r.Logger.Error("username is not set")
 84		return false
 85	}
 86
 87	ff, _ := r.DBPool.FindFeatureForUser(user.ID, "plus")
 88	// we have free tiers so users might not have a feature flag
 89	// in which case we set sane defaults
 90	if ff == nil {
 91		ff = db.NewFeatureFlag(
 92			user.ID,
 93			"plus",
 94			r.Cfg.MaxSize,
 95			r.Cfg.MaxAssetSize,
 96			r.Cfg.MaxSpecialFileSize,
 97		)
 98	}
 99	// this is jank
100	ff.Data.StorageMax = ff.FindStorageMax(r.Cfg.MaxSize)
101	ff.Data.FileMax = ff.FindFileMax(r.Cfg.MaxAssetSize)
102	ff.Data.SpecialFileMax = ff.FindSpecialFileMax(r.Cfg.MaxSpecialFileSize)
103
104	SetUser(ctx, user)
105	SetFeatureFlag(ctx, ff)
106	return true
107}