fix(pipe): end pipe-web sessions with their requests

...	...		@@ -14,18 +14,17 @@ import (
14	14		"sync"
15	15		"time"
16	16
17		-	"github.com/google/uuid"
18	17		"github.com/gorilla/websocket"
19	18		"github.com/picosh/pico/pkg/db/postgres"
20	19		"github.com/picosh/pico/pkg/shared"
21	20		"github.com/picosh/pico/pkg/shared/router"
22	21		"github.com/picosh/utils/pipe"
23	22		"github.com/prometheus/client_golang/prometheus/promhttp"
	23	+	"golang.org/x/crypto/ssh"
24	24		)
25	25
26	26		var (
27	27		cleanRegex = regexp.MustCompile(`[^0-9a-zA-Z,/]`)
28		-	sshClient *pipe.Client
29	28		upgrader = websocket.Upgrader{
30	29		CheckOrigin: func(r *http.Request) bool {
31	30		return true

...	...		@@ -48,6 +47,13 @@ func (w writeFlusher) Write(p []byte) (n int, err error) {
48	47
49	48		var _ io.Writer = writeFlusher{}
50	49
	50	+	const (
	51	+	// pongWait is how long a websocket client has to answer a ping.
	52	+	pongWait = 60 * time.Second
	53	+	pingPeriod = pongWait / 2
	54	+	writeWait = 10 * time.Second
	55	+	)
	56	+
51	57		func (web *webRouter) handleSub(pubsub bool) http.HandlerFunc {
52	58		return func(w http.ResponseWriter, r *http.Request) {
53	59		logger := web.cfg.Logger

...	...		@@ -70,22 +76,15 @@ func (web *webRouter) handleSub(pubsub bool) http.HandlerFunc {
70	76		params += fmt.Sprintf(" -a=%s", cleanList)
71	77		}
72	78
73		-	id := uuid.NewString()
74		-
75		-	p, err := sshClient.AddSession(id, fmt.Sprintf("sub %s %s", params, topic), 0, -1, -1)
	79	+	p, err := web.conn.run(fmt.Sprintf("sub %s %s", params, topic))
76	80		if err != nil {
77	81		logger.Error("sub error", "topic", topic, "info", clientInfo, "err", err.Error())
78	82		http.Error(w, "server error", http.StatusInternalServerError)
79	83		return
80	84		}
81		-
82		-	go func() {
83		-	<-r.Context().Done()
84		-	err := sshClient.RemoveSession(id)
85		-	if err != nil {
86		-	logger.Error("sub remove error", "topic", topic, "info", clientInfo, "err", err.Error())
87		-	}
88		-	}()
	85	+	defer func() { _ = p.Close() }()
	86	+	// A client that leaves while nothing is published ends the sub.
	87	+	defer context.AfterFunc(r.Context(), func() { _ = p.Close() })()
89	88
90	89		if mime := r.URL.Query().Get("mime"); mime != "" {
91	90		w.Header().Add("Content-Type", r.URL.Query().Get("mime"))

...	...		@@ -142,22 +141,14 @@ func (web *webRouter) handlePub(pubsub bool) http.HandlerFunc {
142	141		params += " -e"
143	142		}
144	143
145		-	id := uuid.NewString()
146		-
147		-	p, err := sshClient.AddSession(id, fmt.Sprintf("pub %s %s", params, topic), 0, -1, -1)
	144	+	p, err := web.conn.run(fmt.Sprintf("pub %s %s", params, topic))
148	145		if err != nil {
149	146		logger.Error("pub error", "topic", topic, "info", clientInfo, "err", err.Error())
150	147		http.Error(w, "server error", http.StatusInternalServerError)
151	148		return
152	149		}
153		-
154		-	go func() {
155		-	<-r.Context().Done()
156		-	err := sshClient.RemoveSession(id)
157		-	if err != nil {
158		-	logger.Error("pub remove error", "topic", topic, "info", clientInfo, "err", err.Error())
159		-	}
160		-	}()
	150	+	defer func() { _ = p.Close() }()
	151	+	defer context.AfterFunc(r.Context(), func() { _ = p.Close() })()
161	152
162	153		var scanErr error
163	154		scanStatus := http.StatusInternalServerError

...	...		@@ -206,30 +197,16 @@ func (web *webRouter) handlePub(pubsub bool) http.HandlerFunc {
206	197		return
207	198		}
208	199
209		-	outer:
210		-	for {
211		-	select {
212		-	case <-r.Context().Done():
213		-	break outer
214		-	default:
215		-	messageToWrite := first
216		-	if prefix != "" {
217		-	messageToWrite = append([]byte(prefix), messageToWrite...)
218		-	}
219		-
220		-	n, err := p.Write(messageToWrite)
221		-	if err != nil {
222		-	logger.Error("pub write error", "topic", topic, "info", clientInfo, "err", err.Error())
223		-	http.Error(w, "server error", http.StatusInternalServerError)
224		-	return
225		-	}
226		-
227		-	if n > 0 {
228		-	break outer
229		-	}
	200	+	messageToWrite := first
	201	+	if prefix != "" {
	202	+	messageToWrite = append([]byte(prefix), messageToWrite...)
	203	+	}
230	204
231		-	time.Sleep(10 * time.Millisecond)
232		-	}
	205	+	_, err = p.Write(messageToWrite)
	206	+	if err != nil {
	207	+	logger.Error("pub write error", "topic", topic, "info", clientInfo, "err", err.Error())
	208	+	http.Error(w, "server error", http.StatusInternalServerError)
	209	+	return
233	210		}
234	211
235	212		_, err = io.Copy(p, reader)

...	...		@@ -239,9 +216,17 @@ func (web *webRouter) handlePub(pubsub bool) http.HandlerFunc {
239	216		return
240	217		}
241	218
242		-	w.WriteHeader(http.StatusOK)
	219	+	// pub sends the message once its stdin ends, and exits when it has.
	220	+	// pipe closes its sessions without an exit status.
	221	+	_ = p.CloseWrite()
	222	+	var missing *ssh.ExitMissingError
	223	+	if err := p.Wait(); err != nil && !errors.As(err, &missing) {
	224	+	logger.Error("pub exit error", "topic", topic, "info", clientInfo, "err", err.Error())
	225	+	http.Error(w, "server error", http.StatusInternalServerError)
	226	+	return
	227	+	}
243	228
244		-	time.Sleep(10 * time.Millisecond)
	229	+	w.WriteHeader(http.StatusOK)
245	230		}
246	231		}
247	232

...	...		@@ -288,22 +273,36 @@ func (web *webRouter) handlePipe() http.HandlerFunc {
288	273
289	274		prefix := r.URL.Query().Get("prefix")
290	275
291		-	id := uuid.NewString()
292		-
293		-	p, err := sshClient.AddSession(id, fmt.Sprintf("pipe %s %s", params, topic), 0, -1, -1)
	276	+	p, err := web.conn.run(fmt.Sprintf("pipe %s %s", params, topic))
294	277		if err != nil {
295	278		logger.Error("pipe error", "topic", topic, "info", clientInfo, "err", err.Error())
296		-	http.Error(w, "server error", http.StatusInternalServerError)
	279	+	_ = c.WriteControl(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseInternalServerErr, "server error"), time.Now().Add(time.Second))
297	280		return
298	281		}
299		-
	282	+	defer func() { _ = p.Close() }()
	283	+
	284	+	// A client that stops answering pings is gone, even if its
	285	+	// connection never closed.
	286	+	_ = c.SetReadDeadline(time.Now().Add(pongWait))
	287	+	c.SetPongHandler(func(string) error {
	288	+	return c.SetReadDeadline(time.Now().Add(pongWait))
	289	+	})
	290	+	done := make(chan struct{})
	291	+	defer close(done)
300	292		go func() {
301		-	<-r.Context().Done()
302		-	err := sshClient.RemoveSession(id)
303		-	if err != nil {
304		-	logger.Error("pipe remove error", "topic", topic, "info", clientInfo, "err", err.Error())
	293	+	ticker := time.NewTicker(pingPeriod)
	294	+	defer ticker.Stop()
	295	+	for {
	296	+	select {
	297	+	case <-done:
	298	+	return
	299	+	case <-ticker.C:
	300	+	if err := c.WriteControl(websocket.PingMessage, nil, time.Now().Add(writeWait)); err != nil {
	301	+	_ = c.Close()
	302	+	return
	303	+	}
	304	+	}
305	305		}
306		-	_ = c.Close()
307	306		}()
308	307
309	308		var wg sync.WaitGroup

...	...		@@ -339,10 +338,9 @@ func (web *webRouter) handlePipe() http.HandlerFunc {
339	338		}()
340	339
341	340		var messageBuffer []byte
	341	+	buf := make([]byte, 32*1024)
342	342
343	343		for {
344		-	buf := make([]byte, 32*1024)
345		-
346	344		n, err := p.Read(buf)
347	345		if err != nil {
348	346		logger.Error("pipe read error", "topic", topic, "info", clientInfo, "err", err.Error())

...	...		@@ -435,8 +433,9 @@ func rssHandler(cfg *Config, dbpool DB) http.HandlerFunc {
435	433		}
436	434
437	435		type webRouter struct {
438		-	cfg *Config
439		-	db DB
	436	+	cfg *Config
	437	+	db DB
	438	+	conn *pipeConn
440	439		}
441	440
442	441		func (web webRouter) static(mux http.ServeMux) {

...	...		@@ -489,14 +488,12 @@ func StartApiServer() {
489	488
490	489		info := shared.NewPicoPipeClient()
491	490
492		-	client, err := pipe.NewClient(context.Background(), logger.With("info", info), info)
	491	+	pingClient, err := pipe.NewClient(context.Background(), logger.With("info", info), info)
493	492		if err != nil {
494	493		panic(err)
495	494		}
496	495
497		-	sshClient = client
498		-
499		-	pingSession, err := sshClient.AddSession("ping", "pub -b=false -c ping", 0, -1, -1)
	496	+	pingSession, err := pingClient.AddSession("ping", "pub -b=false -c ping", 0, -1, -1)
500	497		if err != nil {
501	498		panic(err)
502	499		}

...	...		@@ -512,7 +509,7 @@ func StartApiServer() {
512	509		}
513	510		}()
514	511
515		-	web := &webRouter{cfg: cfg, db: db}
	512	+	web := &webRouter{cfg: cfg, db: db, conn: newPipeConn(info)}
516	513		handler := router.NewHostRouter(cfg.Domain, cfg.Space, web.mainMux(), web.sitesMux())
517	514
518	515		portStr := fmt.Sprintf(":%s", cfg.Port)

...	...		@@ -0,0 +1,100 @@
	1	+	package pipe
	2	+
	3	+	import (
	4	+	"io"
	5	+	"sync"
	6	+
	7	+	"github.com/picosh/utils/pipe"
	8	+	"golang.org/x/crypto/ssh"
	9	+	)
	10	+
	11	+	// pipeConn is pipe-web's ssh connection to pipe. Each request runs its
	12	+	// command in a session of its own, which is never reopened, so the command
	13	+	// ends when the request does.
	14	+	type pipeConn struct {
	15	+	info *pipe.SSHClientInfo
	16	+
	17	+	mu sync.Mutex
	18	+	client *ssh.Client
	19	+	}
	20	+
	21	+	func newPipeConn(info pipe.SSHClientInfo) pipeConn {
	22	+	return &pipeConn{info: info}
	23	+	}
	24	+
	25	+	// connect returns the open connection, dialing a new one when there is none
	26	+	// or when the open one is dead.
	27	+	func (c pipeConn) connect(dead ssh.Client) (*ssh.Client, error) {
	28	+	c.mu.Lock()
	29	+	defer c.mu.Unlock()
	30	+	if c.client != nil && c.client != dead {
	31	+	return c.client, nil
	32	+	}
	33	+	if c.client != nil {
	34	+	_ = c.client.Close()
	35	+	c.client = nil
	36	+	}
	37	+	client, err := pipe.NewSSHClient(c.info)
	38	+	if err != nil {
	39	+	return nil, err
	40	+	}
	41	+	c.client = client
	42	+	return client, nil
	43	+	}
	44	+
	45	+	// run starts cmd on pipe. A connection that can't open a session is dialed
	46	+	// again once.
	47	+	func (c pipeConn) run(cmd string) (cmdSession, error) {
	48	+	client, err := c.connect(nil)
	49	+	if err != nil {
	50	+	return nil, err
	51	+	}
	52	+	s, err := startCmd(client, cmd)
	53	+	if err == nil {
	54	+	return s, nil
	55	+	}
	56	+	client, err = c.connect(client)
	57	+	if err != nil {
	58	+	return nil, err
	59	+	}
	60	+	return startCmd(client, cmd)
	61	+	}
	62	+
	63	+	// cmdSession is a command running on pipe, read from its stdout and written
	64	+	// to its stdin.
	65	+	type cmdSession struct {
	66	+	*ssh.Session
	67	+	stdin io.WriteCloser
	68	+	stdout io.Reader
	69	+	}
	70	+
	71	+	func startCmd(client ssh.Client, cmd string) (cmdSession, error) {
	72	+	session, err := client.NewSession()
	73	+	if err != nil {
	74	+	return nil, err
	75	+	}
	76	+	s := &cmdSession{Session: session}
	77	+	if s.stdin, err = session.StdinPipe(); err == nil {
	78	+	if s.stdout, err = session.StdoutPipe(); err == nil {
	79	+	err = session.Start(cmd)
	80	+	}
	81	+	}
	82	+	if err != nil {
	83	+	_ = session.Close()
	84	+	return nil, err
	85	+	}
	86	+	return s, nil
	87	+	}
	88	+
	89	+	func (s *cmdSession) Read(p []byte) (int, error) {
	90	+	return s.stdout.Read(p)
	91	+	}
	92	+
	93	+	func (s *cmdSession) Write(p []byte) (int, error) {
	94	+	return s.stdin.Write(p)
	95	+	}
	96	+
	97	+	// CloseWrite sends the command EOF on its stdin.
	98	+	func (s *cmdSession) CloseWrite() error {
	99	+	return s.stdin.Close()
	100	+	}

...	...		@@ -351,8 +351,8 @@ func waitForPort(t *testing.T, port string, open bool) {
351	351		}
352	352		}
353	353
354		-	// waitForClients waits until n subscribers and pipes have joined their
355		-	// topics, so a message published next reaches them.
	354	+	// waitForClients waits until n subscribers and pipes are on their topics,
	355	+	// so a message published next reaches them.
356	356		func (s TestSSHServer) waitForClients(t testing.T, n int) {
357	357		t.Helper()
358	358		ps := s.PipeHandler.PubSub.(*psub.Multicast)

...	...		@@ -365,11 +365,11 @@ func (s TestSSHServer) waitForClients(t testing.T, n int) {
365	365		for range ps.GetPipes() {
366	366		joined++
367	367		}
368		-	if joined >= n {
	368	+	if joined == n {
369	369		return
370	370		}
371	371		if time.Now().After(deadline) {
372		-	t.Fatalf("%d of %d subscribers and pipes joined", joined, n)
	372	+	t.Fatalf("%d subscribers and pipes, want %d", joined, n)
373	373		}
374	374		time.Sleep(5 * time.Millisecond)
375	375		}

...	...		@@ -0,0 +1,108 @@
	1	+	package pipe
	2	+
	3	+	import (
	4	+	"context"
	5	+	"encoding/pem"
	6	+	"io"
	7	+	"net/http"
	8	+	"net/http/httptest"
	9	+	"os"
	10	+	"path/filepath"
	11	+	"strings"
	12	+	"testing"
	13	+
	14	+	"github.com/gorilla/websocket"
	15	+	"github.com/picosh/utils/pipe"
	16	+	)
	17	+
	18	+	// newWebTest runs pipe-web against a test pipe server.
	19	+	func newWebTest(t testing.T) (TestSSHServer, *httptest.Server) {
	20	+	t.Helper()
	21	+	server := NewTestSSHServer(t)
	22	+	t.Cleanup(server.Shutdown)
	23	+
	24	+	user := GenerateUser("pico")
	25	+	RegisterUserWithServer(server, user)
	26	+	keyPath := filepath.Join(t.TempDir(), "key")
	27	+	key := pem.EncodeToMemory(&pem.Block{Type: "OPENSSH PRIVATE KEY", Bytes: user.privateKey})
	28	+	if err := os.WriteFile(keyPath, key, 0o600); err != nil {
	29	+	t.Fatal(err)
	30	+	}
	31	+
	32	+	conn := newPipeConn(&pipe.SSHClientInfo{
	33	+	RemoteHost: "localhost:" + server.Cfg.Port,
	34	+	RemoteHostname: "localhost",
	35	+	RemoteUser: user.username,
	36	+	KeyLocation: keyPath,
	37	+	})
	38	+	web := &webRouter{cfg: server.Cfg, db: server.DBPool, conn: conn}
	39	+	ts := httptest.NewServer(web.mainMux())
	40	+	t.Cleanup(ts.Close)
	41	+	return server, ts
	42	+	}
	43	+
	44	+	func TestWebSubEndsWithRequest(t *testing.T) {
	45	+	server, ts := newWebTest(t)
	46	+
	47	+	ctx, cancel := context.WithCancel(context.Background())
	48	+	req, err := http.NewRequestWithContext(ctx, http.MethodGet, ts.URL+"/topic/quiet", nil)
	49	+	if err != nil {
	50	+	t.Fatal(err)
	51	+	}
	52	+	go func() {
	53	+	res, err := http.DefaultClient.Do(req)
	54	+	if err == nil {
	55	+	_, _ = io.Copy(io.Discard, res.Body)
	56	+	_ = res.Body.Close()
	57	+	}
	58	+	}()
	59	+	server.waitForClients(t, 1)
	60	+
	61	+	cancel()
	62	+	server.waitForClients(t, 0)
	63	+	}
	64	+
	65	+	func TestWebPubReachesSub(t *testing.T) {
	66	+	server, ts := newWebTest(t)
	67	+
	68	+	received := make(chan string, 1)
	69	+	go func() {
	70	+	res, err := http.Get(ts.URL + "/topic/news")
	71	+	if err != nil {
	72	+	received <- err.Error()
	73	+	return
	74	+	}
	75	+	defer func() { _ = res.Body.Close() }()
	76	+	body, _ := io.ReadAll(res.Body)
	77	+	received <- string(body)
	78	+	}()
	79	+	server.waitForClients(t, 1)
	80	+
	81	+	res, err := http.Post(ts.URL+"/topic/news", "text/plain", strings.NewReader("hello"))
	82	+	if err != nil {
	83	+	t.Fatal(err)
	84	+	}
	85	+	_ = res.Body.Close()
	86	+	if res.StatusCode != http.StatusOK {
	87	+	t.Errorf("pub status = %d", res.StatusCode)
	88	+	}
	89	+	if got := <-received; got != "hello" {
	90	+	t.Errorf("sub got %q, want %q", got, "hello")
	91	+	}
	92	+	server.waitForClients(t, 0)
	93	+	}
	94	+
	95	+	func TestWebPipeEndsWhenClientLeaves(t *testing.T) {
	96	+	server, ts := newWebTest(t)
	97	+
	98	+	c, res, err := websocket.DefaultDialer.Dial("ws"+strings.TrimPrefix(ts.URL, "http")+"/pipe/chat", nil)
	99	+	if err != nil {
	100	+	t.Fatal(err)
	101	+	}
	102	+	_ = res.Body.Close()
	103	+	server.waitForClients(t, 1)
	104	+
	105	+	// Drop the connection without a close frame, as a client that crashed.
	106	+	_ = c.NetConn().Close()
	107	+	server.waitForClients(t, 0)
	108	+	}

Commit af2fdb0

Commit `af2fdb0`