如何为grpc-gateway扩展http转proto.Message的自定义解码逻辑？

我们业务里需要对所有Protobuf的string类型字段执行strings.TrimSpace操作，于是写了个Protobuf插件，给每个生成的结构体加了两个函数：

• trimAll()：遍历结构体的string字段做修剪
• UnmarshalJSON()：原以为grpc-gateway会用这个方法把http.Request.Body转成proto.Message，反序列化后自动触发修剪

生成的示例代码如下：

func (m *GetLogbookCall_Request) trimAll() {
    if m == nil {
        return
    }
    m.Id = strings.TrimSpace(m.Id)
}

func (m *GetLogbookCall_Request) UnmarshalJSON(data []byte) error {
    err := proto.Unmarshal(data, m)
    if err != nil {
        return err
    }

    m.trimAll()

    return nil
}

但实际跑起来发现这个方案没用。现在想找个可行的办法，在grpc-gateway把http.Request.Body转成proto.Message的过程中，插入自定义的字段修剪逻辑。

附grpc-gateway生成的请求处理函数示例：

func request_AccountService_UpdatePassword_0(ctx context.Context, marshaler runtime.Marshaler, client AccountServiceClient, req *http.Request, pathParams map[string]string) (proto.Message, runtime.ServerMetadata, error) {
    var protoReq account.UpdatePasswordCall_Request
    var metadata runtime.ServerMetadata

    if err := marshaler.NewDecoder(req.Body).Decode(&protoReq); err != nil && err != io.EOF {
        return nil, metadata, status.Errorf(codes.InvalidArgument, "%v", err)
    }

    msg, err := client.UpdatePassword(ctx, &protoReq, grpc.Header(&metadata.HeaderMD), grpc.Trailer(&metadata.TrailerMD))
    return msg, metadata, err

}

方法1：自定义Marshaler #

grpc-gateway默认用的runtime.JSONPb marshaler不会调用结构体自己的UnmarshalJSON方法，而是自己实现了反序列化逻辑。所以可以自定义一个marshaler，继承默认的JSONPb，重写NewDecoder方法，在反序列化完成后自动调用trimAll()：

type TrimJSONMarshaler struct {
    runtime.JSONPb
}

func (m *TrimJSONMarshaler) NewDecoder(r io.Reader) runtime.Decoder {
    dec := m.JSONPb.NewDecoder(r)
    return &trimDecoder{dec: dec}
}

type trimDecoder struct {
    dec runtime.Decoder
}

func (d *trimDecoder) Decode(v interface{}) error {
    if err := d.dec.Decode(v); err != nil {
        return err
    }
    // 尝试调用对象的trimAll方法，存在就执行
    if trimmer, ok := v.(interface{ trimAll() }); ok {
        trimmer.trimAll()
    }
    return nil
}

然后在grpc-gateway的注册代码里替换默认marshaler：

mux := runtime.NewServeMux(
    runtime.WithMarshalerOption(runtime.MIMEWildcard, &TrimJSONMarshaler{
        JSONPb: runtime.JSONPb{
            MarshalOptions: protojson.MarshalOptions{
                UseProtoNames: true,
            },
            UnmarshalOptions: protojson.UnmarshalOptions{
                DiscardUnknown: true,
            },
        },
    }),
)
// 后续注册服务的逻辑不变

方法2：修改grpc-gateway生成的代码（基于现有插件） #

既然你已经有写Protobuf插件的基础，可以直接修改插件逻辑，在grpc-gateway生成的请求处理函数里，紧接Decode操作之后插入trimAll()调用。比如把示例代码改成这样：

func request_AccountService_UpdatePassword_0(ctx context.Context, marshaler runtime.Marshaler, client AccountServiceClient, req *http.Request, pathParams map[string]string) (proto.Message, runtime.ServerMetadata, error) {
    var protoReq account.UpdatePasswordCall_Request
    var metadata runtime.ServerMetadata

    if err := marshaler.NewDecoder(req.Body).Decode(&protoReq); err != nil && err != io.EOF {
        return nil, metadata, status.Errorf(codes.InvalidArgument, "%v", err)
    }

    // 新增的修剪逻辑
    protoReq.trimAll()

    msg, err := client.UpdatePassword(ctx, &protoReq, grpc.Header(&metadata.HeaderMD), grpc.Trailer(&metadata.TrailerMD))
    return msg, metadata, err

}

这种方法更直接，不用改marshaler，只要插件能定位到生成的请求处理函数中Decode后的位置，插入调用代码就行。

方法3：结合Protobuf验证器插件 #

如果你的项目已经在用protoc-gen-validate这类Protobuf验证插件，可以添加自定义验证规则，在验证阶段自动修剪string字段。不过这种方法需要额外配置验证规则，适合同时需要字段校验的场景。

内容的提问来源于stack exchange，提问作者Alexandr Rutkowski