feat: stop all active connections on Proxy.Stop()

x/httpproxy/connect_handler.go

@@ -95,22 +95,24 @@ func (h *connectHandler) ServeHTTP(proxyResp http.ResponseWriter, proxyReq *http
 		return
 	}
 
-	httpConn, _, err := hijacker.Hijack()
+	httpConn, clientRW, err := hijacker.Hijack()
 	if err != nil {
 		http.Error(proxyResp, "Failed to hijack connection", http.StatusInternalServerError)
 		return
 	}
 	defer httpConn.Close()
 
 	// Inform the client that the connection has been established.
-	httpConn.Write([]byte("HTTP/1.1 200 Connection established\r\n\r\n"))
+	clientRW.Write([]byte("HTTP/1.1 200 Connection established\r\n\r\n"))
+	clientRW.Flush()
 
 	// Relay data between client and target in both directions.
 	go func() {
-		io.Copy(targetConn, httpConn)
+		io.Copy(targetConn, clientRW)
 		targetConn.CloseWrite()
 	}()
-	io.Copy(httpConn, targetConn)
+	io.Copy(clientRW, targetConn)
+	clientRW.Flush()
 	// httpConn is closed by the defer httpConn.Close() above.
 }
 

x/httpproxy/connect_handler_test.go

@@ -0,0 +1,40 @@
+// Copyright 2024 Jigsaw Operations LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     https://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package httpproxy
+
+import (
+	"context"
+	"errors"
+	"net/http/httptest"
+	"testing"
+
+	"github.com/Jigsaw-Code/outline-sdk/transport"
+	"github.com/stretchr/testify/require"
+)
+
+func TestNewConnectHandler(t *testing.T) {
+	h := NewConnectHandler(transport.FuncStreamDialer(func(ctx context.Context, addr string) (transport.StreamConn, error) {
+		return nil, errors.New("not implemented")
+	}))
+
+	ch, ok := h.(*connectHandler)
+	require.True(t, ok)
+	require.NotNil(t, ch.dialer)
+
+	req := httptest.NewRequest("CONNECT", "example.invalid:0", nil)
+	resp := httptest.NewRecorder()
+	h.ServeHTTP(resp, req)
+	require.Equal(t, 503, resp.Result().StatusCode)
+}

x/mobileproxy/mobileproxy.go

@@ -20,6 +20,7 @@ package mobileproxy
 
 import (
 	"context"
+	"errors"
 	"fmt"
 	"io"
 	"log"
@@ -92,6 +93,7 @@ func (p *Proxy) Stop(timeoutSeconds int) {
 		log.Fatalf("Failed to shutdown gracefully: %v", err)
 		p.server.Close()
 	}
+	p.proxyHandler = nil
 }
 
 // RunProxy runs a local web proxy that listens on localAddress, and handles proxy requests by
