go.uber.org/yarpc@v1.72.1/serialize/serialize.go (about)

     1  // Copyright (c) 2022 Uber Technologies, Inc.
     2  //
     3  // Permission is hereby granted, free of charge, to any person obtaining a copy
     4  // of this software and associated documentation files (the "Software"), to deal
     5  // in the Software without restriction, including without limitation the rights
     6  // to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
     7  // copies of the Software, and to permit persons to whom the Software is
     8  // furnished to do so, subject to the following conditions:
     9  //
    10  // The above copyright notice and this permission notice shall be included in
    11  // all copies or substantial portions of the Software.
    12  //
    13  // THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
    14  // IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
    15  // FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
    16  // AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
    17  // LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
    18  // OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
    19  // THE SOFTWARE.
    20  
    21  package serialize
    22  
    23  import (
    24  	"bytes"
    25  	"errors"
    26  	"fmt"
    27  	"io/ioutil"
    28  
    29  	"github.com/opentracing/opentracing-go"
    30  	"go.uber.org/thriftrw/protocol/binary"
    31  	"go.uber.org/thriftrw/wire"
    32  	"go.uber.org/yarpc/api/transport"
    33  	"go.uber.org/yarpc/serialize/internal"
    34  )
    35  
    36  // version indicates which underlying serialization method will be used
    37  // '0' indicates:
    38  // 		thrift serialization (request) + jaeger.binary format (ctx/tracing)
    39  const version = byte(0)
    40  
    41  // ToBytes encodes an opentracing.SpanContext and transport.Request into bytes
    42  func ToBytes(tracer opentracing.Tracer, spanContext opentracing.SpanContext, req *transport.Request) ([]byte, error) {
    43  	spanBytes, err := spanContextToBytes(tracer, spanContext)
    44  	if err != nil {
    45  		return nil, err
    46  	}
    47  
    48  	body, err := ioutil.ReadAll(req.Body)
    49  	if err != nil {
    50  		return nil, err
    51  	}
    52  
    53  	rpc := internal.RPC{
    54  		SpanContext: spanBytes,
    55  
    56  		CallerName:      req.Caller,
    57  		ServiceName:     req.Service,
    58  		Encoding:        string(req.Encoding),
    59  		Procedure:       req.Procedure,
    60  		Headers:         req.Headers.Items(),
    61  		ShardKey:        &req.ShardKey,
    62  		RoutingKey:      &req.RoutingKey,
    63  		RoutingDelegate: &req.RoutingDelegate,
    64  		Body:            body,
    65  	}
    66  
    67  	wireValue, err := rpc.ToWire()
    68  	if err != nil {
    69  		return nil, err
    70  	}
    71  
    72  	var writer bytes.Buffer
    73  	// use the first byte to version the serialization
    74  	if err := writer.WriteByte(version); err != nil {
    75  		return nil, err
    76  	}
    77  	err = binary.Default.Encode(wireValue, &writer)
    78  	return writer.Bytes(), err
    79  }
    80  
    81  // FromBytes decodes bytes into a opentracing.SpanContext and transport.Request
    82  func FromBytes(tracer opentracing.Tracer, request []byte) (opentracing.SpanContext, *transport.Request, error) {
    83  	if len(request) <= 1 {
    84  		return nil, nil, errors.New("cannot deserialize empty request")
    85  	}
    86  
    87  	// check valid thrift serialization byte
    88  	if request[0] != 0 {
    89  		return nil, nil,
    90  			fmt.Errorf(
    91  				"unsupported YARPC serialization version '%v' found during deserialization",
    92  				request[0])
    93  	}
    94  
    95  	reader := bytes.NewReader(request[1:])
    96  	wireValue, err := binary.Default.Decode(reader, wire.TStruct)
    97  	if err != nil {
    98  		return nil, nil, err
    99  	}
   100  
   101  	var rpc internal.RPC
   102  	if err = rpc.FromWire(wireValue); err != nil {
   103  		return nil, nil, err
   104  	}
   105  
   106  	req := transport.Request{
   107  		Caller:    rpc.CallerName,
   108  		Service:   rpc.ServiceName,
   109  		Encoding:  transport.Encoding(rpc.Encoding),
   110  		Procedure: rpc.Procedure,
   111  		Headers:   transport.HeadersFromMap(rpc.Headers),
   112  		Body:      bytes.NewBuffer(rpc.Body),
   113  	}
   114  
   115  	if rpc.ShardKey != nil {
   116  		req.ShardKey = *rpc.ShardKey
   117  	}
   118  	if rpc.RoutingKey != nil {
   119  		req.RoutingKey = *rpc.RoutingKey
   120  	}
   121  	if rpc.RoutingDelegate != nil {
   122  		req.RoutingDelegate = *rpc.RoutingDelegate
   123  	}
   124  
   125  	spanContext, err := spanContextFromBytes(tracer, rpc.SpanContext)
   126  	if err != nil {
   127  		return nil, nil, err
   128  	}
   129  
   130  	return spanContext, &req, nil
   131  }
   132  
   133  func spanContextToBytes(tracer opentracing.Tracer, spanContext opentracing.SpanContext) ([]byte, error) {
   134  	carrier := bytes.NewBuffer([]byte{})
   135  	err := tracer.Inject(spanContext, opentracing.Binary, carrier)
   136  	return carrier.Bytes(), err
   137  }
   138  
   139  func spanContextFromBytes(tracer opentracing.Tracer, spanContextBytes []byte) (opentracing.SpanContext, error) {
   140  	carrier := bytes.NewBuffer(spanContextBytes)
   141  	spanContext, err := tracer.Extract(opentracing.Binary, carrier)
   142  	// If no SpanContext was given, we return nil instead of erroring
   143  	// transport.ExtractOpenTracingSpan() safely accepts nil
   144  	if err == opentracing.ErrSpanContextNotFound {
   145  		return nil, nil
   146  	}
   147  	return spanContext, err
   148  }