use client specified flush policy if set over bd registrar provided f…

…lush policy
refraction-networking · Oct 27, 2023 · 1929d00 · 1929d00
1 parent d9d5826
commit 1929d00
Show file tree

Hide file tree

Showing 2 changed files with 82 additions and 0 deletions.
diff --git a/pkg/transports/wrapping/prefix/client.go b/pkg/transports/wrapping/prefix/client.go
@@ -213,6 +213,14 @@ func (t *ClientTransport) SetSessionParams(incoming *anypb.Any, unchecked ...boo
 		return fmt.Errorf("%w, nil params", ErrBadParams)
 	}
 
+	// If the client set a custom flush policy, use it over whatever the bidirectional registrar
+	// is trying to set.
+	if t.parameters.CustomFlushPolicy != nil {
+		if t.parameters.GetCustomFlushPolicy() != DefaultFlush {
+			prefixParams.CustomFlushPolicy = t.parameters.CustomFlushPolicy
+		}
+	}
+
 	if len(unchecked) != 0 && unchecked[0] {
 		// Overwrite the prefix bytes and type without checking the default set. This is used for
 		// RegResponse where the registrar may override the chosen prefix with a prefix outside of

diff --git a/pkg/transports/wrapping/prefix/prefix_test.go b/pkg/transports/wrapping/prefix/prefix_test.go
@@ -14,6 +14,8 @@ import (
 
 	"github.com/stretchr/testify/require"
 	"golang.org/x/crypto/curve25519"
+	"google.golang.org/protobuf/proto"
+	"google.golang.org/protobuf/types/known/anypb"
 
 	"github.com/refraction-networking/conjure/internal/conjurepath"
 	tests "github.com/refraction-networking/conjure/internal/testutils"
@@ -498,3 +500,75 @@ func TestPrefixEndToEnd(t *testing.T) {
 		}
 	}
 }
+
+func TestPrefixClientSetSessionParams(t *testing.T) {
+
+	ct := &ClientTransport{Prefix: DefaultPrefixes[0], parameters: nil}
+	err := ct.Prepare(context.Background(), nil)
+	require.Nil(t, err)
+	pp := defaultParams()
+	pp.PrefixId = proto.Int32(int32(OpenSSH2))
+	pp.CustomFlushPolicy = proto.Int32(int32(NoAddedFlush))
+	require.False(t, pp.GetRandomizeDstPort())
+
+	app, err := anypb.New(pp)
+	require.Nil(t, err)
+	err = ct.SetSessionParams(app)
+	require.Nil(t, err)
+	// Set session params should overwrite the prefix and port randomization parameters
+	// since the fliush policy was unset (DefaultFlush), the session flush policy should be set to
+	// the flush policy indicated by the prefix.
+	require.Equal(t, OpenSSH2, PrefixID(ct.sessionParams.GetPrefixId()))
+	require.False(t, ct.sessionParams.GetRandomizeDstPort())
+	require.Equal(t, NoAddedFlush, ct.sessionParams.GetCustomFlushPolicy())
+
+	// the dialer client parameters should remain unchanged
+	require.Equal(t, DefaultPrefixes[0].ID(), PrefixID(ct.parameters.GetPrefixId()))
+
+	// =================================== //
+
+	ct = &ClientTransport{Prefix: DefaultPrefixes[0], parameters: &pb.PrefixTransportParams{CustomFlushPolicy: proto.Int32(int32(FlushAfterPrefix))}}
+	err = ct.Prepare(context.Background(), nil)
+	require.Nil(t, err)
+	pp = defaultParams()
+	pp.PrefixId = proto.Int32(int32(OpenSSH2))
+	require.False(t, pp.GetRandomizeDstPort())
+
+	app, err = anypb.New(pp)
+	require.Nil(t, err)
+	err = ct.SetSessionParams(app)
+	require.Nil(t, err)
+	// Set session params should overwrite the prefix and port randomization parameters
+	// since the flush policy was SET (FlushAfterPrefix), the session flush policy should be set to
+	// the flush policy indicated by the client params.
+	require.Equal(t, OpenSSH2, PrefixID(ct.sessionParams.GetPrefixId()))
+	require.False(t, ct.sessionParams.GetRandomizeDstPort())
+	require.Equal(t, FlushAfterPrefix, ct.sessionParams.GetCustomFlushPolicy())
+
+	// the dialer client parameters should remain unchanged
+	require.Equal(t, DefaultPrefixes[0].ID(), PrefixID(ct.parameters.GetPrefixId()))
+
+	// =================================== //
+
+	ct = &ClientTransport{Prefix: DefaultPrefixes[0], parameters: &pb.PrefixTransportParams{CustomFlushPolicy: proto.Int32(int32(FlushAfterPrefix))}}
+	err = ct.Prepare(context.Background(), nil)
+	require.Nil(t, err)
+	pp = defaultParams()
+	pp.PrefixId = proto.Int32(int32(OpenSSH2))
+	require.False(t, pp.GetRandomizeDstPort())
+
+	app, err = anypb.New(pp)
+	require.Nil(t, err)
+	err = ct.SetSessionParams(app, true)
+	require.Nil(t, err)
+	// Set session params should overwrite the prefix and port randomization parameters
+	// since the flush policy was SET (FlushAfterPrefix), the session flush policy should be set to
+	// the flush policy indicated by the client params even though we are setting with the
+	// unchecked flag enabled.
+	require.Equal(t, OpenSSH2, PrefixID(ct.sessionParams.GetPrefixId()))
+	require.False(t, ct.sessionParams.GetRandomizeDstPort())
+	require.Equal(t, FlushAfterPrefix, ct.sessionParams.GetCustomFlushPolicy())
+
+	// the dialer client parameters should remain unchanged
+	require.Equal(t, DefaultPrefixes[0].ID(), PrefixID(ct.parameters.GetPrefixId()))
+}