Go语言gRPC双向流式通信与拦截器中间件链式设计实战

package chatpb;
option go_package = "chatpb/chatpb";

message ChatMessage {
string user_id = 1;
string content = 2;
int64 timestamp = 3;
string room_id = 4;
}

message JoinRequest {
string user_id = 1;
string room_id = 2;
}

message ServerResponse {
string message_id = 1;
bool ack = 2;
string error = 3;
}

func (s *ChatServer) JoinChat(stream chatpb.ChatService_JoinChatServer) error {
// 读取客户端首条消息获取用户信息和房间号
firstMsg, err := stream.Recv()
if err != nil {
return fmt.Errorf("读取首条消息失败: %w", err)
}
userID := firstMsg.UserId
roomID := firstMsg.RoomId

// 注册当前流到房间
s.registerStream(roomID, userID, stream)
defer s.unregisterStream(roomID, userID)

// 启动发送协程
sendErr := make(chan error, 1)
go func() {
sendErr <- s.forwardMessages(roomID, userID, stream)
}()

// 接收循环
recvErr := make(chan error, 1)
go func() {
for {
msg, err := stream.Recv()
if err != nil {
recvErr <- err
return
}
// 广播消息到房间内其他用户
s.broadcastMessage(roomID, userID, msg)
}
}()

client := chatpb.NewChatServiceClient(conn)
stream, err := client.JoinChat(context.Background())
if err != nil {
return err
}

// 发送首条消息注册房间
stream.Send(&chatpb.ChatMessage{
UserId: userID,
RoomId: roomID,
Content: "__join__",
Timestamp: time.Now().Unix(),
})

// 接收协程
done := make(chan struct{})
go func() {
defer close(done)
for {
resp, err := stream.Recv()
if err != nil {
log.Printf("接收结束: %v", err)
return
}
log.Printf("收到消息: ack=%v msg_id=%s", resp.Ack, resp.MessageId)
}
}()

// 主协程读取标准输入并发送
scanner := bufio.NewScanner(os.Stdin)
for scanner.Scan() {
text := scanner.Text()
if text == "quit" {
stream.CloseSend()
break
}
err := stream.Send(&chatpb.ChatMessage{
UserId: userID,
RoomId: roomID,
Content: text,
Timestamp: time.Now().Unix(),
})
if err != nil {
log.Printf("发送失败: %v", err)
break
}
}

// 日志拦截器
func LoggingInterceptor(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (interface{}, error) {
start := time.Now()
resp, err := handler(ctx, req)
duration := time.Since(start)

log.Printf("method=%s duration=%s error=%v", info.FullMethod, duration, err)
return resp, err
}

// 认证拦截器
func AuthInterceptor(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (interface{}, error) {
md, ok := metadata.FromIncomingContext(ctx)
if !ok {
return nil, status.Error(codes.Unauthenticated, "缺少元数据")
}
tokens := md.Get("authorization")
if len(tokens) == 0 {
return nil, status.Error(codes.Unauthenticated, "缺少认证token")
}
userID, err := validateToken(tokens[0])
if err != nil {
return nil, status.Error(codes.Unauthenticated, "token无效")
}
// 注入用户ID到context
ctx = context.WithValue(ctx, "userID", userID)
return handler(ctx, req)
}

limiter := rate.NewLimiter(rate.Every(time.Second/100), 200)

server := grpc.NewServer(
grpc.UnaryInterceptor(
ChainUnaryInterceptors(
LoggingInterceptor,
AuthInterceptor,
RateLimitInterceptor(limiter),
),
),
grpc.MaxRecvMsgSize(16 * 1024 * 1024),
grpc.MaxSendMsgSize(16 * 1024 * 1024),
)

chatpb.RegisterChatServiceServer(server, &ChatServer{
rooms: make(map[string]map[string]chatpb.ChatService_JoinChatServer),
})

// 从context获取peer信息
peerInfo, _ := peer.FromContext(ctx)
clientAddr := "unknown"
if peerInfo != nil {
clientAddr = peerInfo.Addr.String()
}

err := handler(srv, ss)
duration := time.Since(start)

log.Printf("stream method=%s client=%s duration=%s error=%v",
info.FullMethod, clientAddr, duration, err)
return err
}

// 包装ServerStream注入context值
type wrappedStream struct {
grpc.ServerStream
ctx context.Context
}

func (w *wrappedStream) Context() context.Context {
return w.ctx
}

func StreamAuthInterceptor(srv interface{}, ss grpc.ServerStream, info *grpc.StreamServerInfo, handler grpc.StreamHandler) error {
md, ok := metadata.FromIncomingContext(ss.Context())
if !ok {
return status.Error(codes.Unauthenticated, "缺少元数据")
}
tokens := md.Get("authorization")
if len(tokens) == 0 {
return status.Error(codes.Unauthenticated, "缺少认证token")
}
userID, err := validateToken(tokens[0])
if err != nil {
return status.Error(codes.Unauthenticated, "token无效")
}