diff --git a/go.mod b/go.mod index 1a9e61197..3ce787a97 100644 --- a/go.mod +++ b/go.mod @@ -56,6 +56,7 @@ require ( k8s.io/apimachinery v0.36.1 k8s.io/client-go v0.36.1 k8s.io/metrics v0.36.1 + k8s.io/streaming v0.36.1 k8s.io/utils v0.0.0-20260319190234-28399d86e0b5 sigs.k8s.io/controller-runtime v0.24.1 sigs.k8s.io/yaml v1.6.0 @@ -186,7 +187,6 @@ require ( gotest.tools/v3 v3.5.2 // indirect k8s.io/klog/v2 v2.140.0 // indirect k8s.io/kube-openapi v0.0.0-20260317180543-43fb72c5454a // indirect - k8s.io/streaming v0.36.1 // indirect sigs.k8s.io/json v0.0.0-20250730193827-2d320260d730 // indirect sigs.k8s.io/randfill v1.0.0 // indirect sigs.k8s.io/structured-merge-diff/v6 v6.3.2 // indirect diff --git a/internal/ateclient/builder.go b/internal/ateclient/builder.go index ed72a7599..caab69da6 100644 --- a/internal/ateclient/builder.go +++ b/internal/ateclient/builder.go @@ -19,7 +19,7 @@ import ( "crypto/tls" "crypto/x509" "fmt" - "io" + "net" "net/http" "os" "sync" @@ -40,7 +40,6 @@ import ( "k8s.io/client-go/kubernetes" "k8s.io/client-go/rest" "k8s.io/client-go/tools/clientcmd" - "k8s.io/client-go/tools/portforward" "k8s.io/client-go/transport/spdy" metricsv1beta1 "k8s.io/metrics/pkg/client/clientset/versioned" ) @@ -190,59 +189,22 @@ func dialPortForward(ctx context.Context, kubeconfigPath, k8sContext string, tra return nil, fmt.Errorf("failed to create SPDY transport: %w", err) } - dialer := spdy.NewDialer(upgrader, &http.Client{Transport: transport}, http.MethodPost, req.URL()) - - stopCh := make(chan struct{}) - readyCh := make(chan struct{}) - - ports := []string{"0:443"} // Port 0 asks OS for a random available local port - - fw, err := portforward.New(dialer, ports, stopCh, readyCh, io.Discard, io.Discard) - if err != nil { - return nil, fmt.Errorf("failed to create port forwarder: %w", err) - } - - errCh := make(chan error, 1) - var wg sync.WaitGroup - wg.Add(1) - go func() { - defer wg.Done() - if err := fw.ForwardPorts(); err != nil { - errCh <- fmt.Errorf("port forwarding failed: %w", err) - } - }() - - // Wait for the tunnel to be ready, an error, or context cancellation - select { - case <-readyCh: - // Tunnel is ready! - case err := <-errCh: - return nil, err - case <-ctx.Done(): - return nil, ctx.Err() - } - - forwardedPorts, err := fw.GetPorts() - if err != nil || len(forwardedPorts) == 0 { - close(stopCh) - return nil, fmt.Errorf("failed to get forwarded ports: %w", err) - } - - localPort := forwardedPorts[0].Local - localEndpoint := fmt.Sprintf("127.0.0.1:%d", localPort) + dialer := newSPDYPortForwardDialer(upgrader, &http.Client{Transport: transport}, http.MethodPost, req.URL()) tlsCfg, err := serverTLSConfig(ctx, clientset) if err != nil { - close(stopCh) return nil, err } var opts []grpc.DialOption opts = append(opts, grpc.WithTransportCredentials(credentials.NewTLS(tlsCfg))) opts = append(opts, grpc.WithStatsHandler(otelgrpc.NewClientHandler())) + opts = append(opts, grpc.WithAuthority(apiServerName)) + opts = append(opts, grpc.WithContextDialer(func(ctx context.Context, _ string) (net.Conn, error) { + return dialPodPort(ctx, dialer, 443) + })) tokenOpt, err := bearerTokenDialOption(ctx, clientset) if err != nil { - close(stopCh) return nil, err } opts = append(opts, tokenOpt) @@ -251,9 +213,8 @@ func dialPortForward(ctx context.Context, kubeconfigPath, k8sContext string, tra opts = append(opts, grpc.WithUnaryInterceptor(newTraceInterceptor())) } - conn, err := grpc.NewClient(localEndpoint, opts...) + conn, err := grpc.NewClient("passthrough:///"+apiServerName+":443", opts...) if err != nil { - close(stopCh) return nil, fmt.Errorf("failed to dial gRPC over tunnel: %w", err) } @@ -261,10 +222,7 @@ func dialPortForward(ctx context.Context, kubeconfigPath, k8sContext string, tra ControlClient: ateapipb.NewControlClient(conn), DebugClient: ateapipb.NewDebugClient(conn), conn: conn, - cancel: func() { - close(stopCh) - wg.Wait() - }, + cancel: func() {}, }, nil } diff --git a/internal/ateclient/portforward.go b/internal/ateclient/portforward.go new file mode 100644 index 000000000..75e7e169f --- /dev/null +++ b/internal/ateclient/portforward.go @@ -0,0 +1,197 @@ +// Copyright 2021 The Kubernetes Authors. +// Copyright 2026 Google 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 +// +// http://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 ateclient + +import ( + "context" + "errors" + "fmt" + "io" + "log/slog" + "net" + "net/http" + "net/url" + "strconv" + "sync" + "sync/atomic" + "time" + + corev1 "k8s.io/api/core/v1" + "k8s.io/client-go/tools/portforward" + "k8s.io/client-go/transport/spdy" + "k8s.io/streaming/pkg/httpstream" +) + +type portForwardDialer func(context.Context, ...string) (httpstream.Connection, string, error) + +func newSPDYPortForwardDialer(upgrader spdy.Upgrader, client *http.Client, method string, target *url.URL) portForwardDialer { + return func(ctx context.Context, protocols ...string) (httpstream.Connection, string, error) { + req, err := http.NewRequestWithContext(ctx, method, target.String(), nil) + if err != nil { + return nil, "", fmt.Errorf("creating port-forward request: %w", err) + } + return spdy.NegotiateStreaming(upgrader, client, req, protocols...) + } +} + +// dialPodPort exposes a Kubernetes port-forward data stream directly as a +// net.Conn. This avoids an intermediate localhost TCP listener and its noisy +// connection-reset errors when the gRPC client shuts down. +func dialPodPort(ctx context.Context, dialer portForwardDialer, remotePort int) (conn net.Conn, finalErr error) { + if err := ctx.Err(); err != nil { + return nil, err + } + + streamConn, protocol, err := dialer(ctx, portforward.PortForwardProtocolV1Name) + if err != nil { + return nil, fmt.Errorf("dialing port-forward connection: %w", err) + } + defer func() { + if finalErr != nil { + _ = streamConn.Close() + } + }() + + setupDone := make(chan struct{}) + monitorDone := make(chan struct{}) + var setupFinished atomic.Bool + go func() { + defer close(monitorDone) + select { + case <-ctx.Done(): + if setupFinished.CompareAndSwap(false, true) { + _ = streamConn.Close() + } + case <-setupDone: + } + }() + defer func() { + close(setupDone) + <-monitorDone + }() + + if err := ctx.Err(); err != nil { + return nil, err + } + if protocol != portforward.PortForwardProtocolV1Name { + return nil, fmt.Errorf("port-forward protocol mismatch: server selected %q", protocol) + } + + headers := http.Header{} + headers.Set(corev1.StreamType, corev1.StreamTypeError) + headers.Set(corev1.PortHeader, strconv.Itoa(remotePort)) + headers.Set(corev1.PortForwardRequestIDHeader, "0") + + errorStream, err := streamConn.CreateStream(headers) + if err != nil { + return nil, fmt.Errorf("creating port-forward error stream: %w", err) + } + // The error stream is read-only from the client's perspective. + _ = errorStream.Close() + + headers.Set(corev1.StreamType, corev1.StreamTypeData) + dataStream, err := streamConn.CreateStream(headers) + if err != nil { + return nil, fmt.Errorf("creating port-forward data stream: %w", err) + } + + stream := &portForwardStream{ + Stream: dataStream, + errorStream: errorStream, + streamConn: streamConn, + remotePort: remotePort, + } + + if err := ctx.Err(); err != nil { + return nil, err + } + if !setupFinished.CompareAndSwap(false, true) { + return nil, ctx.Err() + } + go stream.watchErrors() + return stream, nil +} + +type portForwardStream struct { + httpstream.Stream + errorStream httpstream.Stream + streamConn httpstream.Connection + remotePort int + closeOnce sync.Once + closed atomic.Bool + closeErr error +} + +var _ net.Conn = (*portForwardStream)(nil) + +func (s *portForwardStream) watchErrors() { + message, err := io.ReadAll(s.errorStream) + if s.closed.Load() { + return + } + switch { + case err != nil: + slog.Error("Error reading from port-forward error stream", + slog.Int("remotePort", s.remotePort), slog.Any("err", err)) + case len(message) > 0: + slog.Error("Port-forward connection failed", + slog.Int("remotePort", s.remotePort), slog.String("error", string(message))) + default: + return + } + _ = s.Close() +} + +func (s *portForwardStream) Close() error { + s.closeOnce.Do(func() { + s.closed.Store(true) + s.closeErr = errors.Join(s.Stream.Close(), s.streamConn.Close()) + }) + return s.closeErr +} + +func (s *portForwardStream) LocalAddr() net.Addr { + return portForwardAddr("kubernetes-api") +} + +func (s *portForwardStream) RemoteAddr() net.Addr { + return portForwardAddr(fmt.Sprintf("pod:%d", s.remotePort)) +} + +// SPDY streams do not expose deadline controls. gRPC applies RPC deadlines at +// the transport layer, so these methods intentionally match Kubernetes' +// direct port-forward net.Conn adapter and remain no-ops. +func (s *portForwardStream) SetDeadline(time.Time) error { + return nil +} + +func (s *portForwardStream) SetReadDeadline(time.Time) error { + return nil +} + +func (s *portForwardStream) SetWriteDeadline(time.Time) error { + return nil +} + +type portForwardAddr string + +func (a portForwardAddr) Network() string { + return "port-forward" +} + +func (a portForwardAddr) String() string { + return string(a) +} diff --git a/internal/ateclient/portforward_test.go b/internal/ateclient/portforward_test.go new file mode 100644 index 000000000..22aa3b309 --- /dev/null +++ b/internal/ateclient/portforward_test.go @@ -0,0 +1,333 @@ +// Copyright 2026 Google 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 +// +// http://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 ateclient + +import ( + "bytes" + "context" + "errors" + "io" + "net/http" + "net/url" + "strings" + "sync" + "testing" + "time" + + corev1 "k8s.io/api/core/v1" + "k8s.io/client-go/tools/portforward" + "k8s.io/streaming/pkg/httpstream" +) + +type roundTripFunc func(*http.Request) (*http.Response, error) + +func (f roundTripFunc) RoundTrip(req *http.Request) (*http.Response, error) { + return f(req) +} + +type fakePortForwardDialer struct { + conn httpstream.Connection + protocol string + protocols []string + dialCount int +} + +func (d *fakePortForwardDialer) Dial(_ context.Context, protocols ...string) (httpstream.Connection, string, error) { + d.dialCount++ + d.protocols = append(d.protocols, protocols...) + return d.conn, d.protocol, nil +} + +type fakePortForwardConnection struct { + closeCount int + closeOnce sync.Once + closed chan struct{} + headers []http.Header + streams []*fakePortForwardStream + streamReaders []io.Reader + createErrAt int + createErr error + onCreate func(int) +} + +func (c *fakePortForwardConnection) CreateStream(headers http.Header) (httpstream.Stream, error) { + streamNumber := len(c.headers) + 1 + c.headers = append(c.headers, headers.Clone()) + if c.onCreate != nil { + c.onCreate(streamNumber) + } + if c.createErrAt == streamNumber { + return nil, c.createErr + } + stream := &fakePortForwardStream{} + if len(c.streamReaders) >= streamNumber { + stream.reader = c.streamReaders[streamNumber-1] + } + c.streams = append(c.streams, stream) + return stream, nil +} + +func (c *fakePortForwardConnection) Close() error { + c.closeOnce.Do(func() { + c.closeCount++ + if c.closed != nil { + close(c.closed) + } + }) + return nil +} + +func (*fakePortForwardConnection) CloseChan() <-chan bool { + return make(chan bool) +} + +func (*fakePortForwardConnection) SetIdleTimeout(time.Duration) {} + +func (*fakePortForwardConnection) RemoveStreams(...httpstream.Stream) {} + +type fakePortForwardStream struct { + closeCount int + reader io.Reader +} + +func (s *fakePortForwardStream) Read(p []byte) (int, error) { + if s.reader == nil { + return 0, io.EOF + } + return s.reader.Read(p) +} + +func (s *fakePortForwardStream) Write(p []byte) (int, error) { + return len(p), nil +} + +func (s *fakePortForwardStream) Close() error { + s.closeCount++ + return nil +} + +func (*fakePortForwardStream) Reset() error { + return nil +} + +func (*fakePortForwardStream) Headers() http.Header { + return nil +} + +func (*fakePortForwardStream) Identifier() uint32 { + return 0 +} + +func TestDialPodPort(t *testing.T) { + streamConn := &fakePortForwardConnection{} + dialer := &fakePortForwardDialer{ + conn: streamConn, + protocol: portforward.PortForwardProtocolV1Name, + } + + conn, err := dialPodPort(context.Background(), dialer.Dial, 443) + if err != nil { + t.Fatalf("dialPodPort: %v", err) + } + + if dialer.dialCount != 1 { + t.Errorf("dial count = %d, want 1", dialer.dialCount) + } + if len(dialer.protocols) != 1 || dialer.protocols[0] != portforward.PortForwardProtocolV1Name { + t.Errorf("dial protocols = %q, want [%q]", dialer.protocols, portforward.PortForwardProtocolV1Name) + } + if len(streamConn.headers) != 2 { + t.Fatalf("created %d streams, want 2", len(streamConn.headers)) + } + for i, streamType := range []string{corev1.StreamTypeError, corev1.StreamTypeData} { + headers := streamConn.headers[i] + if got := headers.Get(corev1.StreamType); got != streamType { + t.Errorf("stream %d type = %q, want %q", i, got, streamType) + } + if got := headers.Get(corev1.PortHeader); got != "443" { + t.Errorf("stream %d port = %q, want 443", i, got) + } + if got := headers.Get(corev1.PortForwardRequestIDHeader); got != "0" { + t.Errorf("stream %d request ID = %q, want 0", i, got) + } + } + + if got := conn.LocalAddr().Network(); got != "port-forward" { + t.Errorf("local network = %q, want port-forward", got) + } + if got := conn.RemoteAddr().String(); got != "pod:443" { + t.Errorf("remote address = %q, want pod:443", got) + } + + if err := conn.Close(); err != nil { + t.Fatalf("Close: %v", err) + } + if err := conn.Close(); err != nil { + t.Fatalf("second Close: %v", err) + } + if streamConn.closeCount != 1 { + t.Errorf("stream connection close count = %d, want 1", streamConn.closeCount) + } + if streamConn.streams[1].closeCount != 1 { + t.Errorf("data stream close count = %d, want 1", streamConn.streams[1].closeCount) + } +} + +func TestDialPodPortCanceledContext(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + cancel() + dialer := &fakePortForwardDialer{} + + if _, err := dialPodPort(ctx, dialer.Dial, 443); err == nil { + t.Fatal("dialPodPort: got nil error, want context cancellation") + } + if dialer.dialCount != 0 { + t.Errorf("dial count = %d, want 0", dialer.dialCount) + } +} + +func TestSPDYPortForwardDialerHonorsContext(t *testing.T) { + requestStarted := make(chan struct{}) + client := &http.Client{Transport: roundTripFunc(func(req *http.Request) (*http.Response, error) { + close(requestStarted) + <-req.Context().Done() + return nil, req.Context().Err() + })} + + target := &url.URL{ + Scheme: "https", + Host: "example.test", + } + dialer := newSPDYPortForwardDialer(nil, client, http.MethodPost, target) + ctx, cancel := context.WithCancel(context.Background()) + go func() { + <-requestStarted + cancel() + }() + + if _, _, err := dialer(ctx, portforward.PortForwardProtocolV1Name); err == nil || !strings.Contains(err.Error(), context.Canceled.Error()) { + t.Fatalf("dial error = %v, want context cancellation", err) + } +} + +func TestDialPodPortClosesConnectionOnSetupFailure(t *testing.T) { + setupError := errors.New("setup failed") + tests := []struct { + name string + dialer *fakePortForwardDialer + streamConn *fakePortForwardConnection + wantCreated int + }{ + { + name: "protocol mismatch", + dialer: &fakePortForwardDialer{ + protocol: "unexpected.example", + }, + streamConn: &fakePortForwardConnection{}, + }, + { + name: "error stream", + dialer: &fakePortForwardDialer{ + protocol: portforward.PortForwardProtocolV1Name, + }, + streamConn: &fakePortForwardConnection{ + createErrAt: 1, + createErr: setupError, + }, + wantCreated: 1, + }, + { + name: "data stream", + dialer: &fakePortForwardDialer{ + protocol: portforward.PortForwardProtocolV1Name, + }, + streamConn: &fakePortForwardConnection{ + createErrAt: 2, + createErr: setupError, + }, + wantCreated: 2, + }, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + test.dialer.conn = test.streamConn + if _, err := dialPodPort(context.Background(), test.dialer.Dial, 443); err == nil { + t.Fatal("dialPodPort: got nil error, want setup failure") + } + if test.streamConn.closeCount != 1 { + t.Errorf("stream connection close count = %d, want 1", test.streamConn.closeCount) + } + if len(test.streamConn.headers) != test.wantCreated { + t.Errorf("created stream count = %d, want %d", len(test.streamConn.headers), test.wantCreated) + } + }) + } +} + +func TestDialPodPortClosesConnectionWhenContextCanceledDuringSetup(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + closed := make(chan struct{}) + streamConn := &fakePortForwardConnection{ + closed: closed, + onCreate: func(streamNumber int) { + if streamNumber == 2 { + cancel() + <-closed + } + }, + } + dialer := &fakePortForwardDialer{ + conn: streamConn, + protocol: portforward.PortForwardProtocolV1Name, + } + + if _, err := dialPodPort(ctx, dialer.Dial, 443); !errors.Is(err, context.Canceled) { + t.Fatalf("dialPodPort error = %v, want context cancellation", err) + } + if streamConn.closeCount != 1 { + t.Errorf("stream connection close count = %d, want 1", streamConn.closeCount) + } +} + +func TestPortForwardErrorClosesConnection(t *testing.T) { + closed := make(chan struct{}) + streamConn := &fakePortForwardConnection{ + closed: closed, + streamReaders: []io.Reader{ + bytes.NewBufferString("remote port unavailable"), + nil, + }, + } + dialer := &fakePortForwardDialer{ + conn: streamConn, + protocol: portforward.PortForwardProtocolV1Name, + } + + conn, err := dialPodPort(context.Background(), dialer.Dial, 443) + if err != nil { + t.Fatalf("dialPodPort: %v", err) + } + defer conn.Close() + + select { + case <-closed: + case <-time.After(time.Second): + t.Fatal("timed out waiting for port-forward error to close connection") + } + if streamConn.streams[1].closeCount != 1 { + t.Errorf("data stream close count = %d, want 1", streamConn.streams[1].closeCount) + } +}