github.com/ava-labs/avalanchego@v1.11.11/vms/rpcchainvm/ghttp/gresponsewriter/writer_client.go (about)

     1  // Copyright (C) 2019-2024, Ava Labs, Inc. All rights reserved.
     2  // See the file LICENSE for licensing terms.
     3  
     4  package gresponsewriter
     5  
     6  import (
     7  	"bufio"
     8  	"context"
     9  	"net"
    10  	"net/http"
    11  
    12  	"google.golang.org/protobuf/types/known/emptypb"
    13  
    14  	"github.com/ava-labs/avalanchego/vms/rpcchainvm/ghttp/gconn"
    15  	"github.com/ava-labs/avalanchego/vms/rpcchainvm/ghttp/greader"
    16  	"github.com/ava-labs/avalanchego/vms/rpcchainvm/ghttp/gwriter"
    17  	"github.com/ava-labs/avalanchego/vms/rpcchainvm/grpcutils"
    18  
    19  	responsewriterpb "github.com/ava-labs/avalanchego/proto/pb/http/responsewriter"
    20  	readerpb "github.com/ava-labs/avalanchego/proto/pb/io/reader"
    21  	writerpb "github.com/ava-labs/avalanchego/proto/pb/io/writer"
    22  	connpb "github.com/ava-labs/avalanchego/proto/pb/net/conn"
    23  )
    24  
    25  var (
    26  	_ http.ResponseWriter = (*Client)(nil)
    27  	_ http.Flusher        = (*Client)(nil)
    28  	_ http.Hijacker       = (*Client)(nil)
    29  )
    30  
    31  // Client is an http.ResponseWriter that talks over RPC.
    32  type Client struct {
    33  	client responsewriterpb.WriterClient
    34  	header http.Header
    35  }
    36  
    37  // NewClient returns a response writer connected to a remote response writer
    38  func NewClient(header http.Header, client responsewriterpb.WriterClient) *Client {
    39  	return &Client{
    40  		client: client,
    41  		header: header,
    42  	}
    43  }
    44  
    45  func (c *Client) Header() http.Header {
    46  	return c.header
    47  }
    48  
    49  func (c *Client) Write(payload []byte) (int, error) {
    50  	req := &responsewriterpb.WriteRequest{
    51  		Headers: make([]*responsewriterpb.Header, 0, len(c.header)),
    52  		Payload: payload,
    53  	}
    54  	for key, values := range c.header {
    55  		req.Headers = append(req.Headers, &responsewriterpb.Header{
    56  			Key:    key,
    57  			Values: values,
    58  		})
    59  	}
    60  	resp, err := c.client.Write(context.Background(), req)
    61  	if err != nil {
    62  		return 0, err
    63  	}
    64  	return int(resp.Written), nil
    65  }
    66  
    67  func (c *Client) WriteHeader(statusCode int) {
    68  	req := &responsewriterpb.WriteHeaderRequest{
    69  		Headers:    make([]*responsewriterpb.Header, 0, len(c.header)),
    70  		StatusCode: int32(statusCode),
    71  	}
    72  	for key, values := range c.header {
    73  		req.Headers = append(req.Headers, &responsewriterpb.Header{
    74  			Key:    key,
    75  			Values: values,
    76  		})
    77  	}
    78  	// TODO: Is there a way to handle the error here?
    79  	_, _ = c.client.WriteHeader(context.Background(), req)
    80  }
    81  
    82  func (c *Client) Flush() {
    83  	// TODO: is there a way to handle the error here?
    84  	_, _ = c.client.Flush(context.Background(), &emptypb.Empty{})
    85  }
    86  
    87  type addr struct {
    88  	network string
    89  	str     string
    90  }
    91  
    92  func (a *addr) Network() string {
    93  	return a.network
    94  }
    95  
    96  func (a *addr) String() string {
    97  	return a.str
    98  }
    99  
   100  func (c *Client) Hijack() (net.Conn, *bufio.ReadWriter, error) {
   101  	resp, err := c.client.Hijack(context.Background(), &emptypb.Empty{})
   102  	if err != nil {
   103  		return nil, nil, err
   104  	}
   105  
   106  	clientConn, err := grpcutils.Dial(resp.ServerAddr)
   107  	if err != nil {
   108  		return nil, nil, err
   109  	}
   110  
   111  	conn := gconn.NewClient(
   112  		connpb.NewConnClient(clientConn),
   113  		&addr{
   114  			network: resp.LocalNetwork,
   115  			str:     resp.LocalString,
   116  		},
   117  		&addr{
   118  			network: resp.RemoteNetwork,
   119  			str:     resp.RemoteString,
   120  		},
   121  		clientConn,
   122  	)
   123  
   124  	reader := greader.NewClient(readerpb.NewReaderClient(clientConn))
   125  	writer := gwriter.NewClient(writerpb.NewWriterClient(clientConn))
   126  
   127  	readWriter := bufio.NewReadWriter(
   128  		bufio.NewReader(reader),
   129  		bufio.NewWriter(writer),
   130  	)
   131  
   132  	return conn, readWriter, nil
   133  }