Repository: go-kratos/kratos (go)

Issue:
Kratos middleware seems to only apply to unary gRPC operations
#### What happened:
We've started using Kratos middleware (recovery, logging, validate) in our project, and noticed that, on the gRPC side, they only seem to apply to unary operations. Our streaming operations don't get validated or logged, and a panic takes down the server (versus on a unary endpoint, they're validated and logged as expected, and panics are handled.)

#### What you expected to happen:
We expected to either have the middleware apply seamlessly to streaming requests or to have separate streaming middleware we can register (which, to be fair, may exist, but we weren't able to find it.)

#### How to reproduce it (as minimally and precisely as possible):
Repro project: https://github.com/wscalf/streaming-repro

- "SayHello" is the normal hello world operation + a 50/50 chance of panicking to test the recovery middleware.
- "KeepSayingHello" is an added operation that streams the response every second until the client terminates the connection, at which point it panics (intentionally, to test 
Commands:
- `grpcurl -plaintext -d '{"name":""}' localhost:9000 helloworld.v1.Greeter.SayHello`: gets correctly rejected with a validation error
- `grpcurl -plaintext -d '{"name":""}' localhost:9000 helloworld.v1.Greeter.KeepSayingHello`: should get rejected (it's using the same message) but does not. Also, is coded to panic when the client closes, which should trigger the recovery middleware, but it doesn't- it just exits
- `grpcurl -plaintext -d '{"name":"Kratos"}' localhost:9000 helloworld.v1.Greeter.SayHello`: is valid, and is coded to have a 50/50 chance of panicking, which does correctly trigger the recovery middleware, returning an error to the client and not exiting.

#### Anything else we need to know?:

#### Environment:
- Kratos version (use `kratos -v`): kratos version v2.7.3
- Go version (use `go version`): go version go1.22.2 linux/amd64
- OS (e.g: `cat /etc/os-release`):
```
NAME="Fedora Linux"
VERSION="40 (Workstation Edition)"
ID=fedora
VERSION_ID=40
VERSION_CODENAME=""
PLATFORM_ID="platform:f40"
PRETTY_NAME="Fedora Linux 40 (Workstation Edition)"
ANSI_COLOR="0;38;2;60;110;180"
LOGO=fedora-logo-icon
CPE_NAME="cpe:/o:fedoraproject:fedora:40"
DEFAULT_HOSTNAME="fedora"
HOME_URL="https://fedoraproject.org/"
DOCUMENTATION_URL="https://docs.fedoraproject.org/en-US/fedora/f40/system-administrators-guide/"
SUPPORT_URL="https://ask.fedoraproject.org/"
BUG_REPORT_URL="https://bugzilla.redhat.com/"
REDHAT_BUGZILLA_PRODUCT="Fedora"
REDHAT_BUGZILLA_PRODUCT_VERSION=40
REDHAT_SUPPORT_PRODUCT="Fedora"
REDHAT_SUPPORT_PRODUCT_VERSION=40
SUPPORT_END=2025-05-13
VARIANT="Workstation Edition"
VARIANT_ID=workstation

```
- Others:


Pull request:
adding stream interceptor for logging middleware
#### Description (what this PR does / why we need it):

Update the stream interceptor to handle stream use-case

<!--
🎉 Thanks for sending a pull request to Kratos! Here are some tips for you:

1. If this is your first time contributing to Kratos, please read our contribution guide: https://go-kratos.dev/en/docs/community/contribution/
2. Ensure you have added or ran the appropriate tests and lint for your PR, please use `make lint` and `make test` before filing your PR, use `make clean` to tidy your go mod.
3. If the PR is unfinished, you may need to mark it as a WIP(Work In Progress) PR or Draft PR
4. Please use a semantic commits format title, such as `<type>[optional scope]: <description>`, see: https://go-kratos.dev/docs/community/contribution#type
5. at the same time, please note that similar work should be submitted in one PR as far as possible to reduce the workload of reviewers. Do not split a work into multiple PR unless it should.
-->

<!--
🎉 感谢您向 Kratos 发送 PR！以下是一些提示：
如果这是你第一次为 Kratos 贡献，请阅读我们的贡献指南：https://go-kratos.dev/en/docs/community/contribution/
2、确保您已经为您的 PR 添加或运行了适当的测试和lint，请在提交PR之前使用“make lint”和“make test”，使用“make clean”整理您的 go.mod。
3、如果 PR 未完成，您可能需要将其标记为 WIP（Work In Progress）PR 或 Draft PR
4、请使用语义提交格式标题，如“<类型>[可选范围]：<说明>`，请参阅：https://go-kratos.dev/docs/community/contribution#type
5. 同时请注意，同类的工作请尽量在一个PR中提交，以减轻 review 者的工作负担，不要把一项工作拆分成很多个PR，除非它应该这样做。
-->



<!--
* The description should include the motivation for this PR or contrast this with previous behavior
-->


#### Which issue(s) this PR fixes (resolves / be part of):
<!--
* Automatically closes linked issue when PR is merged.
* If your PR is not fully resolved the issue, please use `part of #<issue number>` instead.

Usage: `fixes/resolves #<issue number>`, or `fixes/resolves (paste link of issue)`.
-->
fixes #3351

#### Other special notes for the reviewers:
<!--
* Somethings that need extra attention for the reviewers
* Some additional notes, TODO list, etc.
-->


Patch:
diff --git a/transport/grpc/client.go b/transport/grpc/client.go
index edc57d3e..5182b180 100644
--- a/transport/grpc/client.go
+++ b/transport/grpc/client.go
@@ -11,6 +11,7 @@ import (
 	grpcinsecure "google.golang.org/grpc/credentials/insecure"
 	grpcmd "google.golang.org/grpc/metadata"
 
+	"github.com/go-kratos/kratos/v2/internal/matcher"
 	"github.com/go-kratos/kratos/v2/log"
 	"github.com/go-kratos/kratos/v2/middleware"
 	"github.com/go-kratos/kratos/v2/registry"
@@ -132,6 +133,7 @@ type clientOptions struct {
 	timeout                time.Duration
 	discovery              registry.Discovery
 	middleware             []middleware.Middleware
+	streamMiddleware       []middleware.Middleware
 	ints                   []grpc.UnaryClientInterceptor
 	streamInts             []grpc.StreamClientInterceptor
 	grpcOpts               []grpc.DialOption
@@ -166,7 +168,7 @@ func dial(ctx context.Context, insecure bool, opts ...ClientOption) (*grpc.Clien
 		unaryClientInterceptor(options.middleware, options.timeout, options.filters),
 	}
 	sints := []grpc.StreamClientInterceptor{
-		streamClientInterceptor(options.filters),
+		streamClientInterceptor(options.streamMiddleware, options.filters),
 	}
 
 	if len(options.ints) > 0 {
@@ -239,7 +241,54 @@ func unaryClientInterceptor(ms []middleware.Middleware, timeout time.Duration, f
 	}
 }
 
-func streamClientInterceptor(filters []selector.NodeFilter) grpc.StreamClientInterceptor {
+// wrappedClientStream wraps the grpc.ClientStream and applies middleware
+type wrappedClientStream struct {
+	grpc.ClientStream
+	ctx        context.Context
+	middleware matcher.Matcher
+}
+
+func (w *wrappedClientStream) Context() context.Context {
+	return w.ctx
+}
+
+func (w *wrappedClientStream) SendMsg(m interface{}) error {
+	h := func(ctx context.Context, req interface{}) (interface{}, error) {
+		return req, w.ClientStream.SendMsg(m)
+	}
+
+	info, ok := transport.FromClientContext(w.ctx)
+	if !ok {
+		return fmt.Errorf("transport value stored in ctx returns: %v", ok)
+	}
+
+	if next := w.middleware.Match(info.Operation()); len(next) > 0 {
+		h = middleware.Chain(next...)(h)
+	}
+
+	_, err := h(w.ctx, m)
+	return err
+}
+
+func (w *wrappedClientStream) RecvMsg(m interface{}) error {
+	h := func(ctx context.Context, req interface{}) (interface{}, error) {
+		return req, w.ClientStream.RecvMsg(m)
+	}
+
+	info, ok := transport.FromClientContext(w.ctx)
+	if !ok {
+		return fmt.Errorf("transport value stored in ctx returns: %v", ok)
+	}
+
+	if next := w.middleware.Match(info.Operation()); len(next) > 0 {
+		h = middleware.Chain(next...)(h)
+	}
+
+	_, err := h(w.ctx, m)
+	return err
+}
+
+func streamClientInterceptor(ms []middleware.Middleware, filters []selector.NodeFilter) grpc.StreamClientInterceptor {
 	return func(ctx context.Context, desc *grpc.StreamDesc, cc *grpc.ClientConn, method string, streamer grpc.Streamer, opts ...grpc.CallOption) (grpc.ClientStream, error) { // nolint
 		ctx = transport.NewClientContext(ctx, &Transport{
 			endpoint:    cc.Target(),
@@ -249,6 +298,28 @@ func streamClientInterceptor(filters []selector.NodeFilter) grpc.StreamClientInt
 		})
 		var p selector.Peer
 		ctx = selector.NewPeerContext(ctx, &p)
-		return streamer(ctx, desc, cc, method, opts...)
+
+		clientStream, err := streamer(ctx, desc, cc, method, opts...)
+		if err != nil {
+			return nil, err
+		}
+
+		h := func(ctx context.Context, req interface{}) (interface{}, error) {
+			return streamer, nil
+		}
+
+		m := matcher.New()
+		if len(ms) > 0 {
+			m.Use(ms...)
+			middleware.Chain(ms...)(h)
+		}
+
+		wrappedStream := &wrappedClientStream{
+			ClientStream: clientStream,
+			ctx:          ctx,
+			middleware:   m,
+		}
+
+		return wrappedStream, nil
 	}
 }
diff --git a/transport/grpc/interceptor.go b/transport/grpc/interceptor.go
index 6cc33154..6261442d 100644
--- a/transport/grpc/interceptor.go
+++ b/transport/grpc/interceptor.go
@@ -2,11 +2,13 @@ package grpc
 
 import (
 	"context"
+	"fmt"
 
 	"google.golang.org/grpc"
 	grpcmd "google.golang.org/grpc/metadata"
 
 	ic "github.com/go-kratos/kratos/v2/internal/context"
+	"github.com/go-kratos/kratos/v2/internal/matcher"
 	"github.com/go-kratos/kratos/v2/middleware"
 	"github.com/go-kratos/kratos/v2/transport"
 )
@@ -48,13 +50,15 @@ func (s *Server) unaryServerInterceptor() grpc.UnaryServerInterceptor {
 // wrappedStream is rewrite grpc stream's context
 type wrappedStream struct {
 	grpc.ServerStream
-	ctx context.Context
+	ctx        context.Context
+	middleware matcher.Matcher
 }
 
-func NewWrappedStream(ctx context.Context, stream grpc.ServerStream) grpc.ServerStream {
+func NewWrappedStream(ctx context.Context, stream grpc.ServerStream, m matcher.Matcher) grpc.ServerStream {
 	return &wrappedStream{
 		ServerStream: stream,
 		ctx:          ctx,
+		middleware:   m,
 	}
 }
 
@@ -76,7 +80,19 @@ func (s *Server) streamServerInterceptor() grpc.StreamServerInterceptor {
 			replyHeader: headerCarrier(replyHeader),
 		})
 
-		ws := NewWrappedStream(ctx, ss)
+		h := func(ctx context.Context, req interface{}) (interface{}, error) {
+			return handler(srv, ss), nil
+		}
+
+		if next := s.streamMiddleware.Match(info.FullMethod); len(next) > 0 {
+			middleware.Chain(next...)(h)
+		}
+
+		ctx = context.WithValue(ctx, stream{
+			ServerStream:     ss,
+			streamMiddleware: s.streamMiddleware,
+		}, ss)
+		ws := NewWrappedStream(ctx, ss, s.streamMiddleware)
 
 		err := handler(srv, ws)
 		if len(replyHeader) > 0 {
@@ -85,3 +101,48 @@ func (s *Server) streamServerInterceptor() grpc.StreamServerInterceptor {
 		return err
 	}
 }
+
+type stream struct {
+	grpc.ServerStream
+	streamMiddleware matcher.Matcher
+}
+
+func GetStream(ctx context.Context) grpc.ServerStream {
+	return ctx.Value(stream{}).(grpc.ServerStream)
+}
+
+func (w *wrappedStream) SendMsg(m interface{}) error {
+	h := func(_ context.Context, req interface{}) (interface{}, error) {
+		return req, w.ServerStream.SendMsg(m)
+	}
+
+	info, ok := transport.FromServerContext(w.ctx)
+	if !ok {
+		return fmt.Errorf("transport value stored in ctx returns: %v", ok)
+	}
+
+	if next := w.middleware.Match(info.Operation()); len(next) > 0 {
+		h = middleware.Chain(next...)(h)
+	}
+
+	_, err := h(w.ctx, m)
+	return err
+}
+
+func (w *wrappedStream) RecvMsg(m interface{}) error {
+	h := func(_ context.Context, req interface{}) (interface{}, error) {
+		return req, w.ServerStream.RecvMsg(m)
+	}
+
+	info, ok := transport.FromServerContext(w.ctx)
+	if !ok {
+		return fmt.Errorf("transport value stored in ctx returns: %v", ok)
+	}
+
+	if next := w.middleware.Match(info.Operation()); len(next) > 0 {
+		h = middleware.Chain(next...)(h)
+	}
+
+	_, err := h(w.ctx, m)
+	return err
+}
diff --git a/transport/grpc/server.go b/transport/grpc/server.go
index c8c3f4f3..a1d74ace 100644
--- a/transport/grpc/server.go
+++ b/transport/grpc/server.go
@@ -72,6 +72,12 @@ func Middleware(m ...middleware.Middleware) ServerOption {
 	}
 }
 
+func StreamMiddleware(m ...middleware.Middleware) ServerOption {
+	return func(s *Server) {
+		s.streamMiddleware.Use(m...)
+	}
+}
+
 // CustomHealth Checks server.
 func CustomHealth() ServerOption {
 	return func(s *Server) {
@@ -117,33 +123,35 @@ func Options(opts ...grpc.ServerOption) ServerOption {
 // Server is a gRPC server wrapper.
 type Server struct {
 	*grpc.Server
-	baseCtx      context.Context
-	tlsConf      *tls.Config
-	lis          net.Listener
-	err          error
-	network      string
-	address      string
-	endpoint     *url.URL
-	timeout      time.Duration
-	middleware   matcher.Matcher
-	unaryInts    []grpc.UnaryServerInterceptor
-	streamInts   []grpc.StreamServerInterceptor
-	grpcOpts     []grpc.ServerOption
-	health       *health.Server
-	customHealth bool
-	metadata     *apimd.Server
-	adminClean   func()
+	baseCtx          context.Context
+	tlsConf          *tls.Config
+	lis              net.Listener
+	err              error
+	network          string
+	address          string
+	endpoint         *url.URL
+	timeout          time.Duration
+	middleware       matcher.Matcher
+	streamMiddleware matcher.Matcher
+	unaryInts        []grpc.UnaryServerInterceptor
+	streamInts       []grpc.StreamServerInterceptor
+	grpcOpts         []grpc.ServerOption
+	health           *health.Server
+	customHealth     bool
+	metadata         *apimd.Server
+	adminClean       func()
 }
 
 // NewServer creates a gRPC server by options.
 func NewServer(opts ...ServerOption) *Server {
 	srv := &Server{
-		baseCtx:    context.Background(),
-		network:    "tcp",
-		address:    ":0",
-		timeout:    1 * time.Second,
-		health:     health.NewServer(),
-		middleware: matcher.New(),
+		baseCtx:          context.Background(),
+		network:          "tcp",
+		address:          ":0",
+		timeout:          1 * time.Second,
+		health:           health.NewServer(),
+		middleware:       matcher.New(),
+		streamMiddleware: matcher.New(),
 	}
 	for _, o := range opts {
 		o(srv)


Test patch:
diff --git a/transport/grpc/server_test.go b/transport/grpc/server_test.go
index d9d2c7af..067c4859 100644
--- a/transport/grpc/server_test.go
+++ b/transport/grpc/server_test.go
@@ -12,6 +12,7 @@ import (
 	"time"
 
 	"google.golang.org/grpc"
+	"google.golang.org/grpc/metadata"
 
 	"github.com/go-kratos/kratos/v2/errors"
 	"github.com/go-kratos/kratos/v2/internal/matcher"
@@ -280,6 +281,82 @@ func TestServer_unaryServerInterceptor(t *testing.T) {
 	}
 }
 
+type mockServerStream struct {
+	ctx      context.Context
+	sentMsg  interface{}
+	recvMsg  interface{}
+	metadata metadata.MD
+	grpc.ServerStream
+}
+
+func (m *mockServerStream) SetHeader(md metadata.MD) error {
+	m.metadata = md
+	return nil
+}
+
+func (m *mockServerStream) SendHeader(md metadata.MD) error {
+	m.metadata = md
+	return nil
+}
+
+func (m *mockServerStream) SetTrailer(md metadata.MD) {
+	m.metadata = md
+}
+
+func (m *mockServerStream) Context() context.Context {
+	return m.ctx
+}
+
+func (m *mockServerStream) SendMsg(msg interface{}) error {
+	m.sentMsg = msg
+	return nil
+}
+
+func (m *mockServerStream) RecvMsg(msg interface{}) error {
+	m.recvMsg = msg
+	return nil
+}
+
+func TestServer_streamServerInterceptor(t *testing.T) {
+	u, err := url.Parse("grpc://hello/world")
+	if err != nil {
+		t.Errorf("expect %v, got %v", nil, err)
+	}
+	srv := &Server{
+		baseCtx:          context.Background(),
+		endpoint:         u,
+		timeout:          time.Duration(10),
+		middleware:       matcher.New(),
+		streamMiddleware: matcher.New(),
+	}
+
+	srv.streamMiddleware.Use(EmptyMiddleware())
+
+	mockStream := &mockServerStream{
+		ctx: srv.baseCtx,
+	}
+
+	handler := func(_ interface{}, stream grpc.ServerStream) error {
+		resp := &testResp{Data: "stream hi"}
+		return stream.SendMsg(resp)
+	}
+
+	info := &grpc.StreamServerInfo{
+		FullMethod: "/grpc.reflection.v1.ServerReflection/ServerReflectionInfo",
+	}
+
+	err = srv.streamServerInterceptor()(nil, mockStream, info, handler)
+	if err != nil {
+		t.Errorf("expect %v, got %v", nil, err)
+	}
+
+	// Check response
+	resp := mockStream.sentMsg.(*testResp)
+	if !reflect.DeepEqual("stream hi", resp.Data) {
+		t.Errorf("expect %s, got %s", "stream hi", resp.Data)
+	}
+}
+
 func TestListener(t *testing.T) {
 	lis, err := net.Listen("tcp", ":0")
 	if err != nil {