fix: http-lp forwarding - a download the agent finished reading reaches the browser whole; the session closing on its end of stream cut off what was still queued (up to 32 MB), so browsers restarted big downloads

internal/server/forward.go

@@ -129,20 +129,39 @@ func (m *forwardManager) bridgeHTTP(rule ForwardRule, conn net.Conn) {
 		}
 	}()
 
-	for {
-		select {
-		case <-sess.Done():
-			return
-		case out, ok := <-sess.Out():
-			if !ok || out.Exit {
-				return
+	// write passes one message on; false: the stream is over.
+	write := func(out proto.OutMessage) bool {
+		if out.Exit {
+			return false
 		}
 		if out.Data != "" {
 			b, err := base64.StdEncoding.DecodeString(out.Data)
 			if err != nil {
-					return
+				return false
 			}
 			if _, err := conn.Write(b); err != nil {
+				return false
+			}
+		}
+		return true
+	}
+	for {
+		select {
+		case out, ok := <-sess.Out():
+			if !ok || !write(out) {
+				return
+			}
+		case <-sess.Done():
+			// The agent's end of stream closes the session at once, while
+			// what it sent before may still be queued (a download read
+			// faster than the browser takes it): that goes out first.
+			for {
+				select {
+				case out := <-sess.Out():
+					if !write(out) {
+						return
+					}
+				default:
 					return
 				}
 			}

internal/server/forward_http_test.go

@@ -0,0 +1,49 @@
+package server
+
+import (
+	"bytes"
+	"encoding/base64"
+	"io"
+	"net"
+	"testing"
+	"time"
+
+	"github.com/x-grk/longpolling-remote-access/internal/proto"
+)
+
+// An http-lp download the agent has finished (its last chunk, then the end)
+// reaches the browser whole: the session closing does not cut off what is
+// still queued for it.
+func TestForwardHTTPDrainsAfterClose(t *testing.T) {
+	hub := NewHub()
+	m := &forwardManager{hub: hub}
+	user, host := net.Pipe()
+	go m.bridgeHTTP(ForwardRule{AgentID: "a1", ClientPort: 80, Proto: "http"}, host)
+
+	var sess *Session
+	for i := 0; i < 100 && sess == nil; i++ {
+		hub.mu.Lock()
+		for _, s := range hub.sessions {
+			sess = s
+		}
+		hub.mu.Unlock()
+		time.Sleep(10 * time.Millisecond)
+	}
+	if sess == nil {
+		t.Fatal("no session")
+	}
+	want := bytes.Repeat([]byte("0123456789abcdef"), 4096) // 64 KiB in 64 chunks
+	go func() {
+		for i := 0; i < len(want); i += 1024 {
+			sess.PushOut(proto.OutMessage{Data: base64.StdEncoding.EncodeToString(want[i : i+1024])})
+		}
+		sess.PushOut(proto.OutMessage{Exit: true})
+		sess.close() // as the server does on the agent's end of stream
+	}()
+	time.Sleep(50 * time.Millisecond) // all of it queued, the session closed, before the browser reads
+	_ = user.SetReadDeadline(time.Now().Add(5 * time.Second))
+	got, _ := io.ReadAll(user)
+	if !bytes.Equal(got, want) {
+		t.Fatalf("got %d of %d bytes", len(got), len(want))
+	}
+}