github.com/laof/lite-speed-test@v0.0.0-20230930011949-1f39b7037845/transport/vmess/h2.go (about)

     1  package vmess
     2  
     3  import (
     4  	"io"
     5  	"math/rand"
     6  	"net"
     7  	"net/http"
     8  	"net/url"
     9  
    10  	"golang.org/x/net/http2"
    11  )
    12  
    13  type h2Conn struct {
    14  	net.Conn
    15  	*http2.ClientConn
    16  	pwriter *io.PipeWriter
    17  	res     *http.Response
    18  	cfg     *H2Config
    19  }
    20  
    21  type H2Config struct {
    22  	Hosts []string
    23  	Path  string
    24  }
    25  
    26  func (hc *h2Conn) establishConn() error {
    27  	preader, pwriter := io.Pipe()
    28  
    29  	host := hc.cfg.Hosts[rand.Intn(len(hc.cfg.Hosts))]
    30  	path := hc.cfg.Path
    31  	// TODO: connect use VMess Host instead of H2 Host
    32  	req := http.Request{
    33  		Method: "PUT",
    34  		Host:   host,
    35  		URL: &url.URL{
    36  			Scheme: "https",
    37  			Host:   host,
    38  			Path:   path,
    39  		},
    40  		Proto:      "HTTP/2",
    41  		ProtoMajor: 2,
    42  		ProtoMinor: 0,
    43  		Body:       preader,
    44  		Header: map[string][]string{
    45  			"Accept-Encoding": {"identity"},
    46  		},
    47  	}
    48  
    49  	// it will be close at :  `func (hc *h2Conn) Close() error`
    50  	res, err := hc.ClientConn.RoundTrip(&req)
    51  	if err != nil {
    52  		return err
    53  	}
    54  
    55  	hc.pwriter = pwriter
    56  	hc.res = res
    57  
    58  	return nil
    59  }
    60  
    61  // Read implements net.Conn.Read()
    62  func (hc *h2Conn) Read(b []byte) (int, error) {
    63  	if hc.res != nil && !hc.res.Close {
    64  		n, err := hc.res.Body.Read(b)
    65  		return n, err
    66  	}
    67  
    68  	if err := hc.establishConn(); err != nil {
    69  		return 0, err
    70  	}
    71  	return hc.res.Body.Read(b)
    72  }
    73  
    74  // Write implements io.Writer.
    75  func (hc *h2Conn) Write(b []byte) (int, error) {
    76  	if hc.pwriter != nil {
    77  		return hc.pwriter.Write(b)
    78  	}
    79  
    80  	if err := hc.establishConn(); err != nil {
    81  		return 0, err
    82  	}
    83  	return hc.pwriter.Write(b)
    84  }
    85  
    86  func (hc *h2Conn) Close() error {
    87  	if err := hc.pwriter.Close(); err != nil {
    88  		return err
    89  	}
    90  	if err := hc.ClientConn.Shutdown(hc.res.Request.Context()); err != nil {
    91  		return err
    92  	}
    93  	return hc.Conn.Close()
    94  }
    95  
    96  func StreamH2Conn(conn net.Conn, cfg *H2Config) (net.Conn, error) {
    97  	transport := &http2.Transport{}
    98  
    99  	cconn, err := transport.NewClientConn(conn)
   100  	if err != nil {
   101  		return nil, err
   102  	}
   103  
   104  	return &h2Conn{
   105  		Conn:       conn,
   106  		ClientConn: cconn,
   107  		cfg:        cfg,
   108  	}, nil
   109  }