github.com/ipfans/trojan-go@v0.11.0/common/net.go (about)

     1  package common
     2  
     3  import (
     4  	"fmt"
     5  	"io"
     6  	"io/ioutil"
     7  	"net"
     8  	"net/http"
     9  	"net/url"
    10  	"os"
    11  	"strconv"
    12  	"strings"
    13  	"time"
    14  )
    15  
    16  const (
    17  	KiB = 1024
    18  	MiB = KiB * 1024
    19  	GiB = MiB * 1024
    20  )
    21  
    22  func HumanFriendlyTraffic(bytes uint64) string {
    23  	if bytes <= KiB {
    24  		return fmt.Sprintf("%d B", bytes)
    25  	}
    26  	if bytes <= MiB {
    27  		return fmt.Sprintf("%.2f KiB", float32(bytes)/KiB)
    28  	}
    29  	if bytes <= GiB {
    30  		return fmt.Sprintf("%.2f MiB", float32(bytes)/MiB)
    31  	}
    32  	return fmt.Sprintf("%.2f GiB", float32(bytes)/GiB)
    33  }
    34  
    35  func PickPort(network string, host string) int {
    36  	switch network {
    37  	case "tcp":
    38  		for retry := 0; retry < 16; retry++ {
    39  			l, err := net.Listen("tcp", host+":0")
    40  			if err != nil {
    41  				continue
    42  			}
    43  			defer l.Close()
    44  			_, port, err := net.SplitHostPort(l.Addr().String())
    45  			Must(err)
    46  			p, err := strconv.ParseInt(port, 10, 32)
    47  			Must(err)
    48  			return int(p)
    49  		}
    50  	case "udp":
    51  		for retry := 0; retry < 16; retry++ {
    52  			conn, err := net.ListenPacket("udp", host+":0")
    53  			if err != nil {
    54  				continue
    55  			}
    56  			defer conn.Close()
    57  			_, port, err := net.SplitHostPort(conn.LocalAddr().String())
    58  			Must(err)
    59  			p, err := strconv.ParseInt(port, 10, 32)
    60  			Must(err)
    61  			return int(p)
    62  		}
    63  	default:
    64  		return 0
    65  	}
    66  	return 0
    67  }
    68  
    69  func WriteAllBytes(writer io.Writer, payload []byte) error {
    70  	for len(payload) > 0 {
    71  		n, err := writer.Write(payload)
    72  		if err != nil {
    73  			return err
    74  		}
    75  		payload = payload[n:]
    76  	}
    77  	return nil
    78  }
    79  
    80  func WriteFile(path string, payload []byte) error {
    81  	writer, err := os.Create(path)
    82  	if err != nil {
    83  		return err
    84  	}
    85  	defer writer.Close()
    86  
    87  	return WriteAllBytes(writer, payload)
    88  }
    89  
    90  func FetchHTTPContent(target string) ([]byte, error) {
    91  	parsedTarget, err := url.Parse(target)
    92  	if err != nil {
    93  		return nil, fmt.Errorf("invalid URL: %s", target)
    94  	}
    95  
    96  	if s := strings.ToLower(parsedTarget.Scheme); s != "http" && s != "https" {
    97  		return nil, fmt.Errorf("invalid scheme: %s", parsedTarget.Scheme)
    98  	}
    99  
   100  	client := &http.Client{
   101  		Timeout: 30 * time.Second,
   102  	}
   103  	resp, err := client.Do(&http.Request{
   104  		Method: "GET",
   105  		URL:    parsedTarget,
   106  		Close:  true,
   107  	})
   108  	if err != nil {
   109  		return nil, fmt.Errorf("failed to dial to %s", target)
   110  	}
   111  	defer resp.Body.Close()
   112  
   113  	if resp.StatusCode != http.StatusOK {
   114  		return nil, fmt.Errorf("unexpected HTTP status code: %d", resp.StatusCode)
   115  	}
   116  
   117  	content, err := ioutil.ReadAll(resp.Body)
   118  	if err != nil {
   119  		return nil, fmt.Errorf("failed to read HTTP response")
   120  	}
   121  
   122  	return content, nil
   123  }