support gRPC proxy to remote backends

internal/machine/api/proxy/director.go

@@ -3,35 +3,123 @@ package proxy
 import (
 	"context"
 	"github.com/siderolabs/grpc-proxy/proxy"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/metadata"
+	"google.golang.org/grpc/status"
+	"net"
+	"strconv"
 	"sync"
 )
 
 // Director manages routing of gRPC requests between local and remote backends.
 type Director struct {
-	localTarget    string
+	localBackend   *LocalBackend
-	localBackend   proxy.Backend
+	remotePort     int
 	remoteBackends sync.Map
-
+	// mu synchronizes access to localAddress.
 	mu           sync.RWMutex
+	localAddress string
 }
 
-func NewDirector(localSockPath string) *Director {
+func NewDirector(localSockPath string, remotePort int) *Director {
 	return &Director{
 		localBackend: NewLocalBackend(localSockPath),
+		remotePort:   remotePort,
 	}
 }
 
 // UpdateLocalAddress updates the local machine address used to identify which requests should be proxied
 // to the local gRPC server.
-func (d *Director) UpdateLocalAddress(target string) {
+func (d *Director) UpdateLocalAddress(addr string) {
 	d.mu.Lock()
 	defer d.mu.Unlock()
 
-	d.localTarget = target
+	d.localAddress = addr
 }
 
 // Director implements proxy.StreamDirector for grpc-proxy, routing requests to local or remote backends based
-// on gRPC metadata in the context.
+// on gRPC metadata in the context. Each machine metadata is injected into the response messages by the proxy
+// if the request is proxied to multiple backends.
 func (d *Director) Director(ctx context.Context, fullMethodName string) (proxy.Mode, []proxy.Backend, error) {
+	md, ok := metadata.FromIncomingContext(ctx)
+	if !ok {
 		return proxy.One2One, []proxy.Backend{d.localBackend}, nil
+	}
+	// If the request is already proxied, send it to the local backend.
+	if _, ok = md["proxy-authority"]; ok {
+		return proxy.One2One, []proxy.Backend{d.localBackend}, nil
+	}
+	// If the request metadata doesn't contain machines to proxy to, send it to the local backend.
+	machines, ok := md["machines"]
+	if !ok {
+		return proxy.One2One, []proxy.Backend{d.localBackend}, nil
+	}
+	if len(machines) == 0 {
+		return proxy.One2One, nil, status.Error(codes.InvalidArgument, "no machines specified")
+	}
+
+	d.mu.RLock()
+	localAddress := d.localAddress
+	d.mu.RUnlock()
+
+	backends := make([]proxy.Backend, len(machines))
+	for i, addr := range machines {
+		if addr == localAddress {
+			backends[i] = d.localBackend
+			continue
+		}
+
+		target := net.JoinHostPort(addr, strconv.Itoa(d.remotePort))
+		backend, err := d.remoteBackend(target)
+		if err != nil {
+			return proxy.One2One, nil, status.Error(codes.Internal, err.Error())
+		}
+		backends[i] = backend
+	}
+
+	if len(backends) == 1 {
+		return proxy.One2One, backends, nil
+	}
+	return proxy.One2Many, backends, nil
+}
+
+// remoteBackend returns a RemoteBackend for the given target from the cache or creates a new one.
+func (d *Director) remoteBackend(target string) (*RemoteBackend, error) {
+	b, ok := d.remoteBackends.Load(target)
+	if ok {
+		return b.(*RemoteBackend), nil
+	}
+
+	backend, err := NewRemoteBackend(target)
+	if err != nil {
+		return nil, err
+	}
+	existing, loaded := d.remoteBackends.LoadOrStore(target, backend)
+	if loaded {
+		// A concurrent remoteBackend call built a different backend.
+		backend.Close()
+		return existing.(*RemoteBackend), nil
+	}
+
+	return backend, nil
+}
+
+// FlushRemoteBackends closes all remote backend connections and removes them from the cache.
+func (d *Director) FlushRemoteBackends() {
+	d.remoteBackends.Range(func(key, value interface{}) bool {
+		backend, ok := value.(*RemoteBackend)
+		if !ok {
+			return true
+		}
+
+		backend.Close()
+		d.remoteBackends.Delete(key)
+		return true
+	})
+}
+
+// Close closes all backend connections.
+func (d *Director) Close() {
+	d.localBackend.Close()
+	d.FlushRemoteBackends()
 }

internal/machine/api/proxy/local.go

@@ -2,6 +2,7 @@ package proxy
 
 import (
 	"context"
+	"fmt"
 	"github.com/siderolabs/grpc-proxy/proxy"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/credentials/insecure"
@@ -26,43 +27,56 @@ func NewLocalBackend(sockPath string) *LocalBackend {
 	}
 }
 
-func (l *LocalBackend) String() string {
+func (b *LocalBackend) String() string {
 	return "local"
 }
 
 // GetConnection returns a gRPC connection to the local server listening on the Unix socket.
-func (l *LocalBackend) GetConnection(ctx context.Context, _ string) (context.Context, *grpc.ClientConn, error) {
+func (b *LocalBackend) GetConnection(ctx context.Context, _ string) (context.Context, *grpc.ClientConn, error) {
 	md, _ := metadata.FromIncomingContext(ctx)
+	// TODO: delete
+	fmt.Printf("### local backend metadata: %+v\n", md)
 	outCtx := metadata.NewOutgoingContext(ctx, md)
 
-	l.mu.RLock()
+	b.mu.RLock()
-	if l.conn != nil {
+	if b.conn != nil {
-		l.mu.RUnlock()
+		defer b.mu.RUnlock()
-		return outCtx, l.conn, nil
+		return outCtx, b.conn, nil
 	}
-	l.mu.RUnlock()
+	b.mu.RUnlock()
 
-	l.mu.Lock()
+	b.mu.Lock()
-	defer l.mu.Unlock()
+	defer b.mu.Unlock()
 
 	var err error
-	l.conn, err = grpc.NewClient(
+	b.conn, err = grpc.NewClient(
-		"unix://"+l.sockPath,
+		"unix://"+b.sockPath,
 		grpc.WithTransportCredentials(insecure.NewCredentials()),
 		grpc.WithDefaultCallOptions(
 			grpc.ForceCodecV2(proxy.Codec()),
 		),
 	)
 
-	return outCtx, l.conn, err
+	return outCtx, b.conn, err
 }
 
 // AppendInfo is called to enhance response from the backend with additional data.
-func (l *LocalBackend) AppendInfo(_ bool, resp []byte) ([]byte, error) {
+func (b *LocalBackend) AppendInfo(_ bool, resp []byte) ([]byte, error) {
 	return resp, nil
 }
 
 // BuildError is called to convert error from upstream into response field.
-func (l *LocalBackend) BuildError(bool, error) ([]byte, error) {
+func (b *LocalBackend) BuildError(bool, error) ([]byte, error) {
 	return nil, nil
 }
+
+// Close closes the upstream gRPC connection.
+func (b *LocalBackend) Close() {
+	b.mu.Lock()
+	defer b.mu.Unlock()
+
+	if b.conn != nil {
+		b.conn.Close()
+		b.conn = nil
+	}
+}

internal/machine/api/proxy/remote.go

@@ -0,0 +1,241 @@
+package proxy
+
+import (
+	"context"
+	"fmt"
+	"github.com/siderolabs/grpc-proxy/proxy"
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/backoff"
+	"google.golang.org/grpc/credentials/insecure"
+	"google.golang.org/grpc/metadata"
+	"google.golang.org/grpc/status"
+	"google.golang.org/protobuf/encoding/protowire"
+	"google.golang.org/protobuf/proto"
+	"net"
+	"net/netip"
+	"sync"
+	"time"
+	"uncloud/internal/machine/api/pb"
+)
+
+// RemoteBackend is a proxy.Backend implementation that proxies to a remote gRPC server, injecting machine metadata
+// into the response.
+//
+// Based on the Talos apid implementation:
+// https://github.com/siderolabs/talos/blob/59a78da42cdea8fbccc35d0851f9b0eef928261b/internal/app/apid/pkg/backend/apid.go
+type RemoteBackend struct {
+	target string
+
+	mu   sync.RWMutex
+	conn *grpc.ClientConn
+}
+
+var _ proxy.Backend = (*RemoteBackend)(nil)
+
+// NewRemoteBackend creates a new instance of RemoteBackend for the given target which must have the format [IPv6]:port.
+func NewRemoteBackend(target string) (*RemoteBackend, error) {
+	host, _, err := net.SplitHostPort(target)
+	if err != nil {
+		return nil, fmt.Errorf("target must have the format [IPv6]:port: %s", target)
+	}
+	addr, err := netip.ParseAddr(host)
+	if err != nil || !addr.Is6() {
+		return nil, fmt.Errorf("target host must be a valid IPv6 address: %s", host)
+	}
+
+	return &RemoteBackend{target: target}, nil
+}
+
+func (b *RemoteBackend) String() string {
+	return b.target
+}
+
+// GetConnection returns a gRPC connection to the remote server.
+func (b *RemoteBackend) GetConnection(ctx context.Context, _ string) (context.Context, *grpc.ClientConn, error) {
+	md, _ := metadata.FromIncomingContext(ctx)
+	if authority := md[":authority"]; len(authority) > 0 {
+		md.Set("proxy-authority", authority...)
+	} else {
+		md.Set("proxy-authority", "unknown")
+	}
+	delete(md, ":authority")
+	delete(md, "machines")
+
+	outCtx := metadata.NewOutgoingContext(ctx, md)
+
+	b.mu.RLock()
+	if b.conn != nil {
+		defer b.mu.RUnlock()
+		return outCtx, b.conn, nil
+	}
+	b.mu.RUnlock()
+
+	b.mu.Lock()
+	defer b.mu.Unlock()
+
+	// Override the max delay to avoid excessive backoff when the another node is unavailable, e.g. rebooted
+	// or the WireGuard connection is temporary down.
+	//
+	// Default max delay is 2 minutes, which is too long for our use case.
+	backoffConfig := backoff.DefaultConfig
+	// The maximum wait time between attempts.
+	backoffConfig.MaxDelay = 15 * time.Second
+
+	var err error
+	b.conn, err = grpc.NewClient(
+		b.target,
+		grpc.WithTransportCredentials(insecure.NewCredentials()),
+		grpc.WithConnectParams(grpc.ConnectParams{
+			Backoff: backoffConfig,
+			// Not published as a constant in gRPC library.
+			// See: https://github.com/grpc/grpc-go/blob/d5dee5fdbdeb52f6ea10b37b2cc7ce37814642d7/clientconn.go#L55-L56
+			// Each connection attempt can take up to MinConnectTimeout.
+			MinConnectTimeout: 20 * time.Second,
+		}),
+		grpc.WithDefaultCallOptions(
+			grpc.ForceCodecV2(proxy.Codec()),
+		),
+	)
+
+	return outCtx, b.conn, err
+}
+
+// AppendInfo is called to enhance response from the backend with additional data.
+//
+// AppendInfo enhances upstream response with machine metadata (target).
+//
+// This method depends on grpc protobuf response structure, each response should
+// look like:
+//
+//	message SomeResponse {
+//	  repeated SomeReply messages = 1; // please note field ID == 1
+//	}
+//
+//	message SomeReply {
+//	  common.Metadata metadata = 1;
+//	  <other fields go here ...>
+//	}
+//
+// As 'SomeReply' is repeated in 'SomeResponse', if we concatenate protobuf representation
+// of several 'SomeResponse' messages, we still get valid 'SomeResponse' representation but with more
+// entries (feature of protobuf binary representation).
+//
+// If we look at binary representation of any unary 'SomeResponse' message, it will always contain one
+// protobuf field with field ID 1 (see above) and type 2 (embedded message SomeReply is encoded
+// as string with length). So if we want to add fields to 'SomeReply', we can simply read field
+// header, adjust length for new 'SomeReply' representation, and prepend new field header.
+//
+// At the same time, we can add 'common.Metadata' structure to 'SomeReply' by simply
+// appending or prepending 'common.Metadata' as a single field. This requires 'metadata'
+// field to be not defined in original response. (This is due to the fact that protobuf message
+// representation is concatenation of each field representation).
+//
+// To build only single field (Metadata) we use helper message which contains exactly this
+// field with same field ID as in every other 'SomeReply':
+//
+//	message Empty {
+//	  common.Metadata metadata = 1;
+//	}
+//
+// As streaming replies are not wrapped into 'SomeResponse' with 'repeated', handling is simpler: we just
+// need to append Empty with details.
+//
+// So AppendInfo does the following: validates that response contains field ID 1 encoded as string,
+// cuts field header, rest is representation of some reply. Marshal 'Empty' as protobuf,
+// which builds 'common.Metadata' field, append it to original response message, build new header
+// for new length of some response, and add back new field header.
+func (b *RemoteBackend) AppendInfo(streaming bool, resp []byte) ([]byte, error) {
+	payload, err := proto.Marshal(&pb.Empty{
+		Metadata: &pb.Metadata{
+			Machine: b.target,
+		},
+	})
+
+	if streaming {
+		return append(resp, payload...), err
+	}
+
+	const (
+		metadataField = 1 // field number in proto definition for repeated response
+		metadataType  = 2 // "string" for embedded messages
+	)
+
+	// decode protobuf embedded header
+
+	typ, n1 := protowire.ConsumeVarint(resp)
+	if n1 < 0 {
+		return nil, protowire.ParseError(n1)
+	}
+
+	_, n2 := protowire.ConsumeVarint(resp[n1:]) // length
+	if n2 < 0 {
+		return nil, protowire.ParseError(n2)
+	}
+
+	if typ != (metadataField<<3)|metadataType {
+		return nil, fmt.Errorf("unexpected message format: %d", typ)
+	}
+
+	if n1+n2 > len(resp) {
+		return nil, fmt.Errorf("unexpected message size: %d", len(resp))
+	}
+
+	// cut off embedded message header
+	resp = resp[n1+n2:]
+	// build new embedded message header
+	prefix := protowire.AppendVarint(
+		protowire.AppendVarint(nil, (metadataField<<3)|metadataType),
+		uint64(len(resp)+len(payload)),
+	)
+	resp = append(prefix, resp...)
+
+	return append(resp, payload...), err
+}
+
+// BuildError converts upstream error into message from upstream, so that multiple
+// successful and failure responses might be returned.
+//
+// This simply relies on the fact that any response contains 'Empty' message.
+// So if 'Empty' is unmarshalled into any other reply message, all the fields
+// are undefined but 'Metadata':
+//
+//	message Empty {
+//	  common.Metadata metadata = 1;
+//	}
+//
+//	message EmptyResponse {
+//	  repeated Empty messages = 1;
+//	}
+//
+// Streaming responses are not wrapped into Empty, so we simply marshall EmptyResponse
+// message.
+func (b *RemoteBackend) BuildError(streaming bool, err error) ([]byte, error) {
+	var resp proto.Message = &pb.Empty{
+		Metadata: &pb.Metadata{
+			Machine: b.target,
+			Error:   err.Error(),
+			Status:  status.Convert(err).Proto(),
+		},
+	}
+
+	if !streaming {
+		resp = &pb.EmptyResponse{
+			Messages: []*pb.Empty{
+				resp.(*pb.Empty),
+			},
+		}
+	}
+
+	return proto.Marshal(resp)
+}
+
+// Close closes the upstream gRPC connection.
+func (b *RemoteBackend) Close() {
+	b.mu.Lock()
+	defer b.mu.Unlock()
+
+	if b.conn != nil {
+		b.conn.Close()
+		b.conn = nil
+	}
+}