repos / pico

tunnel.go

 1package shared
 2
 3import (
 4	"fmt"
 5
 6	"github.com/charmbracelet/ssh"
 7	"github.com/picosh/pico/db"
 8)
 9
10type ctxPublicKey struct{}
11type ctxUserKey struct{}
12
13func GetPublicKeyCtx(ctx ssh.Context) (ssh.PublicKey, error) {
14	pk, ok := ctx.Value(ctxPublicKey{}).(ssh.PublicKey)
15	if !ok {
16		return nil, fmt.Errorf("public key not set on `ssh.Context()` for connection")
17	}
18	return pk, nil
19}
20
21func SetPublicKeyCtx(ctx ssh.Context, pk ssh.PublicKey) {
22	ctx.SetValue(ctxPublicKey{}, pk)
23}
24
25func GetUserCtx(ctx ssh.Context) (*db.User, error) {
26	pk, ok := ctx.Value(ctxUserKey{}).(*db.User)
27	if !ok {
28		return nil, fmt.Errorf("user not set on `ssh.Context()` for connection")
29	}
30	return pk, nil
31}
32
33func SetUserCtx(ctx ssh.Context, user *db.User) {
34	ctx.SetValue(ctxUserKey{}, user)
35}