session: encode shared key metadata

frontend/dockerfile/dockerfile_copy_test.go

@@ -1246,6 +1246,50 @@ COPY test+aou.txt /
 	require.Equal(t, "baz", string(dt))
 }
 
+func testLocalUnicodeSharedKey(t *testing.T, sb integration.Sandbox) {
+	f := getFrontend(t, sb)
+	c, err := client.New(sb.Context(), sb.Address())
+	require.NoError(t, err)
+	defer c.Close()
+
+	dockerfile := []byte(integration.UnixOrWindows(
+		`
+FROM scratch
+COPY foo /
+`,
+		`
+FROM nanoserver
+COPY foo /
+`,
+	))
+
+	dir := integration.Tmpdir(
+		t,
+		fstest.CreateFile("Dockerfile", dockerfile, 0600),
+		fstest.CreateFile("foo", []byte("contents"), 0600),
+	)
+	destDir := integration.Tmpdir(t)
+
+	_, err = f.Solve(sb.Context(), c, client.SolveOpt{
+		SharedKey: "context:\u65e9:%2B+plain",
+		Exports: []client.ExportEntry{
+			{
+				Type:      client.ExporterLocal,
+				OutputDir: destDir.Name,
+			},
+		},
+		LocalMounts: map[string]fsutil.FS{
+			dockerui.DefaultLocalNameDockerfile: dir,
+			dockerui.DefaultLocalNameContext:    dir,
+		},
+	}, nil)
+	require.NoError(t, err)
+
+	dt, err := os.ReadFile(filepath.Join(destDir.Name, "foo"))
+	require.NoError(t, err)
+	require.Equal(t, "contents", string(dt))
+}
+
 func testChmodNonOctal(t *testing.T, sb integration.Sandbox) {
 	integration.SkipOnPlatform(t, "windows")
 	f := getFrontend(t, sb)

frontend/dockerfile/dockerfile_test.go

@@ -109,6 +109,7 @@ var allTests = integration.TestFuncs(
 	testCopyWildcards,
 	testCopyRelative,
 	testCopyUnicodePath,
+	testLocalUnicodeSharedKey,
 	testLocalCustomSessionID,
 
 	// dockerfile_core_test.go

session/header.go

@@ -0,0 +1,23 @@
+package session
+
+import "net/url"
+
+func encodeHeaderValue(input string) (string, bool) {
+	for _, r := range input {
+		if r < 0x20 || r > 0x7e {
+			return url.QueryEscape(input), true
+		}
+	}
+	return input, false
+}
+
+func decodeHeaderValue(input string, encoded bool) string {
+	if !encoded {
+		return input
+	}
+	out, err := url.QueryUnescape(input)
+	if err != nil {
+		return input
+	}
+	return out
+}

session/manager.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"net"
 	"net/http"
+	"strconv"
 	"strings"
 	"sync"
 
@@ -105,7 +106,7 @@ func (sm *Manager) handleConn(ctx context.Context, conn net.Conn, opts map[strin
 
 	h := http.Header(opts)
 	id := h.Get(headerSessionID)
-	sharedKey := h.Get(headerSessionSharedKey)
+	sharedKey := decodeHeaderValue(h.Get(headerSessionSharedKey), headerValueIsEncoded(h, headerSessionSharedKeyEncoded))
 
 	ctx, cc, err := grpcClientConn(ctx, conn, opts)
 	if err != nil {
@@ -213,3 +214,12 @@ func canonicalHeaders(in map[string][]string) map[string][]string {
 	}
 	return out
 }
+
+func headerValueIsEncoded(h http.Header, key string) bool {
+	v := h.Get(key)
+	if v == "" {
+		return false
+	}
+	encoded, _ := strconv.ParseBool(v)
+	return encoded
+}

session/session.go

@@ -18,10 +18,11 @@ import (
 )
 
 const (
-	headerSessionID        = "X-Docker-Expose-Session-Uuid"
-	headerSessionName      = "X-Docker-Expose-Session-Name"
-	headerSessionSharedKey = "X-Docker-Expose-Session-Sharedkey"
-	headerSessionMethod    = "X-Docker-Expose-Session-Grpc-Method"
+	headerSessionID               = "X-Docker-Expose-Session-Uuid"
+	headerSessionName             = "X-Docker-Expose-Session-Name"
+	headerSessionSharedKey        = "X-Docker-Expose-Session-Sharedkey"
+	headerSessionSharedKeyEncoded = headerSessionSharedKey + "-Encoded"
+	headerSessionMethod           = "X-Docker-Expose-Session-Grpc-Method"
 )
 
 var propagators = propagation.NewCompositeTextMapPropagator(propagation.TraceContext{}, propagation.Baggage{})
@@ -101,7 +102,11 @@ func (s *Session) Run(ctx context.Context, dialer Dialer) error {
 
 	meta := make(map[string][]string)
 	meta[headerSessionID] = []string{s.id}
-	meta[headerSessionSharedKey] = []string{s.sharedKey}
+	sharedKey, encoded := encodeHeaderValue(s.sharedKey)
+	meta[headerSessionSharedKey] = []string{sharedKey}
+	if encoded {
+		meta[headerSessionSharedKeyEncoded] = []string{"1"}
+	}
 
 	for name, svc := range s.grpcServer.GetServiceInfo() {
 		for _, method := range svc.Methods {

session/session_test.go

@@ -0,0 +1,96 @@
+package session
+
+import (
+	"context"
+	"net"
+	"net/url"
+	"testing"
+
+	"github.com/moby/buildkit/session/testutil"
+	"github.com/pkg/errors"
+	"github.com/stretchr/testify/require"
+	"golang.org/x/sync/errgroup"
+)
+
+func TestSessionSharedKeyMetadata(t *testing.T) {
+	t.Parallel()
+
+	tests := []struct {
+		name      string
+		sharedKey string
+		encoded   bool
+	}{
+		{
+			name:      "ascii",
+			sharedKey: "context:%2B+plain",
+		},
+		{
+			name:      "non-ascii",
+			sharedKey: "context:\u65e9:%2B+plain",
+			encoded:   true,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			t.Parallel()
+
+			s, err := NewSession(t.Context(), tt.sharedKey)
+			require.NoError(t, err)
+
+			errDial := errors.New("stop after metadata capture")
+			var called bool
+			var gotProto string
+			var gotMeta map[string][]string
+			dialer := func(ctx context.Context, proto string, meta map[string][]string) (net.Conn, error) {
+				called = true
+				gotProto = proto
+				gotMeta = meta
+				return nil, errDial
+			}
+
+			err = s.Run(t.Context(), dialer)
+			require.ErrorIs(t, err, errDial)
+			require.True(t, called)
+			require.Equal(t, "h2c", gotProto)
+			require.Equal(t, []string{s.ID()}, gotMeta[headerSessionID])
+			if tt.encoded {
+				require.Equal(t, []string{url.QueryEscape(tt.sharedKey)}, gotMeta[headerSessionSharedKey])
+				require.Equal(t, []string{"1"}, gotMeta[headerSessionSharedKeyEncoded])
+			} else {
+				require.Equal(t, []string{tt.sharedKey}, gotMeta[headerSessionSharedKey])
+				require.NotContains(t, gotMeta, headerSessionSharedKeyEncoded)
+			}
+		})
+	}
+}
+
+func TestSessionSharedKeyRoundTrip(t *testing.T) {
+	t.Parallel()
+
+	sharedKey := "context:\u65e9:%2B+plain"
+	s, err := NewSession(t.Context(), sharedKey)
+	require.NoError(t, err)
+
+	m, err := NewManager()
+	require.NoError(t, err)
+
+	dialer := Dialer(testutil.TestStream(testutil.Handler(m.HandleConn)))
+
+	g, ctx := errgroup.WithContext(t.Context())
+	g.Go(func() error {
+		return s.Run(ctx, dialer)
+	})
+	g.Go(func() error {
+		c, err := m.Get(ctx, s.ID(), false)
+		if err != nil {
+			return err
+		}
+		if c.SharedKey() != sharedKey {
+			return errors.Errorf("expected shared key %q, got %q", sharedKey, c.SharedKey())
+		}
+		return s.Close()
+	})
+
+	require.NoError(t, g.Wait())
+}