Commit 5cc5eff559

5cc5eff55957b67141840311e3ad83c9d9056665

parent: 33045f8a24

Verified · cmc

cmc <hello@cleberg.net> · 2026-09-23 06:54 UTC

sshd: Stop ends open follows; test that a closed channel does

Ref #251
internal/sshd/sshd.go +29 −3
@@ -39,6 +39,8 @@ type Server struct {
3939 sessions sync.WaitGroup // accepted connections still being served
4040 mu sync.Mutex
4141 conns map[*conn]struct{}
42 stopping chan struct{} // closed by Stop
43 stopOnce sync.Once
4244}
4345
4446// conn is one accepted connection and how many sessions it is running.
@@ -51,7 +53,7 @@ type conn struct {
5153}
5254
5355func New(cfg config.Config, st *store.Store) (*Server, error) {
54 s := &Server{cfg: cfg, st: st, authLimiter: newRateLimiter(cfg.Limits.SSHAuthRate, time.Minute), conns: map[*conn]struct{}{}}
56 s := &Server{cfg: cfg, st: st, authLimiter: newRateLimiter(cfg.Limits.SSHAuthRate, time.Minute), conns: map[*conn]struct{}{}, stopping: make(chan struct{})}
5557
5658 sc := &ssh.ServerConfig{
5759 PublicKeyCallback: s.authenticate,
@@ -177,10 +179,18 @@ func (s *Server) Serve(ln net.Listener) error {
177179 }
178180}
179181
182// Stop ends the commands that run until something happens (build log
183// --follow), so a shutdown drain waits only for work that finishes. It
184// does not close connections; Shutdown does.
185func (s *Server) Stop() {
186 s.stopOnce.Do(func() { close(s.stopping) })
187}
188
180189// Shutdown closes every idle connection, then waits for the ones with a
181190// session running, or for ctx. The caller closes the listener first; a
182191// push in flight completes rather than being cut mid-pack.
183192func (s *Server) Shutdown(ctx context.Context) error {
193 s.Stop()
184194 s.mu.Lock()
185195 for c := range s.conns {
186196 if c.active.Load() == 0 {
@@ -240,15 +250,31 @@ func (s *Server) handleSession(sconn *ssh.ServerConn, ch ssh.Channel, reqs <-cha
240250 req.Reply(true, nil)
241251 // x/crypto closes reqs when the client closes the channel. That
242252 // is how a follow learns nobody is reading: the CLI's shared
243 // connection outlives a Ctrl-C, the channel does not.
244 done := make(chan struct{})
253 // connection outlives a Ctrl-C, the channel does not. Stop
254 // ends it too, for a restart.
255 closed := make(chan struct{})
245256 go func() {
246257 for r := range reqs {
247258 r.Reply(false, nil)
248259 }
260 close(closed)
261 }()
262 done := make(chan struct{})
263 go func() {
264 select {
265 case <-closed:
266 case <-s.stopping:
267 }
249268 close(done)
250269 }()
251270 code := s.runExec(sconn, ch, payload.Command, done)
271 select {
272 case <-s.stopping:
273 if code != protocol.ExitOK {
274 fmt.Fprintln(ch.Stderr(), "gitbay is restarting; run the command again in a moment")
275 }
276 default:
277 }
252278 sendExit(ch, code)
253279 return
254280 case "shell":
internal/sshd/sshd_test.go added +178
@@ -0,0 +1,178 @@
1package sshd
2
3import (
4 "bufio"
5 "bytes"
6 "crypto/ed25519"
7 "crypto/rand"
8 "errors"
9 "net"
10 "path/filepath"
11 "strings"
12 "testing"
13 "time"
14
15 "golang.org/x/crypto/ssh"
16
17 "gitbay.org/gitbay/internal/config"
18 "gitbay.org/gitbay/internal/store"
19)
20
21// followServer starts an embedded server holding alice, her public repo
22// alice/app and a queued build 1 whose log has one line, and returns it
23// with a client connected as alice.
24func followServer(t *testing.T) (*Server, *ssh.Client) {
25 t.Helper()
26 root := t.TempDir()
27 st, err := store.Open(filepath.Join(root, "gitbay.db"))
28 if err != nil {
29 t.Fatal(err)
30 }
31 t.Cleanup(func() { st.Close() })
32 if err := st.MigrateUp(); err != nil {
33 t.Fatal(err)
34 }
35 uid, err := st.CreateUser("alice", false)
36 if err != nil {
37 t.Fatal(err)
38 }
39 repoID, err := st.CreateRepo("user", uid, "app", "public")
40 if err != nil {
41 t.Fatal(err)
42 }
43 id, err := st.CreateBuild(repoID, "unit", "abc", "main", `["true"]`, "", "", true)
44 if err != nil {
45 t.Fatal(err)
46 }
47 if err := st.AppendBuildLog(id, []byte("queued\n")); err != nil {
48 t.Fatal(err)
49 }
50 _, priv, err := ed25519.GenerateKey(rand.Reader)
51 if err != nil {
52 t.Fatal(err)
53 }
54 signer, err := ssh.NewSignerFromKey(priv)
55 if err != nil {
56 t.Fatal(err)
57 }
58 pub := signer.PublicKey()
59 if err := st.AddSSHKey(uid, ssh.FingerprintSHA256(pub), pub.Type(), pub.Marshal(), "full", "test"); err != nil {
60 t.Fatal(err)
61 }
62
63 cfg := config.Default()
64 cfg.Server.Root = root
65 srv, err := New(cfg, st)
66 if err != nil {
67 t.Fatal(err)
68 }
69 ln, err := net.Listen("tcp", "127.0.0.1:0")
70 if err != nil {
71 t.Fatal(err)
72 }
73 go srv.Serve(ln)
74 t.Cleanup(func() { ln.Close() })
75
76 client, err := ssh.Dial("tcp", ln.Addr().String(), &ssh.ClientConfig{
77 User: "git",
78 Auth: []ssh.AuthMethod{ssh.PublicKeys(signer)},
79 HostKeyCallback: ssh.InsecureIgnoreHostKey(),
80 Timeout: 5 * time.Second,
81 })
82 if err != nil {
83 t.Fatal(err)
84 }
85 t.Cleanup(func() { client.Close() })
86 return srv, client
87}
88
89// startFollow runs build log --follow on a new session and returns once
90// the stored line has arrived, so the follow is past its first read.
91func startFollow(t *testing.T, client *ssh.Client, stderr *bytes.Buffer) *ssh.Session {
92 t.Helper()
93 sess, err := client.NewSession()
94 if err != nil {
95 t.Fatal(err)
96 }
97 t.Cleanup(func() { sess.Close() })
98 sess.Stderr = stderr
99 out, err := sess.StdoutPipe()
100 if err != nil {
101 t.Fatal(err)
102 }
103 if err := sess.Start("build log alice/app 1 --follow"); err != nil {
104 t.Fatal(err)
105 }
106 line := make(chan string, 1)
107 go func() {
108 l, _ := bufio.NewReader(out).ReadString('\n')
109 line <- l
110 }()
111 select {
112 case l := <-line:
113 if l != "queued\n" {
114 t.Fatalf("first line %q", l)
115 }
116 case <-time.After(5 * time.Second):
117 t.Fatal("the stored log never arrived")
118 }
119 return sess
120}
121
122func activeSessions(s *Server) int32 {
123 s.mu.Lock()
124 defer s.mu.Unlock()
125 var n int32
126 for c := range s.conns {
127 n += c.active.Load()
128 }
129 return n
130}
131
132// Closing the session channel ends a follow of a queued build while the
133// connection stays up, as a Ctrl-C does over the CLI's shared connection.
134// Nothing else would end it for ten minutes.
135func TestFollowEndsWhenChannelCloses(t *testing.T) {
136 srv, client := followServer(t)
137 var stderr bytes.Buffer
138 sess := startFollow(t, client, &stderr)
139 if n := activeSessions(srv); n != 1 {
140 t.Fatalf("%d sessions active while following, want 1", n)
141 }
142 sess.Close()
143
144 deadline := time.Now().Add(5 * time.Second)
145 for activeSessions(srv) != 0 {
146 if time.Now().After(deadline) {
147 t.Fatal("the follow outlived its channel")
148 }
149 time.Sleep(20 * time.Millisecond)
150 }
151 if _, err := client.NewSession(); err != nil {
152 t.Fatalf("the connection did not survive the channel: %v", err)
153 }
154}
155
156// Stop ends a follow with exit 1 and says why, without closing the
157// connection.
158func TestStopEndsFollow(t *testing.T) {
159 srv, client := followServer(t)
160 var stderr bytes.Buffer
161 sess := startFollow(t, client, &stderr)
162 srv.Stop()
163
164 waited := make(chan error, 1)
165 go func() { waited <- sess.Wait() }()
166 select {
167 case err := <-waited:
168 var exit *ssh.ExitError
169 if !errors.As(err, &exit) || exit.ExitStatus() != 1 {
170 t.Fatalf("follow ended with %v, want exit 1", err)
171 }
172 case <-time.After(5 * time.Second):
173 t.Fatal("Stop did not end the follow")
174 }
175 if !strings.Contains(stderr.String(), "gitbay is restarting") {
176 t.Errorf("stderr %q", stderr.String())
177 }
178}