github.com/muhammadn/cortex@v1.9.1-0.20220510110439-46bb7000d03d/pkg/frontend/transport/roundtripper.go (about)

     1  package transport
     2  
     3  import (
     4  	"bytes"
     5  	"context"
     6  	"io"
     7  	"io/ioutil"
     8  	"net/http"
     9  
    10  	"github.com/weaveworks/common/httpgrpc"
    11  	"github.com/weaveworks/common/httpgrpc/server"
    12  )
    13  
    14  // GrpcRoundTripper is similar to http.RoundTripper, but works with HTTP requests converted to protobuf messages.
    15  type GrpcRoundTripper interface {
    16  	RoundTripGRPC(context.Context, *httpgrpc.HTTPRequest) (*httpgrpc.HTTPResponse, error)
    17  }
    18  
    19  func AdaptGrpcRoundTripperToHTTPRoundTripper(r GrpcRoundTripper) http.RoundTripper {
    20  	return &grpcRoundTripperAdapter{roundTripper: r}
    21  }
    22  
    23  // This adapter wraps GrpcRoundTripper and converted it into http.RoundTripper
    24  type grpcRoundTripperAdapter struct {
    25  	roundTripper GrpcRoundTripper
    26  }
    27  
    28  type buffer struct {
    29  	buff []byte
    30  	io.ReadCloser
    31  }
    32  
    33  func (b *buffer) Bytes() []byte {
    34  	return b.buff
    35  }
    36  
    37  func (a *grpcRoundTripperAdapter) RoundTrip(r *http.Request) (*http.Response, error) {
    38  	req, err := server.HTTPRequest(r)
    39  	if err != nil {
    40  		return nil, err
    41  	}
    42  
    43  	resp, err := a.roundTripper.RoundTripGRPC(r.Context(), req)
    44  	if err != nil {
    45  		return nil, err
    46  	}
    47  
    48  	httpResp := &http.Response{
    49  		StatusCode:    int(resp.Code),
    50  		Body:          &buffer{buff: resp.Body, ReadCloser: ioutil.NopCloser(bytes.NewReader(resp.Body))},
    51  		Header:        http.Header{},
    52  		ContentLength: int64(len(resp.Body)),
    53  	}
    54  	for _, h := range resp.Headers {
    55  		httpResp.Header[h.Key] = h.Values
    56  	}
    57  	return httpResp, nil
    58  }