diff --git a/internal/server/forward.go b/internal/server/forward.go index 91ea1b9..6755c04 100644 --- a/internal/server/forward.go +++ b/internal/server/forward.go @@ -129,20 +129,39 @@ func (m *forwardManager) bridgeHTTP(rule ForwardRule, conn net.Conn) { } }() + // 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 false + } + if _, err := conn.Write(b); err != nil { + return false + } + } + return true + } for { select { - case <-sess.Done(): - return case out, ok := <-sess.Out(): - if !ok || out.Exit { + if !ok || !write(out) { return } - if out.Data != "" { - b, err := base64.StdEncoding.DecodeString(out.Data) - if err != nil { - return - } - if _, err := conn.Write(b); err != nil { + 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 } } diff --git a/internal/server/forward_http_test.go b/internal/server/forward_http_test.go new file mode 100644 index 0000000..25639ca --- /dev/null +++ b/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)) + } +}