Merge remote-tracking branch 'origin/main' into feature/shared-service-config-loader

# Conflicts:
#	go.mod
#	go.sum

util/file.go

@@ -56,9 +56,9 @@ func WriteJson(ctx context.Context, file string, obj interface{}) error {
 }
 
 // DirectWriteJson writes JSON config object to a file creating parent directories if required without creating a temporary file
-func DirectWriteJson(ctx context.Context, file string, obj interface{}) error {
+func DirectWriteJson(ctx context.Context, file string, obj interface{}) (err error) {
 
-	_, _, err := prepareConfigFileDir(file)
+	_, _, err = prepareConfigFileDir(file)
 	if err != nil {
 		return err
 	}
@@ -68,11 +68,24 @@ func DirectWriteJson(ctx context.Context, file string, obj interface{}) error {
 		return err
 	}
 
+	// Named return so a failed Close is reported rather than logged and
+	// swallowed: the write is only durable once the file closes cleanly, and a
+	// caller told "written" would carry on with data that never landed.
 	defer func() {
-		err = targetFile.Close()
-		if err != nil {
-			log.Errorf("failed to close file %s: %v", file, err)
+		cerr := targetFile.Close()
+		if cerr == nil {
+			return
 		}
+		if err == nil {
+			// Returned, not logged: the caller reports it once.
+			err = cerr
+			return
+		}
+		// The body already failed and that error is the one the caller gets, so
+		// it is the one that explains the failure. This is then the only place
+		// the close failure can surface — at debug, per the logging rules for
+		// close errors on writes.
+		log.Debugf("failed to close file %s after %v: %v", file, err, cerr)
 	}()
 
 	// make it pretty
@@ -149,7 +162,7 @@ func writeBytes(ctx context.Context, file string, configDir string, configFileNa
 		return fmt.Errorf("after temp file: %w", ctx.Err())
 	}
 
-	if err = os.Rename(tempFileName, file); err != nil {
+	if err = renameFile(tempFileName, file); err != nil {
 		return fmt.Errorf("move %s to %s: %w", tempFileName, file, err)
 	}
 
@@ -182,7 +195,7 @@ func openOrCreateFile(file string) (*os.File, error) {
 // ReadJson reads JSON config file and maps to a provided interface
 func ReadJson(file string, res interface{}) (interface{}, error) {
 
-	f, err := os.Open(file)
+	f, err := openRead(file)
 	if err != nil {
 		return nil, err
 	}
@@ -235,7 +248,7 @@ func ListFiles(dir, pattern string) ([]string, error) {
 func ReadJsonWithEnvSub(file string, res interface{}) (interface{}, error) {
 	envVars := getEnvMap()
 
-	f, err := os.Open(file)
+	f, err := openRead(file)
 	if err != nil {
 		return nil, err
 	}

util/file_nonwindows.go

@@ -0,0 +1,16 @@
+//go:build !windows
+
+package util
+
+import "os"
+
+// openRead opens path for reading. Only Windows needs more than this: there a
+// plain open holds the file against the rename that replaces it.
+func openRead(path string) (*os.File, error) {
+	return os.Open(path)
+}
+
+// renameFile replaces newpath with oldpath.
+func renameFile(oldpath, newpath string) error {
+	return os.Rename(oldpath, newpath)
+}

util/file_read_test.go

@@ -0,0 +1,58 @@
+package util
+
+import (
+	"errors"
+	"os"
+	"path/filepath"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestReadJson_ReadsTheFile(t *testing.T) {
+	path := filepath.Join(t.TempDir(), "state.json")
+	require.NoError(t, os.WriteFile(path, []byte(`{"SomeField": 7}`), 0o600))
+
+	var got TestConfig
+	_, err := ReadJson(path, &got)
+
+	require.NoError(t, err)
+	assert.Equal(t, 7, got.SomeField, "the decoded value")
+}
+
+// Callers tell a missing file from a broken one so they can seed a default in
+// its place. The Windows path opens through a root and rebuilds the error, so
+// the mapping has to survive that.
+func TestReadJson_MissingFileIsErrNotExist(t *testing.T) {
+	dir := t.TempDir()
+
+	for _, tc := range []struct {
+		name string
+		path string
+	}{
+		{"missing file", filepath.Join(dir, "absent.json")},
+		{"missing directory", filepath.Join(dir, "absent", "absent.json")},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			var got TestConfig
+			_, err := ReadJson(tc.path, &got)
+
+			require.Error(t, err)
+			assert.ErrorIs(t, err, os.ErrNotExist)
+			assert.Contains(t, err.Error(), tc.path, "the error names the file the caller asked for")
+		})
+	}
+}
+
+func TestReadJson_MalformedFileIsNotErrNotExist(t *testing.T) {
+	path := filepath.Join(t.TempDir(), "state.json")
+	require.NoError(t, os.WriteFile(path, []byte("{not json"), 0o600))
+
+	var got TestConfig
+	_, err := ReadJson(path, &got)
+
+	require.Error(t, err)
+	assert.False(t, errors.Is(err, os.ErrNotExist),
+		"a file that is there but unreadable must not be seeded over: %v", err)
+}

util/file_windows.go

@@ -0,0 +1,79 @@
+package util
+
+import (
+	"errors"
+	"io/fs"
+	"os"
+	"path/filepath"
+)
+
+// openRead opens path for reading without holding it against a rename.
+//
+// os.Open does not set FILE_SHARE_DELETE on Windows, so you cannot rename an
+// open file like on UNIX. This caused concurrency issues with active state
+// config file.
+//
+// os.Root opens through NtCreateFile with delete sharing, which is the
+// behaviour Unix has.
+// https://cs.opensource.google/go/go/+/refs/tags/go1.27.1:src/os/root_windows.go;drc=a4f5d9bbdbdf42da7e2d7e976ac85753c4db5d75;l=176
+func openRead(path string) (*os.File, error) {
+	root, err := os.OpenRoot(filepath.Dir(path))
+	if err != nil {
+		// Names the file the caller asked for, not the directory the root
+		// failed on, so a missing directory reads like a missing file.
+		return nil, pathError("open", path, err)
+	}
+	defer func() { _ = root.Close() }()
+
+	// The file outlives the root: closing a Root closes the directory handle it
+	// holds, not the files opened through it.
+	f, err := root.Open(filepath.Base(path))
+	if err != nil {
+		return nil, pathError("open", path, err)
+	}
+	return f, nil
+}
+
+// renameFile replaces newpath with oldpath, including while something holds
+// newpath open for reading.
+//
+// os.Root.Rename asks for POSIX semantics, which unlink the destination
+// immediately and leave open handles reading the version they opened.
+// https://cs.opensource.google/go/go/+/master:src/internal/syscall/windows/at_windows.go;drc=a4f5d9bbdbdf42da7e2d7e976ac85753c4db5d75;l=384
+func renameFile(oldpath, newpath string) error {
+	dir := filepath.Dir(newpath)
+	if filepath.Dir(oldpath) != dir {
+		return os.Rename(oldpath, newpath)
+	}
+
+	root, err := os.OpenRoot(dir)
+	if err != nil {
+		return os.Rename(oldpath, newpath)
+	}
+	defer func() { _ = root.Close() }()
+
+	if err := root.Rename(filepath.Base(oldpath), filepath.Base(newpath)); err != nil {
+		return linkError("rename", oldpath, newpath, err)
+	}
+	return nil
+}
+
+// pathError restores the full path on an error from a root, which names the
+// file by the base name it was opened with.
+func pathError(op, path string, err error) error {
+	var perr *fs.PathError
+	if errors.As(err, &perr) {
+		err = perr.Err
+	}
+	return &fs.PathError{Op: op, Path: path, Err: err}
+}
+
+// linkError does the same as pathError for a rename, which reports both files
+// by their base names.
+func linkError(op, oldpath, newpath string, err error) error {
+	var lerr *os.LinkError
+	if errors.As(err, &lerr) {
+		err = lerr.Err
+	}
+	return &os.LinkError{Op: op, Old: oldpath, New: newpath, Err: err}
+}

util/file_windows_test.go

@@ -0,0 +1,116 @@
+package util
+
+import (
+	"context"
+	"io"
+	"os"
+	"path/filepath"
+	"sync"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+// seedReplace lays out a write as writeBytes leaves it: the destination that
+// exists and the temp file that is to take its place.
+func seedReplace(t *testing.T) (src, dst string) {
+	t.Helper()
+	dir := t.TempDir()
+	src = filepath.Join(dir, ".tmpstate.json")
+	dst = filepath.Join(dir, "state.json")
+	require.NoError(t, os.WriteFile(src, []byte(`{"SomeField": 2}`), 0o600))
+	require.NoError(t, os.WriteFile(dst, []byte(`{"SomeField": 1}`), 0o600))
+	return src, dst
+}
+
+// The reader has to share the file for delete, or the rename cannot take
+// delete access on it. Regression test.
+func TestRenameFile_ReplacesAFileBeingRead(t *testing.T) {
+	t.Run("a reader that shares delete", func(t *testing.T) {
+		src, dst := seedReplace(t)
+
+		f, err := openRead(dst)
+		require.NoError(t, err)
+		defer f.Close()
+
+		require.Error(t, os.Rename(src, dst),
+			"delete sharing alone has to be too little, or this test proves nothing")
+		require.NoError(t, renameFile(src, dst), "POSIX semantics have to get the replace through")
+
+		// The handle stays on the file it opened, so a read in flight finishes
+		// on that version instead of seeing the replacement.
+		held, err := io.ReadAll(f)
+		require.NoError(t, err)
+		assert.JSONEq(t, `{"SomeField": 1}`, string(held), "the version the reader opened")
+
+		landed, err := os.ReadFile(dst)
+		require.NoError(t, err)
+		assert.JSONEq(t, `{"SomeField": 2}`, string(landed), "the version the writer put there")
+	})
+
+	t.Run("a reader that does not", func(t *testing.T) {
+		src, dst := seedReplace(t)
+
+		f, err := os.Open(dst)
+		require.NoError(t, err)
+		defer f.Close()
+
+		require.Error(t, renameFile(src, dst),
+			"a plain read still holds the file, and the caller is owed that error")
+	})
+
+	t.Run("no readers at all", func(t *testing.T) {
+		src, dst := seedReplace(t)
+
+		require.NoError(t, renameFile(src, dst))
+
+		landed, err := os.ReadFile(dst)
+		require.NoError(t, err)
+		assert.JSONEq(t, `{"SomeField": 2}`, string(landed), "the destination holds what replaced it")
+	})
+}
+
+// A config rewritten while it is being read, which is the daemon reading the
+// active profile against a profile switch writing it.
+func TestReadJsonWriteJson_Concurrently(t *testing.T) {
+	path := filepath.Join(t.TempDir(), "state.json")
+	require.NoError(t, WriteJson(context.Background(), path, &TestConfig{SomeField: 1}))
+
+	var wg sync.WaitGroup
+	errs := make(chan error, 128)
+
+	for i := 0; i < 8; i++ {
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			for r := 0; r < 50; r++ {
+				var got TestConfig
+				if _, err := ReadJson(path, &got); err != nil {
+					errs <- err
+					return
+				}
+			}
+		}()
+	}
+
+	for i := 0; i < 2; i++ {
+		wg.Add(1)
+		go func(writer int) {
+			defer wg.Done()
+			for r := 0; r < 50; r++ {
+				if err := WriteJson(context.Background(), path, &TestConfig{SomeField: writer}); err != nil {
+					errs <- err
+					return
+				}
+			}
+		}(i)
+	}
+
+	wg.Wait()
+	close(errs)
+
+	for err := range errs {
+		assert.NoError(t, err, "a read and a write of the same config must not collide")
+	}
+}

util/serviceurl.go

@@ -0,0 +1,69 @@
+package util
+
+import (
+	"net/url"
+	"strconv"
+	"strings"
+)
+
+// SameServiceURL reports whether two service URLs address the same endpoint.
+// One endpoint can be written several ways, and every spelling below reaches
+// the same server, so none of them is a divergence from another:
+//
+//	an implicit default port    https://mgmt.example.com   :443
+//	a zero-padded port          https://mgmt.example.com:0443
+//	a different host case       https://MGMT.example.com
+//	a trailing slash            https://mgmt.example.com/
+//
+// A path is otherwise part of the identity: https://mgmt.example.com and
+// https://mgmt.example.com/other are two endpoints.
+//
+// It lives here rather than next to any one caller because several of them
+// compare the same kind of URL — an MDM-enforced management URL against a
+// requested one, a stored profile URL against a command-line one — and every
+// copy of these rules that drifts turns an equivalent URL into a refused
+// request.
+func SameServiceURL(a, b *url.URL) bool {
+	if a == nil || b == nil {
+		return a == b
+	}
+
+	return strings.EqualFold(a.Hostname(), b.Hostname()) &&
+		strings.EqualFold(a.Scheme, b.Scheme) &&
+		ServiceURLPort(a) == ServiceURLPort(b) &&
+		strings.TrimSuffix(a.Path, "/") == strings.TrimSuffix(b.Path, "/")
+}
+
+// SameServiceURLStrings is SameServiceURL for unparsed input. Input that does
+// not parse falls back to string equality, which is the strictest thing left
+// to do with it.
+func SameServiceURLStrings(a, b string) bool {
+	ua, errA := url.ParseRequestURI(a)
+	ub, errB := url.ParseRequestURI(b)
+	if errA != nil || errB != nil {
+		return a == b
+	}
+
+	return SameServiceURL(ua, ub)
+}
+
+// ServiceURLPort is the port a URL addresses: the one it carries, normalized
+// numerically so ":0443" and ":443" are one port, or the scheme's default.
+func ServiceURLPort(u *url.URL) string {
+	port := u.Port()
+	if port == "" {
+		switch strings.ToLower(u.Scheme) {
+		case "https":
+			return "443"
+		case "http":
+			return "80"
+		default:
+			return ""
+		}
+	}
+
+	if n, err := strconv.Atoi(port); err == nil {
+		return strconv.Itoa(n)
+	}
+	return port
+}

util/serviceurl_test.go

@@ -0,0 +1,73 @@
+package util
+
+import (
+	"net/url"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestSameServiceURLSpellings(t *testing.T) {
+	tests := []struct {
+		a, b string
+		want bool
+	}{
+		// One endpoint, written several ways.
+		{a: "https://mgmt.example.com", b: "https://mgmt.example.com:443", want: true},
+		{a: "https://mgmt.example.com", b: "https://mgmt.example.com/", want: true},
+		{a: "https://mgmt.example.com/", b: "https://mgmt.example.com:443/", want: true},
+		{a: "https://MGMT.example.com", b: "https://mgmt.example.com", want: true},
+		{a: "https://mgmt.example.com:0443", b: "https://mgmt.example.com:443", want: true},
+		{a: "http://mgmt.example.com", b: "http://mgmt.example.com:80", want: true},
+		{a: "HTTPS://mgmt.example.com", b: "https://mgmt.example.com", want: true},
+
+		// Different endpoints.
+		{a: "https://mgmt.example.com", b: "http://mgmt.example.com", want: false},
+		{a: "https://mgmt.example.com", b: "https://mgmt.example.com:8443", want: false},
+		{a: "https://mgmt.example.com", b: "https://other.example.com", want: false},
+		{a: "https://mgmt.example.com", b: "https://mgmt.example.com/other", want: false},
+
+		// Unparseable input falls back to string equality.
+		{a: "mgmt.example.com", b: "mgmt.example.com", want: true},
+		{a: "mgmt.example.com", b: "https://mgmt.example.com", want: false},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.a+" vs "+tt.b, func(t *testing.T) {
+			assert.Equal(t, tt.want, SameServiceURLStrings(tt.a, tt.b))
+			assert.Equal(t, tt.want, SameServiceURLStrings(tt.b, tt.a), "the comparison must be symmetric")
+		})
+	}
+}
+
+// The parsed form is the primitive the string form delegates to, so it must
+// answer the same for a spelling that only the parser can tell apart.
+func TestSameServiceURLParsed(t *testing.T) {
+	parse := func(raw string) *url.URL {
+		t.Helper()
+		u, err := url.ParseRequestURI(raw)
+		require.NoError(t, err)
+		return u
+	}
+
+	assert.True(t, SameServiceURL(parse("https://mgmt.example.com:0443/"), parse("https://MGMT.example.com")))
+	assert.False(t, SameServiceURL(parse("https://mgmt.example.com"), parse("https://mgmt.example.com:8443")))
+
+	assert.True(t, SameServiceURL(nil, nil), "two absent URLs are the same absence")
+	assert.False(t, SameServiceURL(nil, parse("https://mgmt.example.com")))
+}
+
+func TestServiceURLPort(t *testing.T) {
+	parse := func(raw string) *url.URL {
+		t.Helper()
+		u, err := url.ParseRequestURI(raw)
+		require.NoError(t, err)
+		return u
+	}
+
+	assert.Equal(t, "443", ServiceURLPort(parse("https://mgmt.example.com")))
+	assert.Equal(t, "80", ServiceURLPort(parse("http://mgmt.example.com")))
+	assert.Equal(t, "443", ServiceURLPort(parse("https://mgmt.example.com:0443")))
+	assert.Equal(t, "8443", ServiceURLPort(parse("https://mgmt.example.com:8443")))
+}

util/wsproxy/server/proxy.go

@@ -1,12 +1,8 @@
 package server
 
 import (
-	"context"
-	"io"
-	"net"
 	"net/http"
-	"sync"
-	"time"
+	"sync/atomic"
 
 	"github.com/coder/websocket"
 	log "github.com/sirupsen/logrus"
@@ -15,11 +11,6 @@ import (
 	"github.com/netbirdio/netbird/util/wsproxy"
 )
 
-const (
-	bufferSize = 32 * 1024
-	ioTimeout  = 5 * time.Second
-)
-
 // Config contains the configuration for the WebSocket proxy.
 type Config struct {
 	Handler         http.Handler
@@ -53,14 +44,23 @@ func New(handler http.Handler, opts ...Option) *Proxy {
 
 // Handler returns an http.Handler that proxies WebSocket connections to the local gRPC server.
 func (p *Proxy) Handler() http.Handler {
-	return http.HandlerFunc(p.handleWebSocket)
+	return &proxyHandler{
+		metrics: p.config.MetricsRecorder,
+		handler: p.config.Handler,
+	}
 }
 
-func (p *Proxy) handleWebSocket(w http.ResponseWriter, r *http.Request) {
+type proxyHandler struct {
+	metrics MetricsRecorder
+	handler http.Handler
+	conn    atomic.Pointer[wsConnAdapter]
+}
+
+func (ph *proxyHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 	ctx := r.Context()
 
-	p.metrics.RecordConnection(ctx)
-	defer p.metrics.RecordDisconnection(ctx)
+	ph.metrics.RecordConnection(ctx)
+	defer ph.metrics.RecordDisconnection(ctx)
 
 	log.Debugf("WebSocket proxy handling connection from %s, forwarding to internal gRPC handler", r.RemoteAddr)
 	acceptOptions := &websocket.AcceptOptions{
@@ -69,121 +69,44 @@ func (p *Proxy) handleWebSocket(w http.ResponseWriter, r *http.Request) {
 
 	wsConn, err := websocket.Accept(w, r, acceptOptions)
 	if err != nil {
-		p.metrics.RecordError(ctx, "websocket_accept_failed")
+		ph.metrics.RecordError(ctx, "websocket_accept_failed")
 		log.Errorf("WebSocket upgrade failed from %s: %v", r.RemoteAddr, err)
 		return
 	}
-	defer func() {
-		_ = wsConn.Close(websocket.StatusNormalClosure, "")
-	}()
+	serverConn := (&wsConnAdapter{
+		ctx:        ctx,
+		conn:       wsConn,
+		metrics:    ph.metrics,
+		clientAddr: r.RemoteAddr,
+	})
 
-	clientConn, serverConn := net.Pipe()
 	defer func() {
-		_ = clientConn.Close()
 		_ = serverConn.Close()
 	}()
 
+	ph.conn.Store(serverConn) // used in tests only
+
 	log.Debugf("WebSocket proxy established: %s -> gRPC handler", r.RemoteAddr)
 
-	go func() {
-		(&http2.Server{}).ServeConn(serverConn, &http2.ServeConnOpts{
-			Context: ctx,
-			Handler: p.config.Handler,
-		})
-	}()
+	(&http2.Server{
+		// TODO (dmitri) we should limit the number of concurrent streams per connection (peer)
+		// and idle timeouts
+		// MaxConcurrentStreams: 20,
+		// IdleTimeout: 10 * time.Second,
+	}).ServeConn(serverConn, &http2.ServeConnOpts{
+		Context:    ctx,
+		Handler:    ph.handler,
+		BaseConfig: &http.Server{
+			// b/c we are wrapping a ws connection, read and write connection deadlines normally set
+			// via ReadTimeout and WriteTimeout http.Server fields aren't available to us. The ws
+			// library doesn't expose connection deadline timer config, and we ignore these calls in "wsConnAdapter".
+			//
+			// Another issue is that Server.ServeConn() call bypasses setting of connection deadlines altogether,
+			// ReadTimeout and Writetimeout set here would only apply to h2 streams, i.e. after a HEADERS frame
+			// arrival and processing, turning ReadTimeout into a request body read deadline, and WriteTimeout into
+			// a response deadline (the latter not useful for streaming requests).
+		},
+	})
 
-	p.proxyData(ctx, wsConn, clientConn, r.RemoteAddr)
-}
-
-func (p *Proxy) proxyData(ctx context.Context, wsConn *websocket.Conn, pipeConn net.Conn, clientAddr string) {
-	proxyCtx, cancel := context.WithCancel(ctx)
-	defer cancel()
-
-	var wg sync.WaitGroup
-	wg.Add(2)
-
-	go p.wsToPipe(proxyCtx, cancel, &wg, wsConn, pipeConn, clientAddr)
-	go p.pipeToWS(proxyCtx, cancel, &wg, wsConn, pipeConn, clientAddr)
-
-	wg.Wait()
-}
-
-func (p *Proxy) wsToPipe(ctx context.Context, cancel context.CancelFunc, wg *sync.WaitGroup, wsConn *websocket.Conn, pipeConn net.Conn, clientAddr string) {
-	defer wg.Done()
-	defer cancel()
-
-	for {
-		msgType, data, err := wsConn.Read(ctx)
-		if err != nil {
-			switch {
-			case ctx.Err() != nil:
-				log.Debugf("WebSocket from %s terminating due to context cancellation", clientAddr)
-			case websocket.CloseStatus(err) != -1:
-				log.Debugf("WebSocket from %s disconnected", clientAddr)
-			default:
-				p.metrics.RecordError(ctx, "websocket_read_error")
-				log.Debugf("WebSocket read error from %s: %v", clientAddr, err)
-			}
-			return
-		}
-
-		if msgType != websocket.MessageBinary {
-			log.Warnf("Unexpected WebSocket message type from %s: %v", clientAddr, msgType)
-			continue
-		}
-
-		if ctx.Err() != nil {
-			log.Tracef("wsToPipe goroutine terminating due to context cancellation before pipe write")
-			return
-		}
-
-		if err := pipeConn.SetWriteDeadline(time.Now().Add(ioTimeout)); err != nil {
-			log.Debugf("Failed to set pipe write deadline: %v", err)
-		}
-
-		n, err := pipeConn.Write(data)
-		if err != nil {
-			p.metrics.RecordError(ctx, "pipe_write_error")
-			log.Warnf("Pipe write error for %s: %v", clientAddr, err)
-			return
-		}
-
-		p.metrics.RecordBytesTransferred(ctx, "ws_to_grpc", int64(n))
-	}
-}
-
-func (p *Proxy) pipeToWS(ctx context.Context, cancel context.CancelFunc, wg *sync.WaitGroup, wsConn *websocket.Conn, pipeConn net.Conn, clientAddr string) {
-	defer wg.Done()
-	defer cancel()
-
-	buf := make([]byte, bufferSize)
-	for {
-		n, err := pipeConn.Read(buf)
-		if err != nil {
-			if ctx.Err() != nil {
-				log.Tracef("pipeToWS goroutine terminating due to context cancellation")
-				return
-			}
-
-			if err != io.EOF {
-				log.Debugf("Pipe read error for %s: %v", clientAddr, err)
-			}
-			return
-		}
-
-		if ctx.Err() != nil {
-			log.Tracef("pipeToWS goroutine terminating due to context cancellation before WebSocket write")
-			return
-		}
-
-		if n > 0 {
-			if err := wsConn.Write(ctx, websocket.MessageBinary, buf[:n]); err != nil {
-				p.metrics.RecordError(ctx, "websocket_write_error")
-				log.Warnf("WebSocket write error for %s: %v", clientAddr, err)
-				return
-			}
-
-			p.metrics.RecordBytesTransferred(ctx, "grpc_to_ws", int64(n))
-		}
-	}
+	log.Debugf("WebSocket proxy closing: %s -> gRPC handler", r.RemoteAddr)
 }

util/wsproxy/server/ws_conn_adapter.go

@@ -0,0 +1,126 @@
+package server
+
+import (
+	"context"
+	"net"
+	"sync/atomic"
+	"time"
+
+	"github.com/coder/websocket"
+	log "github.com/sirupsen/logrus"
+)
+
+type wsConnAdapter struct {
+	prefix       string
+	ctx          context.Context
+	conn         *websocket.Conn
+	metrics      MetricsRecorder
+	clientAddr   string
+	closed       atomic.Bool
+	bufferedRead []byte
+}
+
+var _ net.Conn = &wsConnAdapter{}
+
+type wsAddr struct{ prefix string }
+
+func (wa wsAddr) Network() string { return wa.prefix + "ws-proxy" }
+func (wa wsAddr) String() string  { return wa.prefix + "ws-proxy" }
+
+func (ws *wsConnAdapter) Read(b []byte) (int, error) {
+	if len(ws.bufferedRead) > 0 {
+		return ws.readFromBuffer(b)
+	}
+
+	msgType, data, err := ws.conn.Read(ws.ctx)
+	if err != nil {
+		switch {
+		case ws.ctx.Err() != nil:
+			log.Debugf("WebSocket from %s terminating due to context cancellation", ws.clientAddr)
+		case websocket.CloseStatus(err) != -1:
+			log.Debugf("WebSocket from %s disconnected", ws.clientAddr)
+		default:
+			ws.recordError(ws.ctx, "websocket_read_error")
+			log.Debugf("WebSocket read error from %s: %v", ws.clientAddr, err)
+		}
+		return copy(b, data), err
+	}
+	if msgType != websocket.MessageBinary {
+		log.Warnf("Unexpected WebSocket message type from %s: %v", ws.clientAddr, msgType)
+		return 0, nil
+	}
+
+	ws.bufferedRead = data
+	return ws.readFromBuffer(b)
+}
+
+func (ws *wsConnAdapter) readFromBuffer(b []byte) (int, error) {
+	n := copy(b, ws.bufferedRead)
+
+	ws.recordBytesTransferred(ws.ctx, "ws_to_grpc", n)
+	if n == len(ws.bufferedRead) {
+		ws.bufferedRead = nil
+		return n, nil
+	} else {
+		ws.bufferedRead = ws.bufferedRead[n:]
+	}
+	return n, nil
+}
+
+func (ws *wsConnAdapter) Write(b []byte) (int, error) {
+	maybeErr := ws.ctx.Err()
+
+	n := len(b)
+	if n == 0 {
+		return n, maybeErr
+	}
+	if maybeErr != nil {
+		return 0, maybeErr
+	}
+	if err := ws.conn.Write(ws.ctx, websocket.MessageBinary, b[:n]); err != nil {
+		ws.recordError(ws.ctx, "websocket_write_error")
+		log.Warnf("WebSocket write error for %s: %v", ws.clientAddr, err)
+		return 0, err // we don't know how many bytes have been written
+	}
+
+	ws.recordBytesTransferred(ws.ctx, "grpc_to_ws", n)
+	return n, nil
+}
+
+func (ws *wsConnAdapter) Close() error {
+	ws.closed.Store(true)
+	return ws.conn.Close(websocket.StatusNormalClosure, "")
+}
+
+func (ws *wsConnAdapter) LocalAddr() net.Addr  { return wsAddr{ws.prefix} }
+func (ws *wsConnAdapter) RemoteAddr() net.Addr { return wsAddr{ws.prefix} }
+
+func (ws *wsConnAdapter) SetDeadline(t time.Time) error {
+	return nil
+}
+
+func (ws *wsConnAdapter) SetReadDeadline(t time.Time) error {
+	return nil
+}
+
+func (ws *wsConnAdapter) SetWriteDeadline(t time.Time) error {
+	return nil
+}
+
+func (ws *wsConnAdapter) recordError(ctx context.Context, errorType string) {
+	if ws.metrics == nil {
+		return
+	}
+	ws.metrics.RecordError(ctx, errorType)
+}
+
+func (ws *wsConnAdapter) recordBytesTransferred(ctx context.Context, direction string, bytes int) {
+	if ws.metrics == nil {
+		return
+	}
+	ws.metrics.RecordBytesTransferred(ctx, direction, int64(bytes))
+}
+
+func (ws *wsConnAdapter) IsClosed() bool {
+	return ws.closed.Load()
+}

util/wsproxy/server/ws_conn_adapter_test.go

@@ -0,0 +1,204 @@
+package server
+
+import (
+	"bytes"
+	"context"
+	"crypto/tls"
+	"io"
+	"math/rand/v2"
+	"net"
+	"net/http"
+	"os"
+	"path/filepath"
+	"strconv"
+	"strings"
+	"testing"
+	"time"
+
+	"github.com/coder/websocket"
+	"github.com/stretchr/testify/assert"
+	"golang.org/x/net/http2"
+	"golang.org/x/net/http2/hpack"
+)
+
+func TestAdapterHandlingConnectionClosures(t *testing.T) {
+	var cases = []struct {
+		description string
+		casenum     int
+	}{
+		{"client-side ws connection is closed", 0},
+		{"server-side ws connection is closed", 1},
+		{"client-side context is cancelled", 2},
+		{"server-side context is cancelled", 3},
+	}
+
+	for _, c := range cases {
+		t.Run(c.description, func(t *testing.T) {
+			serversock := filepath.Join(os.TempDir(), "http-server-"+strconv.FormatInt(rand.Int64(), 10)+".sock")
+			t.Cleanup(func() { os.Remove(serversock) })
+
+			l, err := net.Listen("unix", serversock)
+			assert.NoError(t, err)
+
+			proxy := New(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+				buf, _ := io.ReadAll(r.Body)
+				defer r.Body.Close()
+				w.Write([]byte("echo: " + string(buf))) //nolint:errcheck
+			}))
+
+			handler, ok := proxy.Handler().(*proxyHandler)
+			assert.True(t, ok)
+
+			protocols := new(http.Protocols)
+			protocols.SetHTTP1(true)
+			protocols.SetUnencryptedHTTP2(true)
+			httpServer := http.Server{
+				Handler: handler,
+			}
+			go httpServer.Serve(l) //nolint:errcheck
+			t.Cleanup(func() { httpServer.Close() })
+
+			clientconn, _, err := websocket.Dial(context.Background(), "http://whatever", //nolint:bodyclose
+				&websocket.DialOptions{HTTPClient: &http.Client{
+					Transport: &http.Transport{
+						DialContext: func(_ context.Context, _, _ string) (net.Conn, error) {
+							return net.Dial("unix", serversock)
+						},
+					}}})
+			assert.NoError(t, err)
+
+			clientCtx, cancel := context.WithCancel(context.Background()) //nolint:govet
+			h2client := &http.Client{
+				Transport: &http2.Transport{
+					AllowHTTP: true,
+					DialTLSContext: func(_ context.Context, _, _ string, _ *tls.Config) (net.Conn, error) {
+						return &wsConnAdapter{
+							prefix: "test-client",
+							ctx:    clientCtx,
+							conn:   clientconn,
+						}, nil
+					},
+				}}
+
+			resp, err := h2client.Post("http://whatever", "text/html", strings.NewReader("g'day"))
+			assert.NoError(t, err)
+
+			body, err := io.ReadAll(resp.Body)
+			defer resp.Body.Close()
+
+			assert.NoError(t, err)
+			assert.Equal(t, "echo: g'day", string(body))
+
+			switch c.casenum {
+			case 0:
+				clientconn.Close(websocket.StatusNormalClosure, "")
+			case 1:
+				handler.conn.Load().Close()
+			case 2:
+				cancel()
+			case 3:
+				resp.Body.Close()
+				h2client.CloseIdleConnections()
+			}
+
+			assert.EventuallyWithT(t, func(c *assert.CollectT) {
+				assert.True(c, handler.conn.Load().IsClosed())
+			}, 3*time.Second, 100*time.Millisecond)
+		}) //nolint:govet
+	}
+}
+
+func TestAdapterHandlingHttpConnection_NoHeadersSent(t *testing.T) {
+	t.Skip("currently disabled as it requires idle timeout to be set")
+
+	serversock := filepath.Join(os.TempDir(), "http-server-"+strconv.FormatInt(rand.Int64(), 10)+".sock")
+	defer os.Remove(serversock)
+
+	l, err := net.Listen("unix", serversock)
+	assert.NoError(t, err)
+
+	proxy := New(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		buf, _ := io.ReadAll(r.Body)
+		defer r.Body.Close()                    //nolint:errcheck
+		w.Write([]byte("echo: " + string(buf))) //nolint:errcheck
+	}))
+
+	handler, ok := proxy.Handler().(*proxyHandler)
+	assert.True(t, ok)
+
+	protocols := new(http.Protocols)
+	protocols.SetHTTP1(true)
+	protocols.SetUnencryptedHTTP2(true)
+	httpServer := http.Server{
+		Handler: handler,
+	}
+	go httpServer.Serve(l) //nolint:errcheck
+
+	clientconn, _, err := websocket.Dial(context.Background(), "http://whatever", //nolint:bodyclose
+		&websocket.DialOptions{HTTPClient: &http.Client{
+			Transport: &http.Transport{
+				DialContext: func(_ context.Context, _, _ string) (net.Conn, error) {
+					return net.Dial("unix", serversock)
+				},
+			}}})
+	assert.NoError(t, err)
+
+	h2client := &http.Client{
+		Transport: &http2.Transport{
+			AllowHTTP: true,
+			DialTLSContext: func(_ context.Context, _, _ string, _ *tls.Config) (net.Conn, error) {
+				return &h2ConnectionSnooper{wrappedConn: &wsConnAdapter{
+					prefix: "test-client",
+					ctx:    context.Background(),
+					conn:   clientconn,
+				}, shouldDropFrame: func(f http2.FrameType) bool { return f == http2.FrameHeaders || f == http2.FrameData }}, nil
+			},
+		}}
+
+	_, err = h2client.Post("http://whatever", "text/html", strings.NewReader("g'day"))
+	assert.Error(t, err)
+
+	assert.EventuallyWithT(t, func(c *assert.CollectT) {
+		assert.True(c, handler.conn.Load().IsClosed())
+	}, 3*time.Second, 100*time.Millisecond)
+}
+
+type h2ConnectionSnooper struct {
+	wrappedConn     net.Conn
+	shouldDropFrame func(f http2.FrameType) bool
+}
+
+func (hs *h2ConnectionSnooper) Read(b []byte) (n int, err error) {
+	return hs.wrappedConn.Read(b)
+}
+
+func (hs *h2ConnectionSnooper) Write(b []byte) (n int, err error) {
+	fr := http2.NewFramer(nil, bytes.NewReader(b))
+	fr.ReadMetaHeaders = hpack.NewDecoder(0, nil)
+	f, err := fr.ReadFrame()
+	if err != nil {
+		return hs.wrappedConn.Write(b)
+	}
+
+	if hs.shouldDropFrame != nil && hs.shouldDropFrame(f.Header().Type) {
+		return len(b), nil
+	}
+
+	return hs.wrappedConn.Write(b)
+}
+
+func (hs *h2ConnectionSnooper) Close() error { return hs.wrappedConn.Close() }
+
+func (hs *h2ConnectionSnooper) LocalAddr() net.Addr { return hs.wrappedConn.LocalAddr() }
+
+func (hs *h2ConnectionSnooper) RemoteAddr() net.Addr { return hs.wrappedConn.RemoteAddr() }
+
+func (hs *h2ConnectionSnooper) SetDeadline(t time.Time) error { return hs.wrappedConn.SetDeadline(t) }
+
+func (hs *h2ConnectionSnooper) SetReadDeadline(t time.Time) error {
+	return hs.wrappedConn.SetReadDeadline(t)
+}
+
+func (hs *h2ConnectionSnooper) SetWriteDeadline(t time.Time) error {
+	return hs.wrappedConn.SetWriteDeadline(t)
+}