grpc中的攔截器

slowquery發表於2022-10-16
RPC

0.1、索引

waterflow.link/articles/1665853719...

當我們編寫 HTTP 應用程式時，您可以使用 HTTP 中介軟體包裝特定於路由的應用程式處理程式，可以在執行應用程式處理程式之前和之後執行一些常見的邏輯。 我們通常使用中介軟體來編寫跨領域元件，例如授權、日誌記錄、快取等。在 gRPC 中可以使用稱為攔截器的概念來實現相同的功能。

透過使用攔截器，我們可以在客戶端和伺服器上攔截 RPC 方法的執行。 在客戶端和伺服器上，都有兩種型別的攔截器：

  • UnaryInterceptor（一元攔截器）
  • StreamInterceptor（流式攔截器）

UnaryInterceptor 攔截一元 RPC，而 StreamInterceptor 攔截流式 RPC。

在一元 RPC 中，客戶端向伺服器傳送單個請求並返回單個響應。 在流式 RPC 中，客戶端或伺服器，或雙方（雙向流式傳輸），獲取一個流讀取一系列訊息返回，然後客戶端或伺服器從返回的流中讀訊息，直到沒有更多訊息為止。

1、在 gRPC 客戶端中編寫攔截器

我們可以在 gRPC 客戶端應用程式中編寫兩種型別的攔截器：

  • UnaryClientInterceptor：UnaryClientInterceptor 攔截客戶端上一元 RPC 的執行。
  • StreamClientInterceptor：StreamClientInterceptor攔截ClientStream的建立。 它可能會返回一個自定義的 ClientStream 來攔截所有 I/O 操作。

1、UnaryClientInterceptor

為了建立 UnaryClientInterceptor，可以透過提供 UnaryClientInterceptor 函式值呼叫 WithUnaryInterceptor 函式，該函式返回一個 grpc.DialOption 指定一元 RPC 的攔截器：

func WithUnaryInterceptor(f UnaryClientInterceptor) DialOption

然後將返回的 grpc.DialOption 值用作呼叫 grpc.Dial 函式以將攔截器應用於一元 RPC 的引數。

UnaryClientInterceptor func 型別的定義如下：

type UnaryClientInterceptor func(ctx context.Context, method string, req, reply interface{}, cc *ClientConn, invoker UnaryInvoker, opts …CallOption) error

引數呼叫者是完成 RPC 的處理程式，呼叫它是攔截器的責任。 UnaryClientInterceptor 函式值提供攔截器邏輯。 這是一個實現 UnaryClientInterceptor 的示例攔截器：

clientInterceptor := func(ctx context.Context, method string, req, reply interface{}, cc *grpc.ClientConn, invoker grpc.UnaryInvoker, opts ...grpc.CallOption) error {
        start := time.Now().Unix()
        err := invoker(ctx, method, req, reply, cc, opts...)
        end := time.Now().Unix()
    // 獲取呼叫grpc的請求時長
        fmt.Println("request time duration: ", end - start)
        return err
    }

下面的函式返回一個 grpc.DialOption 值，它透過提供 UnaryClientInterceptor 函式值來呼叫 WithUnaryInterceptor 函式：

func WithUnaryInterceptorCustom() grpc.DialOption {
  clientInterceptor := func(ctx context.Context, method string, req, reply interface{}, cc *grpc.ClientConn, invoker grpc.UnaryInvoker, opts ...grpc.CallOption) error {
        start := time.Now().Unix()
        err := invoker(ctx, method, req, reply, cc, opts...)
        end := time.Now().Unix()
    // 獲取呼叫grpc的請求時長
        fmt.Println("request time duration: ", end - start)
        return err
    }
    return grpc.WithUnaryInterceptor(clientInterceptor)
}

返回的 grpc.DialOption 值用作呼叫 grpc.Dial 函式以應用攔截器的引數：

conn, err := grpc.Dial("localhost:1234", grpc.WithInsecure(), WithUnaryInterceptorCustom())

