From 881945d9fff3f3620f78b7dec838ba7ea4262622 Mon Sep 17 00:00:00 2001
From: "J. Yi" <93548144+jyyi1@users.noreply.github.com>
Date: Tue, 12 Sep 2023 15:42:43 -0400
Subject: [PATCH] feat(network): configurable timeout of packet_listener_proxy
 (#66)

---
 network/packet_listener_proxy.go      | 41 ++++++++++++++++++++-------
 network/packet_listener_proxy_test.go | 37 ++++++++++++++++++++++++
 2 files changed, 68 insertions(+), 10 deletions(-)
 create mode 100644 network/packet_listener_proxy_test.go

diff --git a/network/packet_listener_proxy.go b/network/packet_listener_proxy.go
index 047df23a..446d4f80 100644
--- a/network/packet_listener_proxy.go
+++ b/network/packet_listener_proxy.go
@@ -34,10 +34,10 @@ const packetMaxSize = 2048
 var packetBufferPool = slicepool.MakePool(packetMaxSize)
 
 // Compilation guard against interface implementation
-var _ PacketProxy = (*packetListenerProxyAdapter)(nil)
+var _ PacketProxy = (*PacketListenerProxy)(nil)
 var _ PacketRequestSender = (*packetListenerRequestSender)(nil)
 
-type packetListenerProxyAdapter struct {
+type PacketListenerProxy struct {
 	listener         transport.PacketListener
 	writeIdleTimeout time.Duration
 }
@@ -51,23 +51,44 @@ type packetListenerRequestSender struct {
 	writeIdleTimer   *time.Timer
 }
 
-// NewPacketProxyFromPacketListener creates a new [PacketProxy] that uses the existing [transport.PacketListener]
-// to create connections to a proxy. You can use this function if you already have an implementation of
-// [transport.PacketListener] and would like to inject it into one of the network stacks (for example,
-// network/lwip2transport) as UDP traffic handlers.
-func NewPacketProxyFromPacketListener(pl transport.PacketListener) (PacketProxy, error) {
+// NewPacketProxyFromPacketListener creates a new [PacketProxy] that uses the existing [transport.PacketListener] to
+// create connections to a proxy. You can also specify additional options.
+// This function is useful if you already have an implementation of [transport.PacketListener] and you want to use it
+// with one of the network stacks (for example, network/lwip2transport) as a UDP traffic handler.
+func NewPacketProxyFromPacketListener(pl transport.PacketListener, options ...func(*PacketListenerProxy) error) (*PacketListenerProxy, error) {
 	if pl == nil {
 		return nil, errors.New("pl must not be nil")
 	}
-	return &packetListenerProxyAdapter{
+	p := &PacketListenerProxy{
 		listener:         pl,
 		writeIdleTimeout: 30 * time.Second,
-	}, nil
+	}
+	for _, opt := range options {
+		if err := opt(p); err != nil {
+			return nil, err
+		}
+	}
+	return p, nil
+}
+
+// WithPacketListenerWriteIdleTimeout sets the write idle timeout of the [PacketListenerProxy].
+// This means that if there are no WriteTo operations on the UDP session created by NewSession for the specified amount
+// of time, the proxy will end this session.
+//
+// This should be used together with the [NewPacketProxyFromPacketListenerWithOptions] function.
+func WithPacketListenerWriteIdleTimeout(timeout time.Duration) func(*PacketListenerProxy) error {
+	return func(p *PacketListenerProxy) error {
+		if timeout <= 0 {
+			return errors.New("timeout must be greater than 0")
+		}
+		p.writeIdleTimeout = timeout
+		return nil
+	}
 }
 
 // NewSession implements [PacketProxy].NewSession function. It uses [transport.PacketListener].ListenPacket to create
 // a [net.PacketConn], and constructs a new [PacketRequestSender] that is based on this [net.PacketConn].
-func (proxy *packetListenerProxyAdapter) NewSession(respWriter PacketResponseReceiver) (PacketRequestSender, error) {
+func (proxy *PacketListenerProxy) NewSession(respWriter PacketResponseReceiver) (PacketRequestSender, error) {
 	if respWriter == nil {
 		return nil, errors.New("respWriter must not be nil")
 	}
diff --git a/network/packet_listener_proxy_test.go b/network/packet_listener_proxy_test.go
new file mode 100644
index 00000000..cd7e0d78
--- /dev/null
+++ b/network/packet_listener_proxy_test.go
@@ -0,0 +1,37 @@
+// Copyright 2023 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 network
+
+import (
+	"testing"
+	"time"
+
+	"github.com/Jigsaw-Code/outline-sdk/transport"
+	"github.com/stretchr/testify/require"
+)
+
+func TestWithWriteTimeoutOptionWorks(t *testing.T) {
+	pl := &transport.UDPPacketListener{}
+
+	defProxy, err := NewPacketProxyFromPacketListener(pl)
+	require.NoError(t, err)
+	require.NotNil(t, defProxy)
+	require.Equal(t, 30*time.Second, defProxy.writeIdleTimeout) // default timeout is 30s
+
+	altProxy, err := NewPacketProxyFromPacketListener(pl, WithPacketListenerWriteIdleTimeout(5*time.Minute))
+	require.NoError(t, err)
+	require.NotNil(t, altProxy)
+	require.Equal(t, 5*time.Minute, altProxy.writeIdleTimeout)
+}