@@ -101,10 +103,16 @@ func RunProxy(localAddress string, dialer *StreamDialer) (*Proxy, error) {
 	if err != nil {
 		return nil, fmt.Errorf("could not listen on address %v: %v", localAddress, err)
 	}
+	if dialer == nil {
+		return nil, errors.New("dialer must not be nil. Please create and pass a valid StreamDialer")
+	}
 
-	proxyHandler := httpproxy.NewProxyHandler(dialer)
+	stopDialer := &stopStreamDialer{Dialer: dialer}
+	proxyHandler := httpproxy.NewProxyHandler(stopDialer)
 	proxyHandler.FallbackHandler = http.NotFoundHandler()
 	server := &http.Server{Handler: proxyHandler}
+	// Make sure all connections are stopped so they don't linger around.
+	server.RegisterOnShutdown(stopDialer.StopConnections)
 	go server.Serve(listener)
 
 	host, portStr, err := net.SplitHostPort(listener.Addr().String())

x/mobileproxy/stop_dialer.go

@@ -0,0 +1,143 @@
+// Copyright 2024 Jigsaw Operations LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     https://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+// Package mobileproxy provides convenience utilities to help applications run a local proxy
+// and use that to configure their networking libraries.
+//
+// This package is suitable for use with Go Mobile, making it a convenient way to integrate with mobile apps.
+package mobileproxy
+
+import (
+	"container/list"
+	"context"
+	"errors"
+	"io"
+	"sync"
+
+	"github.com/Jigsaw-Code/outline-sdk/transport"
+)
+
+var errStopped = errors.New("dialer stopped")
+
+// stopConn is a [transport.StreamConn] that can be stopped and has a close handler.
+type stopConn struct {
+	transport.StreamConn
+	onceClose sync.Once
+	closeErr  error
+	// OnClose is called when Close is called. Close must be called only once.
+	OnClose func()
+}
+
+// WriteTo implements [io.WriterTo], to make sure we don't hide the functionality from the
+// underlying connection.
+func (c *stopConn) WriteTo(w io.Writer) (int64, error) {
+	return io.Copy(w, c.StreamConn)
+}
+
+// ReadFrom implements [io.ReaderFrom], to make sure we don't hide the functionality from
+// the underlying connection.
+func (c *stopConn) ReadFrom(r io.Reader) (int64, error) {
+	// Prefer ReadFrom if requested. Otherwise io.Copy prefers WriteTo.
+	if rf, ok := c.StreamConn.(io.ReaderFrom); ok {
+		return rf.ReadFrom(r)
+	}
+	return io.Copy(c.StreamConn, r)
+}
+
+// Close implements [transport.StreamConn].
+func (c *stopConn) Close() error {
+	c.Stop()
+	c.OnClose()
+	return c.closeErr
+}
+
+// Stop stops the connection, calling its underlying Close if it wasn't called yet, without calling OnClose.
+func (c *stopConn) Stop() {
+	c.onceClose.Do(func() {
+		c.closeErr = c.StreamConn.Close()
+	})
+}
+
+// stopStreamDialer is a [transport.StreamDialer] that provides a method to close all open connections.
+type stopStreamDialer struct {
+	Dialer       transport.StreamDialer
+	stopped      bool
+	cleanupFuncs list.List
+	mu           sync.Mutex
+}
+
+var _ transport.StreamDialer = (*stopStreamDialer)(nil)
+
+// DialStream implements [transport.StreamDialer].
+func (d *stopStreamDialer) DialStream(ctx context.Context, addr string) (transport.StreamConn, error) {
+	d.mu.Lock()
+	defer d.mu.Unlock()
+
+	// Capture any stop that happened before DialStream.
+	if d.stopped {
+		return nil, errStopped
+	}
+
+	// Register dial cancelation to capture a stop during the DialStream.
+	ctx, cancel := context.WithCancel(ctx)
+	defer func() {
+		if !errors.Is(ctx.Err(), context.Canceled) {
+			cancel()
+		}
+	}()
+	cleanDialEl := d.cleanupFuncs.PushBack(func() {
+		cancel()
+	})
+
+	// We release the lock for the dial because it's a slow operation and to allow for parallel dials.
+	d.mu.Unlock()
+	conn, err := d.Dialer.DialStream(ctx, addr)
+	d.mu.Lock()
+
+	// Clean up dial cancelation.
+	d.cleanupFuncs.Remove(cleanDialEl)
+
+	if err != nil {
+		return nil, err
+	}
+
+	// Dialer may not pay attention to context cancellation, so we check if we are stopped here.
+	if d.stopped {
+		conn.Close()
+		return nil, errStopped
+	}
+
+	// We have a connection. Register cleanup to capture stop during the connection lifetime.
+	e := d.cleanupFuncs.PushBack(nil)
+	sConn := &stopConn{
+		StreamConn: conn,
+		OnClose: func() {
+			d.mu.Lock()
+			d.cleanupFuncs.Remove(e)
+			d.mu.Unlock()
+		},
+	}
+	e.Value = sConn.Stop
+	return sConn, nil
+}
+
+// StopConnections stops all active connections created by this [stopStreamDialer].
+func (d *stopStreamDialer) StopConnections() {
+	d.mu.Lock()
+	defer d.mu.Unlock()
+	d.stopped = true
+	for e := d.cleanupFuncs.Front(); e != nil; e = e.Next() {
+		e.Value.(func())()
+	}
+}

x/mobileproxy/stop_dialer_test.go

@@ -0,0 +1,139 @@
+// Copyright 2024 Jigsaw Operations LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     https://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+// Package mobileproxy provides convenience utilities to help applications run a local proxy
+// and use that to configure their networking libraries.
+//
+// This package is suitable for use with Go Mobile, making it a convenient way to integrate with mobile apps.
+package mobileproxy
+
+import (
+	"context"
+	"io"
+	"net"
+	"sync"
+	"testing"
+
+	"github.com/Jigsaw-Code/outline-sdk/transport"
+	"github.com/stretchr/testify/require"
+)
+
+func Test_stopStreamDialer_StopBeforeDial(t *testing.T) {
+	d := &stopStreamDialer{
+		Dialer: transport.FuncStreamDialer(func(ctx context.Context, addr string) (transport.StreamConn, error) {
+			return nil, nil
+		}),
+	}
+	d.StopConnections()
+	conn, err := d.DialStream(context.Background(), "invalid:0")
+	require.Nil(t, conn)
+	require.ErrorIs(t, err, errStopped)
+}
+
+func Test_stopStreamDialer_StopDuringDial(t *testing.T) {
+	dialStarted := make(chan struct{})
+	resumeDial := make(chan struct{})
+	d := &stopStreamDialer{
+		Dialer: transport.FuncStreamDialer(func(ctx context.Context, addr string) (transport.StreamConn, error) {
+			close(dialStarted)
+			<-resumeDial
+			return &net.TCPConn{}, nil
+		}),
+	}
+	go func() {
+		<-dialStarted
+		d.StopConnections()
+		close(resumeDial)
+	}()
+	conn, err := d.DialStream(context.Background(), "invalid:0")
+	require.Nil(t, conn)
+	require.ErrorIs(t, err, errStopped)
+}
+
+// netConnAdaptor converts a net.Conn to transport.StreamConn.
+// TODO(fortuna): Consider moving it to the SDK.
+type netConnAdaptor struct {
+	net.Conn
+	onceClose sync.Once
+	closeErr  error
+}
+
+var _ transport.StreamConn = (*netConnAdaptor)(nil)
+
+func (c *netConnAdaptor) CloseWrite() error {
+	type WriteCloser interface {
+		CloseWrite() error
+	}
+	if wc, ok := c.Conn.(WriteCloser); ok {
+		return wc.CloseWrite()
+	}
+	return c.Close()
+}
+
+func (c *netConnAdaptor) CloseRead() error {
+	type ReadCloser interface {
+		CloseRead() error
+	}
+	if rc, ok := c.Conn.(ReadCloser); ok {
+		return rc.CloseRead()
+	}
+	return nil
+}
+
+func (c *netConnAdaptor) Close() error {
+	c.onceClose.Do(func() {
+		c.closeErr = c.Conn.Close()
+	})
+	return c.closeErr
+}
+
+func asStreamConn(conn net.Conn) transport.StreamConn {
+	if sc, ok := conn.(transport.StreamConn); ok {
+		return sc
+	}
+	return &netConnAdaptor{Conn: conn}
+}
+
+func Test_stopStreamDialer_StopAfterDial(t *testing.T) {
+	conn1, conn2 := net.Pipe()
+	defer conn2.Close()
+	d := &stopStreamDialer{
+		Dialer: transport.FuncStreamDialer(func(ctx context.Context, addr string) (transport.StreamConn, error) {
+			return asStreamConn(conn1), nil
+		}),
+	}
+	conn, err := d.DialStream(context.Background(), "invalid:0")
+	require.NotNil(t, conn)
+	require.NoError(t, err)
+	defer conn.Close()
+	d.StopConnections()
+	_, err = conn.Read([]byte{})
+	require.ErrorIs(t, err, io.ErrClosedPipe)
+}
+
+func Test_stopStreamDialer_StopAfterClose(t *testing.T) {
+	conn1, conn2 := net.Pipe()
+	defer conn2.Close()
+	d := &stopStreamDialer{
+		Dialer: transport.FuncStreamDialer(func(ctx context.Context, addr string) (transport.StreamConn, error) {
+			return asStreamConn(conn1), nil
+		}),
+	}
+	conn, err := d.DialStream(context.Background(), "invalid:0")
+	require.NotNil(t, conn)
+	require.NoError(t, err)
+	err = conn.Close()
+	require.NoError(t, err)
+	d.StopConnections()
+}