搭建簡單grpc服務可以參考這篇文章：waterflow.link/articles/1665674508...

2、StreamClientInterceptor

為了建立 StreamClientInterceptor，可以透過提供 StreamClientInterceptor 函式值呼叫 WithStreamInterceptor 函式，該函式返回一個 grpc.DialOption 指定流 RPC 的攔截器：

func WithStreamInterceptor(f StreamClientInterceptor) DialOption {
    return newFuncDialOption(func(o *dialOptions) {
        o.streamInt = f
    })
}

然後將返回的 grpc.DialOption 值用作呼叫 grpc.Dial 函式的引數，以將攔截器應用於流式 RPC。

下面是 StreamClientInterceptor func 型別的定義：

type StreamClientInterceptor func(ctx context.Context, desc *StreamDesc, cc *ClientConn, method string, streamer Streamer, opts ...CallOption) (ClientStream, error)

下面是StreamClientInterceptor的具體實現：

// 包裝下stream
// 結構體內嵌介面，初始化的時候需要賦值物件實現了該介面的所有方法
type wrappedStream struct {
    grpc.ClientStream
}

// 實現接收訊息方法，並自定義列印
func (w *wrappedStream) RecvMsg(m interface{}) error  {
    log.Printf("====== [Client Stream Interceptor] " +
        "Receive a message (Type: %T) at %v",
        m, time.Now().Format(time.RFC3339))
    return w.ClientStream.RecvMsg(m)
}

// 實現傳送訊息方法，並自定義列印
func (w *wrappedStream) SendMsg(m interface{}) error {
    log.Printf("====== [Client Stream Interceptor] " +
        "Send a message (Type: %T) at %v",
        m, time.Now().Format(time.RFC3339))
    return w.ClientStream.SendMsg(m)
}

// 初始化包裝stream
func newWrappedStream(s grpc.ClientStream) grpc.ClientStream {
    return &wrappedStream{s}
}

func WithStreamInterceptorCustom() grpc.DialOption {
    clientInterceptor := func(ctx context.Context, desc *grpc.StreamDesc, cc *grpc.ClientConn, method string, streamer grpc.Streamer, opts ...grpc.CallOption) (grpc.ClientStream, error) {
        clientStream, err := streamer(ctx, desc, cc, method, opts...)
        if err != nil {
            return nil, err
        }
    // 返回包裝後的stream
    // 這裡clientStream實現了grpc.ClientStream介面
        return newWrappedStream(clientStream), err
    }
    return grpc.WithStreamInterceptor(clientInterceptor)
}

這裡需要注意：

  1. 定義一個包裝stream結構體wrappedStream，這裡用到了結構體內嵌介面的方式，直接實現了介面的所有方法，具體可以看註釋
  2. 重寫RecvMsg和SendMsg方法
  3. WithStreamInterceptorCustom攔截器中染回包裝後的clientStream

為了將 StreamClientInterceptor 應用於流式 RPC，只需將 WithStreamInterceptor 函式返回的 grpc.DialOption 值作為呼叫 grpc.Dial 函式的引數傳遞。 您可以將 UnaryClientInterceptor 和 StreamClientInterceptor 值傳遞給 grpc.Dial 函式。

conn, err := grpc.Dial("localhost:1234", grpc.WithInsecure(), WithUnaryInterceptorCustom(), WithStreamInterceptorCustom())

完整的客戶端程式碼像，下面這樣：

package main

import (
    "context"
    "fmt"
    "google.golang.org/grpc"
    "grpcdemo/helloservice"
    "io"
    "log"
    "time"
)

func main() {
    // 連線grpc服務端，加入攔截器
    conn, err := grpc.Dial("localhost:1234", grpc.WithInsecure(), WithUnaryInterceptorCustom(), WithStreamInterceptorCustom())
    if err != nil {
        log.Fatal(err)
    }
    defer conn.Close()

  // 一元rpc
    unaryRpc(conn)
  // 流式rpc
    streamRpc(conn)

}

