Refactor Client Interceptors and add auth header for streams (#84)

Author: majst01

doc/index.html

@@ -1984,6 +1984,14 @@ <h2>Table of Contents</h2>
                   <a href="#metalstack.infra.v2.UpdateBMCInfoResponse"><span class="badge">M</span>UpdateBMCInfoResponse</a>
                 </li>
 
+                <li>
+                  <a href="#metalstack.infra.v2.WaitForMachineEventRequest"><span class="badge">M</span>WaitForMachineEventRequest</a>
+                </li>
+              
+                <li>
+                  <a href="#metalstack.infra.v2.WaitForMachineEventResponse"><span class="badge">M</span>WaitForMachineEventResponse</a>
+                </li>
+              
 
 
 
@@ -15673,6 +15681,20 @@ <h3 id="metalstack.infra.v2.UpdateBMCInfoResponse">UpdateBMCInfoResponse</h3>
 
 
 
+        <h3 id="metalstack.infra.v2.WaitForMachineEventRequest">WaitForMachineEventRequest</h3>
+        <p>WaitForMachineEventRequest</p>
+
+        
+
+        
+      
+        <h3 id="metalstack.infra.v2.WaitForMachineEventResponse">WaitForMachineEventResponse</h3>
+        <p>WaitForMachineEventResponse</p>
+
+        
+
+        
+      
 
 
 
@@ -15694,6 +15716,13 @@ <h3 id="metalstack.infra.v2.BMCService">BMCService</h3>
                 <td><p>UpdateBMCInfo</p></td>
               </tr>
 
+              <tr>
+                <td>WaitForMachineEvent</td>
+                <td><a href="#metalstack.infra.v2.WaitForMachineEventRequest">WaitForMachineEventRequest</a></td>
+                <td><a href="#metalstack.infra.v2.WaitForMachineEventResponse">WaitForMachineEventResponse</a> stream</td>
+                <td><p>WaitForMachineEvent is called by the metal-bmc and is returned with a bmc command to execute.</p></td>
+              </tr>
+            
           </tbody>
         </table>
 

generate/go_client.tpl

@@ -2,7 +2,6 @@
 package client
 
 import (
-	"context"
 	"sync"
 
 	"connectrpc.com/connect"
@@ -53,33 +52,17 @@ func New(config *DialConfig) (Client, error) {
 		interceptors: []connect.Interceptor{},
 	}
 
-	authInterceptor := connect.UnaryInterceptorFunc(func(next connect.UnaryFunc) connect.UnaryFunc {
-		return connect.UnaryFunc(func(ctx context.Context, request connect.AnyRequest) (connect.AnyResponse, error) {
-			request.Header().Add("Authorization", "Bearer "+config.Token)
-			return next(ctx, request)
-		})
-	})
-
-	loggingInterceptor := connect.UnaryInterceptorFunc(func(next connect.UnaryFunc) connect.UnaryFunc {
-		return connect.UnaryFunc(func(ctx context.Context, request connect.AnyRequest) (connect.AnyResponse, error) {
-			config.Log.Debug("intercept", "request procedure", request.Spec().Procedure, "body", request.Any())
-			response, err := next(ctx, request)
-			if err != nil {
-				return nil, err
-			}
-			config.Log.Debug("intercept", "request procedure", request.Spec().Procedure, "response", response.Any())
-			return response, err
-		})
-	})
-
 	if config.Token != "" {
+		authInterceptor := &authInterceptor{config: config}
 		c.interceptors = append(c.interceptors, authInterceptor)
 	}
 	if config.Log != nil {
+		loggingInterceptor := &loggingInterceptor{config: config}
 		c.interceptors = append(c.interceptors, loggingInterceptor)
 	}
 	c.interceptors = append(c.interceptors, config.Interceptors...)
 
+	// TODO convert to interceptor
 	go c.startTokenRenewal()
 
 	return c, nil

go.mod

@@ -7,6 +7,7 @@ require (
 	buf.build/go/protovalidate v1.1.0
 	connectrpc.com/connect v1.19.1
 	github.com/bufbuild/protocompile v0.14.1
+	github.com/davecgh/go-spew v1.1.1
 	github.com/go-task/slim-sprig/v3 v3.0.0
 	github.com/golang-jwt/jwt/v5 v5.3.0
 	github.com/google/go-cmp v0.7.0
@@ -18,7 +19,6 @@ require (
 require (
 	cel.dev/expr v0.25.1 // indirect
 	github.com/antlr4-go/antlr/v4 v4.13.1 // indirect
-	github.com/davecgh/go-spew v1.1.1 // indirect
 	github.com/google/cel-go v0.26.1 // indirect
 	github.com/klauspost/compress v1.18.2 // indirect
 	github.com/kr/pretty v0.3.1 // indirect

go/client/client-interceptors.go

@@ -0,0 +1,67 @@
+package client
+
+import (
+	"context"
+
+	"connectrpc.com/connect"
+)
+
+// authinterceptor adds the required auth headers
+type authInterceptor struct {
+	config *DialConfig
+}
+
+func (i *authInterceptor) WrapUnary(next connect.UnaryFunc) connect.UnaryFunc {
+	return connect.UnaryFunc(func(ctx context.Context, request connect.AnyRequest) (connect.AnyResponse, error) {
+		request.Header().Add("Authorization", "Bearer "+i.config.Token)
+		return next(ctx, request)
+	})
+}
+
+func (i *authInterceptor) WrapStreamingClient(next connect.StreamingClientFunc) connect.StreamingClientFunc {
+	return func(ctx context.Context, spec connect.Spec) connect.StreamingClientConn {
+		return &streamingInterceptorConn{
+			StreamingClientConn: next(ctx, spec),
+			token:               i.config.Token,
+		}
+	}
+}
+
+func (i *authInterceptor) WrapStreamingHandler(next connect.StreamingHandlerFunc) connect.StreamingHandlerFunc {
+	return next
+}
+
+type streamingInterceptorConn struct {
+	connect.StreamingClientConn
+	token string
+}
+
+func (conn *streamingInterceptorConn) Send(m any) error {
+	conn.RequestHeader().Add("Authorization", "Bearer "+conn.token)
+	return conn.StreamingClientConn.Send(m)
+}
+
+type loggingInterceptor struct {
+	config *DialConfig
+}
+
+func (i *loggingInterceptor) WrapUnary(next connect.UnaryFunc) connect.UnaryFunc {
+	return connect.UnaryFunc(func(ctx context.Context, request connect.AnyRequest) (connect.AnyResponse, error) {
+		i.config.Log.Debug("intercept", "request procedure", request.Spec().Procedure, "body", request.Any())
+		response, err := next(ctx, request)
+		if err != nil {
+			return nil, err
+		}
+		i.config.Log.Debug("intercept", "request procedure", request.Spec().Procedure, "response", response.Any())
+		return response, err
+	})
+}
+
+func (i *loggingInterceptor) WrapStreamingClient(next connect.StreamingClientFunc) connect.StreamingClientFunc {
+	// TODO also log here
+	return next
+}
+
+func (i *loggingInterceptor) WrapStreamingHandler(next connect.StreamingHandlerFunc) connect.StreamingHandlerFunc {
+	return next
+}

go/client/client.go

@@ -20,6 +20,8 @@ import (
 	"github.com/metal-stack/api/go/client"
 	apiv2 "github.com/metal-stack/api/go/metalstack/api/v2"
 	"github.com/metal-stack/api/go/metalstack/api/v2/apiv2connect"
+	infrav2 "github.com/metal-stack/api/go/metalstack/infra/v2"
+	"github.com/metal-stack/api/go/metalstack/infra/v2/infrav2connect"
 	"github.com/stretchr/testify/require"
 )
 
@@ -155,3 +157,74 @@ func (m *mockTokenService) Revoke(context.Context, *apiv2.TokenServiceRevokeRequ
 func (m *mockTokenService) Update(context.Context, *apiv2.TokenServiceUpdateRequest) (*apiv2.TokenServiceUpdateResponse, error) {
 	panic("unimplemented")
 }
+
+func Test_ClientInterceptors(t *testing.T) {
+	var (
+		bs  = &mockBMCService{}
+		mux = http.NewServeMux()
+		log = slog.New(slog.NewJSONHandler(os.Stdout, &slog.HandlerOptions{Level: slog.LevelDebug}))
+		ctx = t.Context()
+	)
+
+	mux.Handle(infrav2connect.NewBMCServiceHandler(bs))
+	server := httptest.NewTLSServer(mux)
+	server.EnableHTTP2 = true
+	defer func() {
+		server.Close()
+	}()
+
+	tokenString, err := generateToken(1 * time.Second)
+	require.NoError(t, err)
+
+	c, err := client.New(&client.DialConfig{
+		BaseURL:   server.URL,
+		Token:     tokenString,
+		Transport: server.Client().Transport,
+		Log:       log,
+	})
+	require.NoError(t, err)
+
+	// Synchronous call has authheader set
+	resp, err := c.Infrav2().BMC().UpdateBMCInfo(ctx, &infrav2.UpdateBMCInfoRequest{})
+	require.NoError(t, err)
+	require.NotNil(t, resp)
+	require.Equal(t, tokenString, bs.token)
+	bs.token = ""
+
+	// Asynchronous call has authheader set
+	stream, err := c.Infrav2().BMC().WaitForMachineEvent(ctx, &infrav2.WaitForMachineEventRequest{})
+	require.NoError(t, err)
+	require.NotNil(t, stream)
+	require.Equal(t, tokenString, bs.token)
+}
+
+type mockBMCService struct {
+	token string
+}
+
+func (m *mockBMCService) UpdateBMCInfo(ctx context.Context, _ *infrav2.UpdateBMCInfoRequest) (*infrav2.UpdateBMCInfoResponse, error) {
+	callinfo, _ := connect.CallInfoForHandlerContext(ctx)
+	authHeader := callinfo.RequestHeader().Get("Authorization")
+
+	_, token, found := strings.Cut(authHeader, "Bearer ")
+
+	if !found {
+		return nil, fmt.Errorf("unable to extract token from header:%s", authHeader)
+	}
+	m.token = token
+	return &infrav2.UpdateBMCInfoResponse{}, nil
+}
+
+func (m *mockBMCService) WaitForMachineEvent(ctx context.Context, _ *infrav2.WaitForMachineEventRequest, stream *connect.ServerStream[infrav2.WaitForMachineEventResponse]) error {
+	callinfo, _ := connect.CallInfoForHandlerContext(ctx)
+	authHeader := callinfo.RequestHeader().Get("Authorization")
+
+	_, token, found := strings.Cut(authHeader, "Bearer ")
+
+	if !found {
+		return fmt.Errorf("unable to extract token from header:%s", authHeader)
+	}
+
+	m.token = token
+	return stream.Send(&infrav2.WaitForMachineEventResponse{})
+}