Antonio Mika
·
2026-10-11
1package pgs
2
3import (
4 "context"
5 "crypto/ed25519"
6 "crypto/rand"
7 "encoding/pem"
8 "fmt"
9 "log/slog"
10 "net"
11 "os"
12 "os/exec"
13 "path/filepath"
14 "testing"
15 "time"
16
17 pgsdb "github.com/picosh/pico/pkg/apps/pgs/db"
18 "github.com/picosh/pico/pkg/db"
19 "github.com/picosh/pico/pkg/pssh"
20 "github.com/picosh/pico/pkg/shared"
21 "github.com/picosh/pico/pkg/storage"
22 "github.com/pkg/sftp"
23 "golang.org/x/crypto/ssh"
24)
25
26func StartSshServerForTesting(cfg *Config, killCh chan error, readyCh chan *pssh.SSHServer) {
27 ctx, cancel := context.WithCancel(context.Background())
28 defer func() {
29 // Cancel is deferred to avoid being called in error path
30 // It will be called from the cleanup goroutine
31 _ = cancel
32 }()
33
34 cacheClearingQueue := make(chan string, 100)
35 logger := cfg.Logger
36
37 server, err := createSshServer(cfg, ctx, cacheClearingQueue)
38 if err != nil {
39 logger.Error("failed to create ssh server", "err", err.Error())
40 cancel() // Clean up if server creation fails
41 readyCh <- nil
42 return
43 }
44 // No test reads metrics, and a fixed port collides between test servers.
45 server.Config.PromListenAddr = ""
46
47 logger.Info("Starting SSH server", "addr", server.Config.ListenAddr)
48
49 // Signal that server is ready once ListenAndServe starts
50 go func() {
51 if err = server.ListenAndServe(); err != nil {
52 logger.Error("serve", "err", err.Error())
53 }
54 }()
55
56 // Send server when listener is created (happens early in ListenAndServe)
57 readyCh <- server
58
59 go func() {
60 // Wait for kill signal and clean up
61 <-killCh
62 logger.Info("stopping ssh server")
63 cancel()
64 }()
65}
66
67func TestSshServerSftp(t *testing.T) {
68 opts := &slog.HandlerOptions{
69 AddSource: true,
70 Level: slog.LevelDebug,
71 }
72 logger := slog.New(
73 slog.NewTextHandler(os.Stdout, opts),
74 )
75 slog.SetDefault(logger)
76 dbpool := pgsdb.NewDBMemory(logger)
77 // setup test data
78 dbpool.SetupTestData()
79 st, err := storage.NewStorageMemory(map[string]map[string]string{})
80 if err != nil {
81 panic(err)
82 }
83 pubsub := NewPubsubChan()
84 defer func() {
85 _ = pubsub.Close()
86 }()
87
88 // Use dynamic port for tests to avoid port conflicts
89 _ = os.Setenv("PGS_SSH_PORT", "0")
90
91 cfg := NewConfig(logger, dbpool, st, pubsub)
92 done := make(chan error)
93 readyCh := make(chan *pssh.SSHServer)
94
95 go StartSshServerForTesting(cfg, done, readyCh)
96
97 // Wait for server to be ready
98 server := <-readyCh
99 if server == nil {
100 t.Fatal("failed to create ssh server")
101 }
102
103 // Wait for listener to be created
104 var actualAddr string
105 for i := 0; i < 100; i++ {
106 server.Mu.Lock()
107 listener := server.Listener
108 server.Mu.Unlock()
109 if listener != nil {
110 actualAddr = listener.Addr().String()
111 break
112 }
113 time.Sleep(10 * time.Millisecond)
114 }
115
116 if actualAddr == "" {
117 t.Fatal("server listener not ready")
118 }
119
120 user := GenerateUser()
121 // add user's pubkey to the default test account
122 dbpool.Pubkeys = append(dbpool.Pubkeys, &db.PublicKey{
123 ID: "nice-pubkey",
124 UserID: dbpool.Users[0].ID,
125 Key: shared.KeyForKeyText(user.signer.PublicKey()),
126 })
127
128 client, err := user.NewClientAddr(actualAddr)
129 if err != nil {
130 t.Error(err)
131 return
132 }
133 defer func() {
134 _ = client.Close()
135 }()
136
137 _, err = WriteFileWithSftp(cfg, client)
138 if err != nil {
139 t.Error(err)
140 return
141 }
142
143 _, err = WriteFilesMultProjectsWithSftp(cfg, client)
144 if err != nil {
145 t.Error(err)
146 return
147 }
148
149 projects, err := dbpool.FindProjectsByUser(dbpool.Users[0].ID)
150 if err != nil {
151 t.Error(err)
152 return
153 }
154
155 names := ""
156 for _, proj := range projects {
157 names += "_" + proj.Name
158 }
159
160 if names != "_test_mult_mult2" {
161 t.Errorf("not all projects created: %s", names)
162 return
163 }
164
165 close(done)
166 time.Sleep(100 * time.Millisecond)
167}
168
169func TestSshServerRsync(t *testing.T) {
170 opts := &slog.HandlerOptions{
171 AddSource: true,
172 Level: slog.LevelDebug,
173 }
174 logger := slog.New(
175 slog.NewTextHandler(os.Stdout, opts),
176 )
177 slog.SetDefault(logger)
178 dbpool := pgsdb.NewDBMemory(logger)
179 // setup test data
180 dbpool.SetupTestData()
181 st, err := storage.NewStorageMemory(map[string]map[string]string{})
182 if err != nil {
183 panic(err)
184 }
185 pubsub := NewPubsubChan()
186 defer func() {
187 _ = pubsub.Close()
188 }()
189
190 // Use dynamic port for tests to avoid port conflicts
191 _ = os.Setenv("PGS_SSH_PORT", "0")
192
193 cfg := NewConfig(logger, dbpool, st, pubsub)
194 done := make(chan error)
195 readyCh := make(chan *pssh.SSHServer)
196
197 go StartSshServerForTesting(cfg, done, readyCh)
198
199 // Wait for server to be ready
200 server := <-readyCh
201 if server == nil {
202 t.Fatal("failed to create ssh server")
203 }
204
205 // Wait for listener to be created
206 var actualAddr string
207 for i := 0; i < 100; i++ {
208 server.Mu.Lock()
209 listener := server.Listener
210 server.Mu.Unlock()
211 if listener != nil {
212 actualAddr = listener.Addr().String()
213 break
214 }
215 time.Sleep(10 * time.Millisecond)
216 }
217
218 if actualAddr == "" {
219 t.Fatal("server listener not ready")
220 }
221
222 user := GenerateUser()
223 key := shared.KeyForKeyText(user.signer.PublicKey())
224 // add user's pubkey to the default test account
225 dbpool.Pubkeys = append(dbpool.Pubkeys, &db.PublicKey{
226 ID: "nice-pubkey",
227 UserID: dbpool.Users[0].ID,
228 Key: key,
229 })
230
231 conn, err := user.NewClientAddr(actualAddr)
232 if err != nil {
233 t.Error(err)
234 return
235 }
236 defer func() {
237 _ = conn.Close()
238 }()
239
240 // open an SFTP session over an existing ssh connection.
241 client, err := sftp.NewClient(conn)
242 if err != nil {
243 cfg.Logger.Error("could not create sftp client", "err", err)
244 panic(err)
245 }
246 defer func() {
247 _ = client.Close()
248 }()
249
250 name, err := os.MkdirTemp("", "rsync-")
251 if err != nil {
252 panic(err)
253 }
254
255 // remove the temporary directory at the end of the program
256 // defer os.RemoveAll(name)
257
258 block := &pem.Block{
259 Type: "OPENSSH PRIVATE KEY",
260 Bytes: user.privateKey,
261 }
262 keyFile := filepath.Join(name, "id_ed25519")
263 err = os.WriteFile(
264 keyFile,
265 pem.EncodeToMemory(block), 0600,
266 )
267 if err != nil {
268 t.Fatal(err)
269 }
270
271 index := "<!doctype html><html><body>index</body></html>"
272 err = os.WriteFile(
273 filepath.Join(name, "index.html"),
274 []byte(index), 0666,
275 )
276 if err != nil {
277 t.Fatal(err)
278 }
279
280 about := "<!doctype html><html><body>about</body></html>"
281 aboutFile := filepath.Join(name, "about.html")
282 err = os.WriteFile(
283 aboutFile,
284 []byte(about), 0666,
285 )
286 if err != nil {
287 t.Fatal(err)
288 }
289
290 contact := "<!doctype html><html><body>contact</body></html>"
291 err = os.WriteFile(
292 filepath.Join(name, "contact.html"),
293 []byte(contact), 0666,
294 )
295 if err != nil {
296 t.Fatal(err)
297 }
298
299 // Extract port from actualAddr (format: "0.0.0.0:XXXXX")
300 _, port, err := net.SplitHostPort(actualAddr)
301 if err != nil {
302 t.Fatalf("failed to parse server address: %v", err)
303 }
304 // Use localhost for rsync (works regardless of IPv4/IPv6 binding)
305 host := "localhost"
306
307 eCmd := fmt.Sprintf(
308 "ssh -p %s -o IdentitiesOnly=yes -i %s -o StrictHostKeyChecking=no",
309 port,
310 keyFile,
311 )
312
313 // copy files
314 cmd := exec.Command("rsync", "-rv", "-e", eCmd, name+"/", host+":/test")
315 result, err := cmd.CombinedOutput()
316 if err != nil {
317 cfg.Logger.Error("cannot upload", "err", err, "result", string(result))
318 t.Error(err)
319 return
320 }
321
322 // check it's there
323 fi, err := client.Lstat("/test/about.html")
324 if err != nil {
325 cfg.Logger.Error("could not get stat for file", "err", err)
326 t.Error("about.html not found")
327 return
328 }
329 if fi.Size() != 46 {
330 cfg.Logger.Error("about.html wrong size", "size", fi.Size())
331 t.Error("about.html wrong size")
332 return
333 }
334
335 // remove about file
336 _ = os.Remove(aboutFile)
337
338 // copy files with delete
339 delCmd := exec.Command("rsync", "-rv", "--delete", "-e", eCmd, name+"/", "localhost:/test")
340 result, err = delCmd.CombinedOutput()
341 if err != nil {
342 cfg.Logger.Error("cannot upload with delete", "err", err, "result", string(result))
343 t.Error(err)
344 return
345 }
346
347 // check it's not there
348 _, err = client.Lstat("/test/about.html")
349 if err == nil {
350 cfg.Logger.Error("file still exists")
351 t.Error("about.html found")
352 return
353 }
354
355 dlName, err := os.MkdirTemp("", "rsync-download")
356 if err != nil {
357 panic(err)
358 }
359 defer func() {
360 _ = os.RemoveAll(dlName)
361 }()
362 // download files
363 downloadCmd := exec.Command("rsync", "-rvvv", "-e", eCmd, host+":/test/", dlName+"/")
364 result, err = downloadCmd.CombinedOutput()
365 if err != nil {
366 cfg.Logger.Error("cannot download files", "err", err, "result", string(result))
367 t.Error(err)
368 return
369 }
370 // check contents
371 idx, err := os.ReadFile(filepath.Join(dlName, "index.html"))
372 if err != nil {
373 cfg.Logger.Error("cannot open file", "file", "index.html", "err", err)
374 t.Error(err)
375 return
376 }
377 if string(idx) != index {
378 t.Error("downloaded index.html file does not match original")
379 return
380 }
381 _, err = os.ReadFile(filepath.Join(dlName, "about.html"))
382 if err == nil {
383 cfg.Logger.Error("about file should not exist", "file", "about.html")
384 t.Error(err)
385 return
386 }
387 cnt, err := os.ReadFile(filepath.Join(dlName, "contact.html"))
388 if err != nil {
389 cfg.Logger.Error("cannot open file", "file", "contact.html", "err", err)
390 t.Error(err)
391 return
392 }
393 if string(cnt) != contact {
394 t.Error("downloaded contact.html file does not match original")
395 return
396 }
397
398 close(done)
399 time.Sleep(100 * time.Millisecond)
400}
401
402type UserSSH struct {
403 username string
404 signer ssh.Signer
405 privateKey []byte
406}
407
408func (s UserSSH) NewClientAddr(addr string) (*ssh.Client, error) {
409 config := &ssh.ClientConfig{
410 User: s.username,
411 Auth: []ssh.AuthMethod{
412 ssh.PublicKeys(s.signer),
413 },
414 HostKeyCallback: ssh.InsecureIgnoreHostKey(),
415 }
416
417 client, err := ssh.Dial("tcp", addr, config)
418 return client, err
419}
420
421func GenerateUser() UserSSH {
422 _, userKey, err := ed25519.GenerateKey(rand.Reader)
423 if err != nil {
424 panic(err)
425 }
426
427 b, err := ssh.MarshalPrivateKey(userKey, "")
428 if err != nil {
429 panic(err)
430 }
431
432 userSigner, err := ssh.NewSignerFromKey(userKey)
433 if err != nil {
434 panic(err)
435 }
436
437 return UserSSH{
438 username: "testuser",
439 signer: userSigner,
440 privateKey: b.Bytes,
441 }
442}
443
444func WriteFileWithSftp(cfg *Config, conn *ssh.Client) (*os.FileInfo, error) {
445 // open an SFTP session over an existing ssh connection.
446 client, err := sftp.NewClient(conn)
447 if err != nil {
448 cfg.Logger.Error("could not create sftp client", "err", err)
449 return nil, err
450 }
451 defer func() {
452 _ = client.Close()
453 }()
454
455 f, err := client.Create("test/hello.txt")
456 if err != nil {
457 cfg.Logger.Error("could not create file", "err", err)
458 return nil, err
459 }
460 if _, err := f.Write([]byte("Hello world!")); err != nil {
461 cfg.Logger.Error("could not write to file", "err", err)
462 return nil, err
463 }
464
465 cfg.Logger.Info("closing", "err", f.Close())
466
467 // check it's there
468 fi, err := client.Lstat("test/hello.txt")
469 if err != nil {
470 cfg.Logger.Error("could not get stat for file", "err", err)
471 return nil, err
472 }
473
474 return &fi, nil
475}
476
477func WriteFilesMultProjectsWithSftp(cfg *Config, conn *ssh.Client) (*os.FileInfo, error) {
478 // open an SFTP session over an existing ssh connection.
479 client, err := sftp.NewClient(conn)
480 if err != nil {
481 cfg.Logger.Error("could not create sftp client", "err", err)
482 return nil, err
483 }
484 defer func() {
485 _ = client.Close()
486 }()
487
488 f, err := client.Create("mult/hello.txt")
489 if err != nil {
490 cfg.Logger.Error("could not create file", "err", err)
491 return nil, err
492 }
493 if _, err := f.Write([]byte("Hello world!")); err != nil {
494 cfg.Logger.Error("could not write to file", "err", err)
495 return nil, err
496 }
497
498 f, err = client.Create("mult2/hello.txt")
499 if err != nil {
500 cfg.Logger.Error("could not create file", "err", err)
501 return nil, err
502 }
503 if _, err := f.Write([]byte("Hello world!")); err != nil {
504 cfg.Logger.Error("could not write to file", "err", err)
505 return nil, err
506 }
507
508 cfg.Logger.Info("closing", "err", f.Close())
509
510 // check it's there
511 fi, err := client.Lstat("test/hello.txt")
512 if err != nil {
513 cfg.Logger.Error("could not get stat for file", "err", err)
514 return nil, err
515 }
516
517 return &fi, nil
518}