grpc中的攔截器

slowquery發表於2022-10-16

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 協議》,轉載必須註明作者和本文連結

相關文章