connect-go icon indicating copy to clipboard operation
connect-go copied to clipboard

[Question] middleware set context not work in stream rpc

Open solarhell opened this issue 1 year ago • 1 comments

my middleware

package middleware

import (
	"context"
	"fmt"

	"connectrpc.com/connect"
	"pkg/ctxuser"
)

func NewAuthInterceptor() connect.UnaryInterceptorFunc {
	interceptor := func(next connect.UnaryFunc) connect.UnaryFunc {
		return func(
			ctx context.Context,
			req connect.AnyRequest,
		) (connect.AnyResponse, error) {
			tokenString := req.Header().Get("token")
			if tokenString == "" {
				return nil, connect.NewError(connect.CodeUnauthenticated, fmt.Errorf("token is empty"))
			}

			userID, err := getUserIDFromToken(tokenString)
			...

			ctx = ctxuser.SetInContext(ctx, ctxuser.RequestInfo{
				UserID: userID,
			})

			return next(ctx, req)
		}
	}
	return interceptor
}
package ctxuser

import (
	"context"
)

type ctxMarker struct{}

var (
	ctxMarkerKey = &ctxMarker{}
)

type RequestInfo struct {
	RequestID string

	UserID string
}

func Extract(ctx context.Context) (RequestInfo, bool) {
	info, ok := ctx.Value(ctxMarkerKey).(RequestInfo)
	return info, ok
}

func SetInContext(ctx context.Context, info RequestInfo) context.Context {
	return context.WithValue(ctx, ctxMarkerKey, info)
}

in my api, i use token in http header to check if user is logged in. and i use middleware to set the context, so i can get user_id easily in handler.

But in such rpc handler i can't get the user_id.

my proto

  rpc AskGPT(AskGPTRequest) returns(stream AskGPTResponse);

Do i need to change my middleware code? Thanks.

solarhell avatar Jun 05 '24 09:06 solarhell

The context code looks correct, however (stream AskGPTResponse); requires you implement the full Interceptor interface. The current interceptor will only work on unary, it doesn't cover streaming use cases.

Would recommend to look at the authn-go library: https://github.com/connectrpc/authn-go/tree/main This will correctly handle setting the context value in the request and has some security benefits doing this as HTTP middleware.

emcfarlane avatar Jun 05 '24 14:06 emcfarlane

After implemented the WrapStreamingHandler(StreamingHandlerFunc) StreamingHandlerFunc, it works fine now.

solarhell avatar Jul 03 '24 06:07 solarhell