github.com/cloudwego/kitex@v0.9.0/pkg/generic/jsonpb_codec.go (about)

     1  /*
     2   * Copyright 2023 CloudWeGo Authors
     3   *
     4   * Licensed under the Apache License, Version 2.0 (the "License");
     5   * you may not use this file except in compliance with the License.
     6   * You may obtain a copy of the License at
     7   *
     8   *     http://www.apache.org/licenses/LICENSE-2.0
     9   *
    10   * Unless required by applicable law or agreed to in writing, software
    11   * distributed under the License is distributed on an "AS IS" BASIS,
    12   * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    13   * See the License for the specific language governing permissions and
    14   * limitations under the License.
    15   */
    16  
    17  package generic
    18  
    19  import (
    20  	"context"
    21  	"fmt"
    22  	"sync/atomic"
    23  
    24  	"github.com/cloudwego/dynamicgo/conv"
    25  	dproto "github.com/cloudwego/dynamicgo/proto"
    26  
    27  	"github.com/cloudwego/kitex/pkg/generic/proto"
    28  	"github.com/cloudwego/kitex/pkg/remote"
    29  	"github.com/cloudwego/kitex/pkg/remote/codec"
    30  	"github.com/cloudwego/kitex/pkg/remote/codec/perrors"
    31  	"github.com/cloudwego/kitex/pkg/serviceinfo"
    32  )
    33  
    34  var (
    35  	_ remote.PayloadCodec = &jsonPbCodec{}
    36  	_ Closer              = &jsonPbCodec{}
    37  )
    38  
    39  type jsonPbCodec struct {
    40  	svcDsc           atomic.Value // *idl
    41  	provider         PbDescriptorProviderDynamicGo
    42  	codec            remote.PayloadCodec
    43  	opts             *Options
    44  	convOpts         conv.Options // used for dynamicgo conversion
    45  	dynamicgoEnabled bool         // currently set to true by default
    46  }
    47  
    48  func newJsonPbCodec(p PbDescriptorProviderDynamicGo, codec remote.PayloadCodec, opts *Options) (*jsonPbCodec, error) {
    49  	svc := <-p.Provide()
    50  	c := &jsonPbCodec{codec: codec, provider: p, opts: opts, dynamicgoEnabled: true}
    51  	convOpts := opts.dynamicgoConvOpts
    52  	c.convOpts = convOpts
    53  
    54  	c.svcDsc.Store(svc)
    55  	go c.update()
    56  	return c, nil
    57  }
    58  
    59  func (c *jsonPbCodec) update() {
    60  	for {
    61  		svc, ok := <-c.provider.Provide()
    62  		if !ok {
    63  			return
    64  		}
    65  		c.svcDsc.Store(svc)
    66  	}
    67  }
    68  
    69  func (c *jsonPbCodec) Marshal(ctx context.Context, msg remote.Message, out remote.ByteBuffer) error {
    70  	method := msg.RPCInfo().Invocation().MethodName()
    71  	if method == "" {
    72  		return perrors.NewProtocolErrorWithMsg("empty methodName in protobuf Marshal")
    73  	}
    74  	if msg.MessageType() == remote.Exception {
    75  		return c.codec.Marshal(ctx, msg, out)
    76  	}
    77  
    78  	pbSvc := c.svcDsc.Load().(*dproto.ServiceDescriptor)
    79  
    80  	wm, err := proto.NewWriteJSON(pbSvc, method, msg.RPCRole() == remote.Client, &c.convOpts)
    81  	if err != nil {
    82  		return err
    83  	}
    84  
    85  	msg.Data().(WithCodec).SetCodec(wm)
    86  
    87  	return c.codec.Marshal(ctx, msg, out)
    88  }
    89  
    90  func (c *jsonPbCodec) Unmarshal(ctx context.Context, msg remote.Message, in remote.ByteBuffer) error {
    91  	if err := codec.NewDataIfNeeded(serviceinfo.GenericMethod, msg); err != nil {
    92  		return err
    93  	}
    94  
    95  	pbSvc := c.svcDsc.Load().(*dproto.ServiceDescriptor)
    96  
    97  	wm, err := proto.NewReadJSON(pbSvc, msg.RPCRole() == remote.Client, &c.convOpts)
    98  	if err != nil {
    99  		return err
   100  	}
   101  
   102  	msg.Data().(WithCodec).SetCodec(wm)
   103  
   104  	return c.codec.Unmarshal(ctx, msg, in)
   105  }
   106  
   107  func (c *jsonPbCodec) getMethod(req interface{}, method string) (*Method, error) {
   108  	fnSvc := c.svcDsc.Load().(*dproto.ServiceDescriptor).LookupMethodByName(method)
   109  	if fnSvc == nil {
   110  		return nil, fmt.Errorf("missing method: %s in service", method)
   111  	}
   112  
   113  	// protobuf does not have oneway methods
   114  	return &Method{method, false}, nil
   115  }
   116  
   117  func (c *jsonPbCodec) Name() string {
   118  	return "JSONPb"
   119  }
   120  
   121  func (c *jsonPbCodec) Close() error {
   122  	return c.provider.Close()
   123  }