// 一元攔截器
func WithUnaryInterceptorCustom() grpc.DialOption {
    clientInterceptor := func(ctx context.Context, method string, req, reply interface{}, cc *grpc.ClientConn, invoker grpc.UnaryInvoker, opts ...grpc.CallOption) error {
        start := time.Now().Unix()
        err := invoker(ctx, method, req, reply, cc, opts...)
        end := time.Now().Unix()
        fmt.Println("invoker request time duration: ", end - start)
        return err
    }
    return grpc.WithUnaryInterceptor(clientInterceptor)
}

type wrappedStream struct {
    grpc.ClientStream
}

func (w *wrappedStream) RecvMsg(m interface{}) error  {
    log.Printf("====== [Client Stream Interceptor] " +
        "Receive a message (Type: %T) at %v",
        m, time.Now().Format(time.RFC3339))
    return w.ClientStream.RecvMsg(m)
}

func (w *wrappedStream) SendMsg(m interface{}) error {
    log.Printf("====== [Client Stream Interceptor] " +
        "Send a message (Type: %T) at %v",
        m, time.Now().Format(time.RFC3339))
    return w.ClientStream.SendMsg(m)
}

func newWrappedStream(s grpc.ClientStream) grpc.ClientStream {
    return &wrappedStream{s}
}

// 流式攔截器
func WithStreamInterceptorCustom() grpc.DialOption {
    clientInterceptor := func(ctx context.Context, desc *grpc.StreamDesc, cc *grpc.ClientConn, method string, streamer grpc.Streamer, opts ...grpc.CallOption) (grpc.ClientStream, error) {
        clientStream, err := streamer(ctx, desc, cc, method, opts...)
        if err != nil {
            return nil, err
        }
        return newWrappedStream(clientStream), err
    }
    return grpc.WithStreamInterceptor(clientInterceptor)
}

func unaryRpc(conn *grpc.ClientConn) {
    client := helloservice.NewHelloServiceClient(conn)
    reply, err := client.Hello(context.Background(), &helloservice.String{Value: "hello"})
    if err != nil {
        log.Fatal(err)
    }
    log.Println("unaryRpc recv: ", reply.Value)
}

func streamRpc(conn *grpc.ClientConn) {
    client := helloservice.NewHelloServiceClient(conn)
    stream, err := client.Channel(context.Background())
    if err != nil {
        log.Fatal(err)
    }

    go func() {
        for {
            if err := stream.Send(&helloservice.String{Value: "hi"}); err != nil {
                log.Fatal(err)
            }
            time.Sleep(time.Second)
        }
    }()

    for {
        recv, err := stream.Recv()
        if err != nil {
            if err == io.EOF {
                break
            }
            log.Fatal(err)
        }

        fmt.Println("streamRpc recv: ", recv.Value)

    }
}

可以結合我之前的文章，把本期的程式碼加進去執行除錯

(搭建簡單grpc服務可以參考這篇文章：waterflow.link/articles/1665674508...)

執行效果如下：

go run helloclient/main.go
invoker request time duration:  1
2022/10/14 23:17:35 unaryRpc recv:  hello:hello
2022/10/14 23:17:35 ====== [Client Stream Interceptor] Receive a message (Type: *helloservice.String) at 2022-10-14T23:17:35+08:00
2022/10/14 23:17:35 ====== [Client Stream Interceptor] Send a message (Type: *helloservice.String) at 2022-10-14T23:17:35+08:00
2022/10/14 23:17:36 ====== [Client Stream Interceptor] Send a message (Type: *helloservice.String) at 2022-10-14T23:17:36+08:00
streamRpc recv:  hello:hi
2022/10/14 23:17:36 ====== [Client Stream Interceptor] Receive a message (Type: *helloservice.String) at 2022-10-14T23:17:36+08:00
2022/10/14 23:17:37 ====== [Client Stream Interceptor] Send a message (Type: *helloservice.String) at 2022-10-14T23:17:37+08:00
streamRpc recv:  hello:hi
2022/10/14 23:17:37 ====== [Client Stream Interceptor] Receive a message (Type: *helloservice.String) at 2022-10-14T23:17:37+08:00
2022/10/14 23:17:38 ====== [Client Stream Interceptor] Send a message (Type: *helloservice.String) at 2022-10-14T23:17:38+08:00
streamRpc recv:  hello:hi
2022/10/14 23:17:38 ====== [Client Stream Interceptor] Receive a message (Type: *helloservice.String) at 2022-10-14T23:17:38+08:00

