main pico / pkg / shared / ssh_test.go
Antonio Mika  ·  2026-10-25
 1package shared
 2
 3import (
 4	"crypto/ed25519"
 5	"crypto/rand"
 6	"io"
 7	"log/slog"
 8	"testing"
 9
10	"github.com/picosh/pico/pkg/db"
11	"golang.org/x/crypto/ssh"
12)
13
14type fakeAuthDB struct {
15	user *db.User
16	logs []*db.AccessLog
17}
18
19func (f *fakeAuthDB) FindUser(string) (*db.User, error)                   { return f.user, nil }
20func (f *fakeAuthDB) FindUserByPubkey(string) (*db.User, error)           { return f.user, nil }
21func (f *fakeAuthDB) FindUserByName(string) (*db.User, error)             { return f.user, nil }
22func (f *fakeAuthDB) FindFeature(string, string) (*db.FeatureFlag, error) { return nil, nil }
23func (f *fakeAuthDB) InsertAccessLog(log *db.AccessLog) error {
24	f.logs = append(f.logs, log)
25	return nil
26}
27
28type fakeConn struct{ ssh.ConnMetadata }
29
30func (fakeConn) User() string { return "erock" }
31
32func TestPubkeyAuthHandlerUsesKeyName(t *testing.T) {
33	pub, _, err := ed25519.GenerateKey(rand.Reader)
34	if err != nil {
35		t.Fatal(err)
36	}
37	key, err := ssh.NewPublicKey(pub)
38	if err != nil {
39		t.Fatal(err)
40	}
41
42	fdb := &fakeAuthDB{user: &db.User{ID: "1", Name: "erock", PublicKey: &db.PublicKey{Name: "laptop"}}}
43	handler := NewSshAuthHandler(fdb, slog.New(slog.NewTextHandler(io.Discard, nil)), "pgs")
44
45	perms, err := handler.PubkeyAuthHandler(fakeConn{}, key)
46	if err != nil {
47		t.Fatal(err)
48	}
49	if got := perms.Extensions["identity"]; got != "laptop" {
50		t.Fatalf("identity = %q, want laptop", got)
51	}
52	if len(fdb.logs) != 1 || fdb.logs[0].Identity != "laptop" {
53		t.Fatalf("access log identity = %+v, want laptop", fdb.logs)
54	}
55}