go.uber.org/yarpc@v1.72.1/internal/firstoutboundmiddleware/middleware.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 firstoutboundmiddleware annotates every outbound request with
    22  // metadata like the request transport protocol.
    23  // These metadata must be avilable to all subsequent middleware.
    24  package firstoutboundmiddleware
    25  
    26  import (
    27  	"context"
    28  
    29  	"go.uber.org/yarpc/api/encoding"
    30  	"go.uber.org/yarpc/api/middleware"
    31  	"go.uber.org/yarpc/api/transport"
    32  )
    33  
    34  var (
    35  	_ middleware.UnaryOutbound  = (*Middleware)(nil)
    36  	_ middleware.StreamOutbound = (*Middleware)(nil)
    37  	_ middleware.OnewayOutbound = (*Middleware)(nil)
    38  )
    39  
    40  // Middleware is the first middleware that MUST be executed in the chain of
    41  // TransportOutboundMiddleware.
    42  type Middleware struct{}
    43  
    44  // New returns middleware to begin any YARPC outbound middleware chain.
    45  func New() *Middleware {
    46  	return &Middleware{}
    47  }
    48  
    49  // Call implements middleware.UnaryOutbound.
    50  func (m *Middleware) Call(ctx context.Context, req *transport.Request, next transport.UnaryOutbound) (*transport.Response, error) {
    51  	update(ctx, req, next)
    52  	return next.Call(ctx, req)
    53  }
    54  
    55  // CallOneway implements middleware.OnewayOutbound.
    56  func (m *Middleware) CallOneway(ctx context.Context, req *transport.Request, next transport.OnewayOutbound) (transport.Ack, error) {
    57  	update(ctx, req, next)
    58  	return next.CallOneway(ctx, req)
    59  }
    60  
    61  // CallStream implements middleware.StreamOutbound.
    62  func (m *Middleware) CallStream(ctx context.Context, req *transport.StreamRequest, next transport.StreamOutbound) (*transport.ClientStream, error) {
    63  	updateStream(ctx, req, next)
    64  	return next.CallStream(ctx, req)
    65  }
    66  
    67  func update(ctx context.Context, req *transport.Request, out transport.Outbound) {
    68  	// TODO(apeatsbond): Setting environment headers and unique IDs should live
    69  	// here too (T1860945).
    70  
    71  	// Reset the transport field to the current outbound transport.
    72  	// Request forwarding in transport layer proxies needs this when copying
    73  	// requests to a different outbound type.
    74  	if namer, ok := out.(transport.Namer); ok {
    75  		req.Transport = namer.TransportName()
    76  	}
    77  
    78  	// Update the caller procedure to the current procedure making this request
    79  	call := encoding.CallFromContext(ctx)
    80  	if call != nil {
    81  		req.CallerProcedure = call.Procedure()
    82  	}
    83  }
    84  
    85  func updateStream(ctx context.Context, req *transport.StreamRequest, out transport.Outbound) {
    86  	// TODO(apeatsbond): Setting environment headers and unique IDs should live
    87  	// here too (T1860945).
    88  
    89  	// Reset the transport field to the current outbound transport.
    90  	// Request forwarding in transport layer proxies needs this when copying
    91  	// requests to a different outbound type.
    92  	if namer, ok := out.(transport.Namer); ok {
    93  		req.Meta.Transport = namer.TransportName()
    94  	}
    95  
    96  	// Update the caller procedure to the current procedure making this request
    97  	call := encoding.CallFromContext(ctx)
    98  	if call != nil {
    99  		req.Meta.CallerProcedure = call.Procedure()
   100  	}
   101  }