2、在 gRPC 客戶端中編寫攔截器

和 gRPC 客戶端應用程式一樣，gRPC 伺服器應用程式提供兩種型別的攔截器：

  • UnaryServerInterceptor：提供了一個鉤子來攔截伺服器上一元 RPC 的執行。
  • StreamServerInterceptor：提供了一個鉤子來攔截伺服器上流式 RPC 的執行。

1、UnaryServerInterceptor

為了建立 UnaryServerInterceptor，可以透過提供 UnaryServerInterceptor 函式值作為引數呼叫 UnaryInterceptor 函式，該引數返回為伺服器設定 UnaryServerInterceptor 的 grpc.ServerOption 值。

func UnaryInterceptor(i UnaryServerInterceptor) ServerOption {
    return newFuncServerOption(func(o *serverOptions) {
        if o.unaryInt != nil {
            panic("The unary server interceptor was already set and may not be reset.")
        }
        o.unaryInt = i
    })
}

然後使用返回的 grpc.ServerOption 值作為引數提供給 grpc.NewServer 函式以註冊為 UnaryServerInterceptor。

UnaryServerInterceptor 函式的定義如下：

func(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (resp interface{}, err error)

引數 info 包含了這個 RPC 的所有資訊，攔截器可以對其進行操作。 而handler是服務方法實現的包裝器。 攔截器負責呼叫處理程式來完成 RPC。

1、定義一個服務端的鑑權一元攔截器
func ServerUnaryInterceptorCustom() grpc.ServerOption {
    serverInterceptor := func(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (resp interface{}, err error) {
        start := time.Now()

    // 如果是非登入請求，需要驗證token
        if info.FullMethod != "/helloservice.HelloService/Login" {
            if err := authorize(ctx); err != nil {
                return nil, err
            }
        }

        h, err := handler(ctx, req)

        log.Printf("Request - Method:%s\tDuration:%s\tError:%v\n",
            info.FullMethod,
            time.Since(start),
            err)
        return h, err
    }
    return grpc.UnaryInterceptor(serverInterceptor)
}

// authorize 從Metadata中獲取token並校驗是否合法
func authorize(ctx context.Context) error {
  // 從context中獲取metadata
    md, ok := metadata.FromIncomingContext(ctx)
    if !ok {
        return status.Errorf(codes.InvalidArgument, "Retrieving metadata is failed")
    }

    authHeader, ok := md["authorization"]
    if !ok {
        return status.Errorf(codes.Unauthenticated, "Authorization token is not supplied")
    }

    token := authHeader[0]
    // validateToken function validates the token
    err := validateToken(token)

    if err != nil {
        return status.Errorf(codes.Unauthenticated, err.Error())
    }
    return nil
}

func validateToken(token string) error {
    // 校驗token
    return nil
}

我們可以看下我們定義的一元攔截器的執行流程：

  1. 首先進來之後我們判斷如果是登入請求，直接轉發請求，並列印日誌
  2. 如果是非登入請求，需要驗證token，呼叫authorize方法
  3. 在authorize方法中，會從context中獲取metadata後設資料，然後解析獲取token並驗證

請注意，前面程式碼塊中的攔截器邏輯使用包 google.golang.org/grpc/codes 和 google.golang.org/grpc/status。

2、grpc客戶端傳入token

gRPC 支援在客戶端和伺服器之間使用 Context 值傳送後設資料。 google.golang.org/grpc/metadata 包提供了後設資料的功能。

其中MD型別是一個k-v的map，想下面這樣

type MD map[string][]string

下面我們在客戶端編寫向服務端傳送token的程式碼，我們修改下客戶端的unaryRpc，構造包含authorization的metadata：

func unaryRpc(conn *grpc.ClientConn) {
    client := helloservice.NewHelloServiceClient(conn)
    ctx := context.Background()
  // 構造後設資料，並返回MD型別的結構
    md := metadata.Pairs("authorization", "mytoken")
  // 後設資料塞入context並返回新的context
    ctx = metadata.NewOutgoingContext(ctx, md)
    reply, err := client.Hello(ctx, &helloservice.String{Value: "hello"})
    if err != nil {
        log.Fatal(err)
    }
    log.Println("unaryRpc recv: ", reply.Value)
}

這樣後設資料的資訊就會跟著context傳送到grpc服務端。

接著我們在服務端grpc中修改如下程式碼，加入一行日誌：

func validateToken(token string) error {
    log.Printf("get the token: %s \n", token)
    // 校驗token
    return nil
}
3、執行服務

我們重新執行下grpc服務端程式，然後執行下客戶端程式碼，可以看到token傳過來了：

go run helloservice/main/main.go
2022/10/15 20:36:04 server started...
2022/10/15 20:36:08 get the token: mytoken 
2022/10/15 20:36:09 Request - Method:/helloservice.HelloService/Hello   Duration:1.001216763s   Error:<nil>

2、StreamServerInterceptor

為了建立 StreamServerInterceptor，透過提供 StreamServerInterceptor func 值作為引數呼叫 StreamInterceptor 函式，該引數返回為伺服器設定 StreamServerInterceptor 的 grpc.ServerOption 值。

func StreamInterceptor(i StreamServerInterceptor) ServerOption {
    return newFuncServerOption(func(o *serverOptions) {
        if o.streamInt != nil {
            panic("The stream server interceptor was already set and may not be reset.")
        }
        o.streamInt = i
    })
}

然後使用返回的 grpc.ServerOption 值作為引數提供給 grpc.NewServer 函式以註冊為 UnaryServerInterceptor。

下面是 StreamServerInterceptor func 型別的定義：

type StreamServerInterceptor func(srv interface{}, ss ServerStream, info *StreamServerInfo, handler StreamHandler) error

我們看下服務端流式攔截器的具體例子：

type wrappedStream struct {
    grpc.ServerStream
}

func (w *wrappedStream) RecvMsg(m interface{}) error  {
    log.Printf("====== [Server Stream Interceptor] " +
        "Receive a message (Type: %T) at %v",
        m, time.Now().Format(time.RFC3339))
    return w.ServerStream.RecvMsg(m)
}

func (w *wrappedStream) SendMsg(m interface{}) error {
    log.Printf("====== [Server Stream Interceptor] " +
        "Send a message (Type: %T) at %v",
        m, time.Now().Format(time.RFC3339))
    return w.ServerStream.SendMsg(m)
}

func newWrappedStream(s grpc.ServerStream) grpc.ServerStream {
    return &wrappedStream{s}
}

func ServerStreamInterceptorCustom() grpc.ServerOption {
    serverInterceptor := func(srv interface{}, ss grpc.ServerStream, info *grpc.StreamServerInfo, handler grpc.StreamHandler) error {
        return handler(srv, newWrappedStream(ss))
    }
    return grpc.StreamInterceptor(serverInterceptor)
}

上面服務端流式攔截器程式碼可參考客戶端流式攔截器的程式碼，基本差不多。

3、多攔截器

go-grpc在v1.28.0之前是不支援多個攔截器。但是可以使用一些第三方的包，攔截器連結允應用多個攔截器。

v1.28.0之後已經可以支援多個攔截器，我們修改下服務端程式碼如下：

...

unaryInterceptors := []grpc.ServerOption {
        ServerUnaryInterceptorCustom(),
        ServerStreamInterceptorCustom(),
    }
grpcServer := grpc.NewServer(unaryInterceptors...)

...
本作品採用《CC 協議》，轉載必須註明作者和本文連結

相關文章