1
0

util.go 4.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182
  1. package tests
  2. import (
  3. "encoding/base64"
  4. "encoding/json"
  5. "errors"
  6. "fmt"
  7. "io/ioutil"
  8. "net"
  9. "net/http"
  10. "net/url"
  11. "strings"
  12. "time"
  13. "github.com/fatedier/frp/client"
  14. frpNet "github.com/fatedier/frp/utils/net"
  15. )
  16. func getProxyStatus(name string) (status *client.ProxyStatusResp, err error) {
  17. req, err := http.NewRequest("GET", "http://"+ADMIN_ADDR+"/api/status", nil)
  18. if err != nil {
  19. return status, err
  20. }
  21. authStr := "Basic " + base64.StdEncoding.EncodeToString([]byte(ADMIN_USER+":"+ADMIN_PWD))
  22. req.Header.Add("Authorization", authStr)
  23. resp, err := http.DefaultClient.Do(req)
  24. if err != nil {
  25. return status, err
  26. } else {
  27. if resp.StatusCode != 200 {
  28. return status, fmt.Errorf("admin api status code [%d]", resp.StatusCode)
  29. }
  30. defer resp.Body.Close()
  31. body, err := ioutil.ReadAll(resp.Body)
  32. if err != nil {
  33. return status, err
  34. }
  35. allStatus := &client.StatusResp{}
  36. err = json.Unmarshal(body, &allStatus)
  37. if err != nil {
  38. return status, fmt.Errorf("unmarshal http response error: %s", strings.TrimSpace(string(body)))
  39. }
  40. for _, s := range allStatus.Tcp {
  41. if s.Name == name {
  42. return &s, nil
  43. }
  44. }
  45. for _, s := range allStatus.Udp {
  46. if s.Name == name {
  47. return &s, nil
  48. }
  49. }
  50. for _, s := range allStatus.Http {
  51. if s.Name == name {
  52. return &s, nil
  53. }
  54. }
  55. for _, s := range allStatus.Https {
  56. if s.Name == name {
  57. return &s, nil
  58. }
  59. }
  60. for _, s := range allStatus.Stcp {
  61. if s.Name == name {
  62. return &s, nil
  63. }
  64. }
  65. for _, s := range allStatus.Xtcp {
  66. if s.Name == name {
  67. return &s, nil
  68. }
  69. }
  70. }
  71. return status, errors.New("no proxy status found")
  72. }
  73. func sendTcpMsg(addr string, msg string) (res string, err error) {
  74. c, err := frpNet.ConnectTcpServer(addr)
  75. if err != nil {
  76. err = fmt.Errorf("connect to tcp server error: %v", err)
  77. return
  78. }
  79. defer c.Close()
  80. return sendTcpMsgByConn(c, msg)
  81. }
  82. func sendTcpMsgByConn(c net.Conn, msg string) (res string, err error) {
  83. timer := time.Now().Add(5 * time.Second)
  84. c.SetDeadline(timer)
  85. c.Write([]byte(msg))
  86. buf := make([]byte, 2048)
  87. n, errRet := c.Read(buf)
  88. if errRet != nil {
  89. err = fmt.Errorf("read from tcp server error: %v", errRet)
  90. return
  91. }
  92. return string(buf[:n]), nil
  93. }
  94. func sendUdpMsg(addr string, msg string) (res string, err error) {
  95. udpAddr, errRet := net.ResolveUDPAddr("udp", addr)
  96. if errRet != nil {
  97. err = fmt.Errorf("resolve udp addr error: %v", err)
  98. return
  99. }
  100. conn, errRet := net.DialUDP("udp", nil, udpAddr)
  101. if errRet != nil {
  102. err = fmt.Errorf("dial udp server error: %v", err)
  103. return
  104. }
  105. defer conn.Close()
  106. _, err = conn.Write([]byte(msg))
  107. if err != nil {
  108. err = fmt.Errorf("write to udp server error: %v", err)
  109. return
  110. }
  111. buf := make([]byte, 2048)
  112. n, errRet := conn.Read(buf)
  113. if errRet != nil {
  114. err = fmt.Errorf("read from udp server error: %v", err)
  115. return
  116. }
  117. return string(buf[:n]), nil
  118. }
  119. func sendHttpMsg(method, urlStr string, host string, header map[string]string, proxy string) (code int, body string, err error) {
  120. req, errRet := http.NewRequest(method, urlStr, nil)
  121. if errRet != nil {
  122. err = errRet
  123. return
  124. }
  125. if host != "" {
  126. req.Host = host
  127. }
  128. for k, v := range header {
  129. req.Header.Set(k, v)
  130. }
  131. tr := &http.Transport{
  132. DialContext: (&net.Dialer{
  133. Timeout: 30 * time.Second,
  134. KeepAlive: 30 * time.Second,
  135. DualStack: true,
  136. }).DialContext,
  137. MaxIdleConns: 100,
  138. IdleConnTimeout: 90 * time.Second,
  139. TLSHandshakeTimeout: 10 * time.Second,
  140. ExpectContinueTimeout: 1 * time.Second,
  141. }
  142. if len(proxy) != 0 {
  143. tr.Proxy = func(req *http.Request) (*url.URL, error) {
  144. return url.Parse(proxy)
  145. }
  146. }
  147. client := http.Client{
  148. Transport: tr,
  149. }
  150. resp, errRet := client.Do(req)
  151. if errRet != nil {
  152. err = errRet
  153. return
  154. }
  155. code = resp.StatusCode
  156. buf, errRet := ioutil.ReadAll(resp.Body)
  157. if errRet != nil {
  158. err = errRet
  159. return
  160. }
  161. body = string(buf)
  162. return
  163. }
  164. func basicAuth(username, passwd string) string {
  165. auth := username + ":" + passwd
  166. return "Basic " + base64.StdEncoding.EncodeToString([]byte(auth))
  167. }