github.com/geph-official/geph2@v0.22.6-0.20210211030601-f527cb59b0df/libs/warpfront/Client.go (about)

     1  package warpfront
     2  
     3  import (
     4  	"bytes"
     5  	"context"
     6  	"crypto/rand"
     7  	"encoding/binary"
     8  	"fmt"
     9  	"io"
    10  	"io/ioutil"
    11  	"net"
    12  	"net/http"
    13  
    14  	"golang.org/x/time/rate"
    15  )
    16  
    17  func getWithHost(client *http.Client, url string, host string) (resp *http.Response, err error) {
    18  	req, err := http.NewRequest("GET", url, nil)
    19  	if err != nil {
    20  		return
    21  	}
    22  	req.Host = host
    23  	return client.Do(req)
    24  }
    25  
    26  func postWithHost(client *http.Client, url string, host string, body io.Reader) (resp *http.Response, err error) {
    27  	req, err := http.NewRequest("POST", url, body)
    28  	if err != nil {
    29  		return
    30  	}
    31  	req.Host = host
    32  	req.Header.Add("Content-Type", "application/octet-stream")
    33  	return client.Do(req)
    34  }
    35  
    36  // Connect returns a warpfront session connected to the given front and real host. The front must contain a protocol scheme (http:// or https://).
    37  func Connect(client *http.Client, frontHost string, realHost string) (net.Conn, error) {
    38  	// generate session number
    39  	num := make([]byte, 32)
    40  	rand.Read(num)
    41  	// register our session
    42  	resp, err := getWithHost(client, fmt.Sprintf("%v/register?id=%x", frontHost, num), realHost)
    43  	if err != nil {
    44  		return nil, err
    45  	}
    46  	if resp.StatusCode != http.StatusOK {
    47  		return nil, fmt.Errorf("bad status code: %v", resp.StatusCode)
    48  	}
    49  	sesh := newSession()
    50  
    51  	emptyGetLimiter := rate.NewLimiter(1, 10)
    52  
    53  	go func() {
    54  		defer sesh.Close()
    55  		// poll and stuff into rx
    56  		for i := 0; ; i++ {
    57  			resp, err := getWithHost(client,
    58  				fmt.Sprintf("%v/%x?serial=%v", frontHost, num, i),
    59  				realHost)
    60  			if err != nil {
    61  				return
    62  			}
    63  			if resp.StatusCode != http.StatusOK {
    64  				err = fmt.Errorf("unexpected status code %v", resp.StatusCode)
    65  				resp.Body.Close()
    66  				return
    67  			}
    68  			totpkts := 0
    69  			for {
    70  				lbts := make([]byte, 4)
    71  				_, err := io.ReadFull(resp.Body, lbts)
    72  				if err != nil {
    73  					resp.Body.Close()
    74  					return
    75  				}
    76  				if binary.BigEndian.Uint32(lbts) == 0 {
    77  					//log.Println("warpfront: client got continuation signal, looping around")
    78  					resp.Body.Close()
    79  					goto OUT
    80  				}
    81  				totpkts++
    82  				buf := make([]byte, binary.BigEndian.Uint32(lbts))
    83  				_, err = io.ReadFull(resp.Body, buf)
    84  				if err != nil {
    85  					resp.Body.Close()
    86  					return
    87  				}
    88  				select {
    89  				case sesh.rx <- buf:
    90  				case <-sesh.ded:
    91  					resp.Body.Close()
    92  					return
    93  				}
    94  				if err != nil {
    95  					resp.Body.Close()
    96  					return
    97  				}
    98  			}
    99  		OUT:
   100  			if totpkts == 0 {
   101  				emptyGetLimiter.Wait(context.Background())
   102  			}
   103  		}
   104  	}()
   105  	go func() {
   106  		defer sesh.Close()
   107  		// drain something from tx
   108  		//timer := time.NewTicker(time.Millisecond * 250)
   109  		buff := bytes.NewBuffer(nil)
   110  		for i := 0; ; i++ {
   111  			select {
   112  			//case <-timer.C:
   113  			case bts := <-sesh.tx:
   114  				buff.Write(bts)
   115  				if buff.Len() > 0 {
   116  					resp, err := postWithHost(client,
   117  						fmt.Sprintf("%v/%x?serial=%v", frontHost, num, i),
   118  						realHost,
   119  						buff)
   120  					if err != nil {
   121  						return
   122  					}
   123  					io.Copy(ioutil.Discard, resp.Body)
   124  					resp.Body.Close()
   125  				}
   126  			case <-sesh.ded:
   127  				//timer.Stop()
   128  				return
   129  			}
   130  		}
   131  	}()
   132  
   133  	// couple closing the session with deletion
   134  	go func() {
   135  		<-sesh.ded
   136  		getWithHost(client, fmt.Sprintf("%v/delete?id=%x", frontHost, num), realHost)
   137  	}()
   138  
   139  	// return the sesh
   140  	return sesh, nil
   141  }